diff --git a/.github/workflows/book.yml b/.github/workflows/book.yml index 732ff8b667..8bc1de2424 100644 --- a/.github/workflows/book.yml +++ b/.github/workflows/book.yml @@ -15,7 +15,8 @@ on: # only ever built after merge, so a PR that broke it could not be caught by # review or by CI. pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - 'cuda-oxide/cuda-oxide-book/**' - 'cutile-rs/cutile-book/**' diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bd0c764954..15a69502bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,17 +12,20 @@ on: - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' - '.github/workflows/**' pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - 'cuda-oxide/**' - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' diff --git a/.github/workflows/clippy.yml b/.github/workflows/clippy.yml index b7a72422a3..b5035eb9fa 100644 --- a/.github/workflows/clippy.yml +++ b/.github/workflows/clippy.yml @@ -19,10 +19,10 @@ jobs: uses: actions/checkout@v4 - name: Install CUDA toolkit - uses: Jimver/cuda-toolkit@3d45d157f327c09c04b50ee6ccdea2d9d017ec76 # v0.2.35 + uses: Jimver/cuda-toolkit@0ca15bb027fc96f7f61ddbb9c5221d699da06bdb # v0.2.36 id: cuda-toolkit with: - cuda: "13.1.1" + cuda: "13.3.1" method: network sub-packages: '["nvcc", "cudart-dev"]' non-cuda-sub-packages: '["libcurand-dev"]' @@ -37,7 +37,7 @@ jobs: working-directory: cuda-oxide run: cargo clippy --workspace --all-targets -- -D warnings - # `rustc-codegen-cuda` and each example under examples/ are their own + # `rustc-codegen-cuda`, the standard and CuTe examples, and CuTe models have # standalone workspaces (the codegen backend is swapped in via # RUSTFLAGS, so they can't live in the root [workspace]). That's 60+ # cargo invocations from different workspace roots. By default each @@ -49,7 +49,7 @@ jobs: # Point all of those invocations at one shared CARGO_TARGET_DIR under # ${RUNNER_TEMP}. Cargo fingerprints each unit by package + features # + workspace_root + toolchain, so co-located workspaces using the - # same pinned nightly deduplicate on identical deps and compile each + # same pinned toolchain deduplicate on identical deps and compile each # one once. This is intentionally scoped *not* to cover the # `cargo clippy --workspace` step above, which already has a # well-defined `target/` under cuda-oxide. @@ -60,18 +60,26 @@ jobs: working-directory: cuda-oxide/crates/rustc-codegen-cuda run: cargo clippy --all-targets -- -D warnings - - name: Run clippy (rustc-codegen-cuda examples) + - name: Run clippy (standard and CuTe examples and models) env: CUDA_TOOLKIT_PATH: ${{ steps.cuda-toolkit.outputs.CUDA_PATH }} CARGO_TARGET_DIR: ${{ runner.temp }}/clippy-shared-target + RUSTC_BOOTSTRAP: "1" run: | set -e - for ex in cuda-oxide/crates/rustc-codegen-cuda/examples/*/; do - if [ ! -f "$ex/Cargo.toml" ]; then continue; fi - name=$(basename "$ex") - echo "::group::clippy: $name" - (cd "$ex" && cargo clippy --all-targets -- -D warnings) - echo "::endgroup::" + example_roots=( + cuda-oxide/crates/rustc-codegen-cuda/examples + cuda-oxide/cute-rs-preview/examples + cuda-oxide/cute-rs-preview/models + ) + for example_root in "${example_roots[@]}"; do + for ex in "$example_root"/*/; do + if [ ! -f "$ex/Cargo.toml" ]; then continue; fi + name=$(basename "$ex") + echo "::group::clippy: $name" + (cd "$ex" && cargo clippy --all-targets -- -D warnings) + echo "::endgroup::" + done done # The loop above is one level deep. A few examples nest a second @@ -97,11 +105,11 @@ jobs: grep -q "\"${sub}\""; then continue fi - echo "::group::clippy: ${dir#cuda-oxide/crates/rustc-codegen-cuda/examples/}" + echo "::group::clippy: ${dir#cuda-oxide/}" (cd "${dir}" && cargo clippy --locked --all-targets -- -D warnings) echo "::endgroup::" done < <( - find cuda-oxide/crates/rustc-codegen-cuda/examples -mindepth 3 -name Cargo.toml \ + find "${example_roots[@]}" -mindepth 3 -name Cargo.toml \ -not -path '*/target/*' | sort ) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 60aaba1f3d..f3fb70e605 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -15,7 +15,8 @@ on: - '**/*.sh' - '**/*.py' pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - '**/*.rs' - '**/Cargo.toml' diff --git a/.github/workflows/cutile-rs.yml b/.github/workflows/cutile-rs.yml index 11d0496217..94d2b581c0 100644 --- a/.github/workflows/cutile-rs.yml +++ b/.github/workflows/cutile-rs.yml @@ -64,6 +64,7 @@ jobs: - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' @@ -155,8 +156,11 @@ jobs: command -v tileiras tileiras --help | grep -q "sm_120" + - uses: $/cutile-rs/.github/actions/setup - uses: actions/checkout@v7 - - uses: ./cutile-rs/.github/actions/setup + + - name: Mark checkout as safe + run: git config --global --add safe.directory /__w/cuda-rust/cuda-rust # Cheapest checks first so a formatting or lint slip fails in seconds. - name: Format check @@ -184,6 +188,71 @@ jobs: - name: CPU only tests run: ./scripts/run_cpu_tests.sh + # nextest archives carry test executables, Cargo metadata, and build + # outputs needed to run on another machine without rebuilding Rust. + # Keep both workspaces separate, then bundle them in one downloadable tar. + - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + name: Install cargo-nextest + with: + tool: cargo-nextest + + - name: Archive Tile runtime test packages + run: | + mkdir -p "$RUNNER_TEMP/cutile-gpu-tests" + cargo nextest archive --locked --features experimental-tune \ + --archive-file "$RUNNER_TEMP/cutile-gpu-tests/tile-tests.tar.zst" + + - name: Archive shared-host runtime test packages + working-directory: ${{ github.workspace }} + run: | + cargo nextest archive --locked \ + -p cuda-bindings -p cuda-core -p cuda-async \ + --archive-file "$RUNNER_TEMP/cutile-gpu-tests/host-tests.tar.zst" + + - name: Bundle runtime test archives + run: | + cp "$(command -v cargo-nextest)" "$RUNNER_TEMP/cutile-gpu-tests/cargo-nextest" + git rev-parse HEAD > "$RUNNER_TEMP/cutile-gpu-tests/revision.txt" + cargo nextest --version > "$RUNNER_TEMP/cutile-gpu-tests/nextest-version.txt" + tar -cf "$RUNNER_TEMP/cutile-gpu-tests.tar" \ + -C "$RUNNER_TEMP/cutile-gpu-tests" . + + - name: Upload runtime test packages + uses: actions/upload-artifact@v4 + with: + name: cutile-gpu-tests + path: ${{ runner.temp }}/cutile-gpu-tests.tar + if-no-files-found: error + compression-level: 0 # The nested nextest archives are already compressed. + retention-days: 1 + + gpu-tests: + needs: build + runs-on: linux-amd64-gpu-rtxpro6000-latest-1 + timeout-minutes: 30 + container: + image: nvidia/cuda:13.3.0-devel-ubuntu24.04 + env: + CUDA_TOOLKIT_PATH: /usr/local/cuda + CUTILE_TILEIRAS_PATH: /usr/local/cuda/bin/tileiras + steps: + - uses: actions/checkout@v7 + + - uses: actions/download-artifact@v4 + with: + name: cutile-gpu-tests + path: ${{ runner.temp }}/gpu-tests + + - name: Extract runtime test packages + run: | + tar -xf "$RUNNER_TEMP/gpu-tests/cutile-gpu-tests.tar" -C "$RUNNER_TEMP/gpu-tests" + test "$(cat "$RUNNER_TEMP/gpu-tests/revision.txt")" = "$GITHUB_SHA" + nvidia-smi + "$CUTILE_TILEIRAS_PATH" --version + + - name: GPU runtime tests + run: ./scripts/run_gpu_tests.sh --archive-dir "$RUNNER_TEMP/gpu-tests" + # Breaking API changes must ship with the right version bump: compare each # publishable crate's API against its latest crates.io release and fail the # PR when a break arrives without one (the 0.3.0 -> 0.3.1 lesson). Same @@ -288,7 +357,7 @@ jobs: - name: Native f16 DeviceCopy working-directory: ${{ github.workspace }} - run: cargo test -p cuda-core --features f16 --test simt_device_copy_impls + run: cargo test -p cuda-core --features f16 --test device_copy_impls - name: loom model check (interleavings + weak memory) working-directory: ${{ github.workspace }} diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index af8fdf3645..89340077a5 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -10,17 +10,20 @@ on: - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' - '.github/workflows/**' pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - 'cuda-oxide/**' - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' @@ -50,10 +53,10 @@ jobs: - name: Install CUDA toolkit # `cargo doc --workspace` compiles cuda-bindings' build-time bindgen, # which needs `cuda.h`. Same toolkit pin as the unit-tests job. - uses: Jimver/cuda-toolkit@3d45d157f327c09c04b50ee6ccdea2d9d017ec76 # v0.2.35 + uses: Jimver/cuda-toolkit@0ca15bb027fc96f7f61ddbb9c5221d699da06bdb # v0.2.36 id: cuda-toolkit with: - cuda: "13.1.1" + cuda: "13.3.1" method: network sub-packages: '["nvcc", "cudart-dev"]' non-cuda-sub-packages: '["libcurand-dev"]' diff --git a/.github/workflows/examples-compile.yml b/.github/workflows/examples-compile.yml index d9ab5a9548..13f3fdc5e9 100644 --- a/.github/workflows/examples-compile.yml +++ b/.github/workflows/examples-compile.yml @@ -10,17 +10,20 @@ on: - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' - '.github/workflows/**' pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - 'cuda-oxide/**' - 'cuda-bindings/**' - 'cuda-core/**' - 'cuda-core-derive/**' + - 'cuda-repr/**' - 'cuda-async/**' - 'Cargo.toml' - 'Cargo.lock' @@ -61,10 +64,10 @@ jobs: # Needed for cuda.h and curand.h: bindgen runs in the shared # cuda-bindings build script. Same setup as unit-tests.yml / clippy.yml. - name: Install CUDA toolkit - uses: Jimver/cuda-toolkit@3d45d157f327c09c04b50ee6ccdea2d9d017ec76 # v0.2.35 + uses: Jimver/cuda-toolkit@0ca15bb027fc96f7f61ddbb9c5221d699da06bdb # v0.2.36 id: cuda-toolkit with: - cuda: "13.1.1" + cuda: "13.3.1" method: network sub-packages: '["nvcc", "cudart-dev"]' # nvJitLink is packaged separately from cuda-nvcc on Linux. Keep it diff --git a/.github/workflows/fmt.yml b/.github/workflows/fmt.yml index 4d8035a248..33ba68059d 100644 --- a/.github/workflows/fmt.yml +++ b/.github/workflows/fmt.yml @@ -20,7 +20,7 @@ jobs: # rustfmt comes from cuda-oxide/rust-toolchain.toml's component list, # like clippy, - # so `cargo fmt` resolves against the same nightly as the rest of CI. + # so `cargo fmt` resolves against the same toolchain as the rest of CI. - name: Check stable host workspace formatting run: cargo fmt --all -- --check @@ -43,19 +43,25 @@ jobs: run: cargo fmt --all -- --check working-directory: cuda-oxide/crates/cuda-macros/tests/device-only - - name: Check example crate formatting + - name: Check example and model crate formatting + env: + RUSTC_BOOTSTRAP: "1" run: | set -euo pipefail # `cargo fmt --all` also formats local path dependencies. Running it # per example repeatedly reformats the large shared workspaces. shopt -s globstar nullglob - manifests=(cuda-oxide/crates/rustc-codegen-cuda/examples/**/Cargo.toml) + manifests=( + cuda-oxide/crates/rustc-codegen-cuda/examples/**/Cargo.toml + cuda-oxide/cute-rs-preview/examples/**/Cargo.toml + cuda-oxide/cute-rs-preview/models/**/Cargo.toml + ) ((${#manifests[@]} > 0)) for manifest in "${manifests[@]}"; do example_dir="$(dirname "$manifest")" echo "::group::cargo fmt --check ($example_dir)" # Run inside the component so rustup discovers CUDA Oxide's - # directory-scoped nightly pin. + # directory-scoped toolchain pin. (cd "$example_dir" && cargo fmt -- --check) echo "::endgroup::" done diff --git a/.github/workflows/naming-guard.yml b/.github/workflows/naming-guard.yml index 00eb30f37e..312790d0e2 100644 --- a/.github/workflows/naming-guard.yml +++ b/.github/workflows/naming-guard.yml @@ -16,7 +16,8 @@ on: - 'cuda-oxide/scripts/check-reserved-prefixes.sh' - '.github/workflows/naming-guard.yml' pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] paths: - 'cuda-oxide/crates/**' - 'cuda-oxide/scripts/check-reserved-prefixes.sh' diff --git a/.github/workflows/status-guard.yml b/.github/workflows/status-guard.yml index 1bb635e1df..1759d4acc5 100644 --- a/.github/workflows/status-guard.yml +++ b/.github/workflows/status-guard.yml @@ -5,7 +5,8 @@ on: push: branches: [main] pull_request: - branches: [main] + # Include release PRs and stacks targeting prerequisite copy-pr-bot mirrors. + branches: [main, "release/v*", "pull-request/[0-9]+"] # Each commit on main gets its own group, so it starts at once. Any other ref # still shares one group and cancels the older run. @@ -60,13 +61,12 @@ jobs: - name: Verify README lists every workspace member run: bash cuda-oxide/scripts/check-crate-inventory.sh - # The pin is copied into places no build compares: the nested backend pin - # (rustup reads it for the backend's own [workspace], so a mismatch there - # is a load failure rather than a warning), the scaffold `cargo oxide new` - # writes into user projects, the devcontainer's preinstalled toolchain, the - # book's quoted [toolchain] blocks, and the dated rustup/cargo commands in - # the book and READMEs. #727 added rustfmt to both real files and left both - # book quotes behind. + # The pin is copied into places no build compares: the scaffold `cargo oxide + # new` writes into user projects, the devcontainer's preinstalled toolchain, + # the book's quoted [toolchain] blocks, the explicit rustup/cargo commands in + # the book and READMEs, and the `rust-overlay` input in `cuda-oxide/flake.lock`. + # #727 added rustfmt to both real files and left both book quotes behind; + # #1609 bumped the pin to 1.99.1 and left `flake.lock` behind. toolchain-pin-parity: name: toolchain pin parity runs-on: ubuntu-latest @@ -130,9 +130,9 @@ jobs: - name: Clone repository uses: actions/checkout@v4 -# rustfmt comes from cuda-oxide/rust-toolchain.toml's component list through the + # rustfmt comes from cuda-oxide/rust-toolchain.toml's component list through the # rustup proxy, as in fmt.yml, so the order checked is the pinned - # nightly's. No cargo, no CUDA toolkit, no network beyond the toolchain. + # toolchain's. No cargo, no CUDA toolkit, no network beyond the toolchain. - name: Verify the book's Rust blocks use rustfmt's import order run: bash cuda-oxide/scripts/check-book-import-order.sh diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 0226b0bfce..066597a4da 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -55,6 +55,18 @@ jobs: - package: dialect-mir - package: dialect-nvvm - package: mir-importer + # CuTe types, semantic IR, and MLIR export need no CUDA toolkit. + # Facade macro tests also compile host metadata via cuda-host. + - package: cute-layout + - package: cute-rs + needs_cuda: true + - package: ctm-rs + needs_cuda: true + - package: ctm-rs + test_args: --all-targets --features host + needs_cuda: true + - package: dialect-cute + - package: cuda-oxide-mlir-export # Includes a real llc -> ptxas spine-kernel validation. - package: cuda-oxide-codegen needs_cuda: true @@ -102,29 +114,14 @@ jobs: - name: Install CUDA toolkit if: matrix.needs_cuda - uses: Jimver/cuda-toolkit@3d45d157f327c09c04b50ee6ccdea2d9d017ec76 # v0.2.35 + uses: Jimver/cuda-toolkit@0ca15bb027fc96f7f61ddbb9c5221d699da06bdb # v0.2.36 id: cuda-toolkit with: - cuda: "13.1.1" + cuda: "13.3.1" method: network sub-packages: '["nvcc", "cudart-dev"]' non-cuda-sub-packages: '["libcurand-dev"]' - - name: Install LLVM 21 (llc/opt for the NVPTX backend) - if: matrix.package == 'cuda-oxide-codegen' - run: | - set -euo pipefail - source /etc/os-release - sudo install -d -m 0755 /etc/apt/keyrings - wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key \ - | sudo tee /etc/apt/keyrings/apt-llvm-org.asc >/dev/null - echo "deb [signed-by=/etc/apt/keyrings/apt-llvm-org.asc] https://apt.llvm.org/${VERSION_CODENAME}/ llvm-toolchain-${VERSION_CODENAME}-21 main" \ - | sudo tee /etc/apt/sources.list.d/llvm-21.list - sudo apt-get update - sudo apt-get install -y --no-install-recommends llvm-21 - llc-21 --version | head -2 - opt-21 --version | head -2 - - name: Run crate unit tests env: CUDA_TOOLKIT_PATH: ${{ steps.cuda-toolkit.outputs.CUDA_PATH }} diff --git a/AGENTS.md b/AGENTS.md index a9836f863c..c1aab632c9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -31,8 +31,9 @@ contribution guide. There is no root `Justfile`. The SIMT recipes live in - `cuda-bindings/`, `cuda-core/`, `cuda-core-derive/`, and `cuda-async/` — shared host crates in the root workspace. Both kernel stacks depend on these by path. - `cuda-oxide/` — the SIMT model. Its `Cargo.toml` explicitly lists workspace - members under `crates/`; its `rust-toolchain.toml` pins the nightly needed for - `rustc_private`. The book is `cuda-oxide/cuda-oxide-book/`, and + members under `crates/`; its `rust-toolchain.toml` pins the stable release + (`1.99.1`) and `rustc-dev` / `llvm-tools` components needed for `rustc_private` + (unlocked via `RUSTC_BOOTSTRAP=1`). The book is `cuda-oxide/cuda-oxide-book/`, and `cuda-oxide/intrinsics/` holds the intrinsic catalog and ABI ledger. `cuda-oxide/crates/rustc-codegen-cuda/` declares its own workspace and inherits the SIMT toolchain pin. Examples and test fixtures can also declare separate @@ -84,15 +85,16 @@ nested workspaces. ## Toolchains The root `rust-toolchain.toml` and `cutile-rs/rust-toolchain.toml` carry matching -stable pins. `cuda-oxide/rust-toolchain.toml` supplies the SIMT nightly, including -for its nested backend workspace. Do not float a pin to work around a failure. +stable pins. `cuda-oxide/rust-toolchain.toml` supplies the SIMT stable pin and +compiler components (`rustc-dev`, `llvm-tools`), including for its nested backend +workspace. Do not float a pin to work around a failure. `cuda-oxide/scripts/check-toolchain-parity.sh` checks the SIMT pin against its scaffold, development container, and documentation. rustup normally resolves the toolchain from the working directory; an explicit override can change that. `--manifest-path` does not change the working directory, so pointing a root Cargo command at the SIMT manifest does not select the SIMT -nightly. Change directory first or select the required toolchain explicitly. +toolchain pin. Change directory first or select the required toolchain explicitly. Nix environments live in `cuda-oxide/flake.nix` and `cutile-rs/flake.nix`. Run `nix develop` from the relevant component. @@ -105,6 +107,9 @@ Run `nix develop` from the relevant component. notice plus `Apache-2.0`, in the form shown under "License Headers" in `CONTRIBUTING.md`. Preserve existing and third-party notices. `cuda-oxide/scripts/check-spdx-headers.sh` enforces the policy. +- **Create new branches from the user's fork of this repository.** Synchronize + the fork's `main` with upstream before branching. For stacked work, base the + next branch on the preceding feature branch in that same fork. - **Never push to the canonical upstream repository, `NVIDIA/cuda-rust`. Treat it as read-only.** Branch and push to a fork. Before pushing, run `git remote -v` and confirm the push remote is a fork of the base, comparing diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5109e9308c..ab909717cc 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -142,13 +142,16 @@ Each product has its own workspace and toolchain under `cuda-oxide/` or `cutile-rs/`. The setup commands below cover cuda-oxide; the Tile checks have their own section further down. -cuda-oxide requires the Rust nightly toolchain with `rustc_private` support. -See the [cuda-oxide README](cuda-oxide/README.md) for setup instructions. +cuda-oxide pins its stable Rust toolchain (`1.99.1`) and required `rustc-dev` +and `llvm-tools` components in `cuda-oxide/rust-toolchain.toml` (with +`RUSTC_BOOTSTRAP=1` configured via `cuda-oxide/.cargo/config.toml` and +`cargo-oxide`). See the [cuda-oxide README](cuda-oxide/README.md) for setup +instructions. -`cuda-oxide/flake.nix` provides the CUDA 13, LLVM 22, Clang, and pinned nightly -development environment. From the repository root, run `cd cuda-oxide` and -then `nix develop` with flakes enabled. Container setup is documented in -[cuda-oxide/.devcontainer/README.md](cuda-oxide/.devcontainer/README.md). +`cuda-oxide/flake.nix` provides the CUDA 13, Clang/libclang, and pinned Rust +(with LLVM 23 `llvm-tools`) development environment. From the repository root, +run `cd cuda-oxide` and then `nix develop` with flakes enabled. Container setup +is documented in [cuda-oxide/.devcontainer/README.md](cuda-oxide/.devcontainer/README.md). ### Running the cuda-oxide checks @@ -189,6 +192,17 @@ cutile-rs/scripts/run_all.sh # both of the above, then the examples and Run `cutile-rs/scripts/run_all.sh` before marking a Tile pull request ready for review. + +The GPU script also accepts `--archive-dir DIR` to run prebuilt tests without +Cargo or a Rust toolchain. The directory must contain `cargo-nextest`, +`tile-tests.tar.zst`, and `host-tests.tar.zst`, as bundled by CI. This mode runs +tests serially and uses the checkout for source fixtures. +See the archive, bundle, and extract steps in +[the Tile CI workflow](.github/workflows/cutile-rs.yml) for how to produce this +directory. Extract `cutile-gpu-tests.tar` first and use a checkout matching its +`revision.txt`; `cargo-nextest` must be executable. The script's opening comments +describe the directory layout and the contents of each test archive. + CI runs the same suites in `.github/workflows/cutile-rs.yml`. Because the Tile lanes use NVIDIA self-hosted runners, that workflow triggers on the `pull-request/` branches that copy-pr-bot mirrors rather than on the pull diff --git a/Cargo.lock b/Cargo.lock index 427b364038..cd529dff7a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -46,6 +46,26 @@ version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cc" version = "1.2.60" @@ -105,7 +125,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -121,7 +141,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -133,11 +153,13 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", "trybuild", @@ -145,13 +167,17 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "dashmap" version = "6.2.1" @@ -315,6 +341,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/Cargo.toml b/Cargo.toml index e27f86ecc4..b8d8c30278 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,11 +3,13 @@ resolver = "3" members = [ # Stable host-side crates shared by the Tile and SIMT components. # Package metadata and dependencies stay in each crate so this workspace - # does not impose CUDA Oxide's nightly toolchain or edition. + # does not impose CUDA Oxide's workspace toolchain pin or edition. "cuda-bindings", "cuda-core", "cuda-core-derive", "cuda-async", + # no_std representation contracts shared by host and device crates. + "cuda-repr", ] # Nested workspaces. Without these, Cargo treats crates under them as part of # this workspace and `edition.workspace = true` in cuda-device (etc.) looks at diff --git a/README.md b/README.md index ecbfb2d9bc..060cc7869a 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ shape its direction by sharing feedback on your experience. | Safety | Data-race freedom for disjoint slice accesses; you control thread synchronization and shared memory | Data-race freedom by construction within a kernel: no explicit thread indexing, shared memory, or synchronization | | Indexing | Explicit thread/block indices | Implicit via partitions | | Compiles | Ahead of time, Rust → PTX | JIT at first launch, Rust → Tile IR → cubin | -| Toolchain | Pinned nightly, managed by `cargo oxide` | Stable Rust 1.89 or newer | +| Toolchain | Pinned stable Rust (`1.99.1`), managed by `cargo oxide` | Stable Rust 1.89 or newer | | Portability | Produces PTX that the driver compiles for compatible GPUs; you control tuning per GPU | CUDA architecture-agnostic; the compiler tunes per GPU | | Reach for it when | You need direct control over threads, warps, shared memory, TMA, or cluster ops | Your problem is naturally expressed as tensor computations and you want the compiler to handle scheduling details | @@ -67,6 +67,7 @@ cuda-rust/ ├── cuda-bindings/ runtime — CUDA driver API bindings ├── cuda-core/ runtime — contexts, streams, memory, modules, events, launch ├── cuda-core-derive/ runtime — derive macros for cuda-core +├── cuda-repr/ shared no_std host/device representation contracts ├── cuda-async/ runtime — lazy ops, futures, streams, graphs ├── cuda-oxide/ SIMT model — compiler, device crates, cuda-oxide-book └── cutile-rs/ Tile model — compiler, tile crates, cutile-book @@ -81,7 +82,7 @@ are documented in the cutile-rs book. ## Contributing -See [CONTRIBUTING.md](CONTRIBUTING.md) for commit sign-off requirements, pull requests, and CI. Toolchain requirements differ between the two: cutile-rs builds on stable Rust, cuda-oxide needs its pinned nightly via `cargo oxide`; each book covers its own. +See [CONTRIBUTING.md](CONTRIBUTING.md) for commit sign-off requirements, pull requests, and CI. Both components build on stable Rust and pin their toolchains in their respective `rust-toolchain.toml` files (with `rustc-dev` and `llvm-tools` for `cuda-oxide`, managed via `cargo oxide`); each book covers its own setup. ## License diff --git a/cuda-async/Cargo.toml b/cuda-async/Cargo.toml index cfb9a48fb4..24fd68a5d9 100644 --- a/cuda-async/Cargo.toml +++ b/cuda-async/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" edition = "2021" rust-version = "1.89" license = "Apache-2.0" @@ -11,9 +11,9 @@ description = "Safe Async CUDA support via Async Rust." [dependencies] # Sibling paths, not the root workspace entries; see cuda-core/Cargo.toml. -cuda-core = { version = "=0.4.0", path = "../cuda-core" } +cuda-core = { version = "=0.5.0", path = "../cuda-core" } rustc-hash = "2.0" -cuda-bindings = { version = "=0.4.0", path = "../cuda-bindings" } +cuda-bindings = { version = "=0.5.0", path = "../cuda-bindings" } half = { version = "2" } futures = "0.3" anyhow = "1" diff --git a/cuda-async/README.md b/cuda-async/README.md index d9cc2c6dc5..5a0255f391 100644 --- a/cuda-async/README.md +++ b/cuda-async/README.md @@ -17,7 +17,7 @@ until `.sync()`, `.sync_on()`, or `.await` is called. use cutile::prelude::*; fn main() -> Result<(), DeviceError> { - let device = cuda_core::Device::new(0)?; + let device = cuda_core::Context::new(0)?; let stream = device.new_stream()?; let mut z = api::zeros::(&[16, 16]).sync_on(&stream)?; @@ -143,21 +143,17 @@ context; a future awaiting through them may not resolve after a fault. ### Cancellation and Forgotten Futures -Dropping a `DeviceFuture` does not cancel submitted GPU work. It controls -when the host releases the resources that work uses. A future dropped after -its first poll but before completion therefore **waits for its stream to -drain** before dropping its undelivered output (tensors, `Vec` DMA -targets, borrowed inputs). If the wait cannot be performed — the context -has faulted, or the stream is recording a graph — the output is leaked -with a message on stderr rather than freed under a running kernel. A -future dropped before its first poll submitted nothing and waits for -nothing. - -cuTile submissions retain storage independently of their outputs. Borrowed -kernel arguments, tensor views, copies, and values discarded by `.then()` -or tuple projection remain alive until completion, including when submission -fails or panics after enqueueing some work. On success these owners are -released before the future returns its result. +Dropping a `DeviceFuture` does not cancel submitted GPU work and does not wait +for it. Its output drops normally. GPU allocations, DMA owners and access leases +are registered independently with the submission before work is enqueued; they +retire after a CUDA event proves completion. If the driver cannot establish +completion, those owners and leases are deliberately leaked. An unpolled future +has submitted no work. + +Tile and SIMT launches use this same runtime. Borrowed arguments, tensor views, +copies and values discarded by `.then()` or tuple projection remain protected +through GPU completion, including partial submission errors and panics. Successful +completion releases submission owners before returning the result. `mem::forget(future)` deliberately leaks those owners and their device access leases. It does not cancel work or release the leases when the GPU finishes. @@ -166,11 +162,15 @@ access on another stream returns an error. Await or drop a future normally to release its leases. Mutable partitions and `memcpy` destinations must also have unique user-visible storage; internal submission owners do not count as aliases. -For 0.4.0, custom `KernelInputStored` and `KernelOutputStored` implementations +Custom `KernelInputStored` and `KernelOutputStored` implementations must implement `retain`. Custom `DeviceOp`s must register owned resources with `ExecutionContext::retain` before enqueueing work. The unsafe `async_on` escape -hatch leaves resource lifetime and access ordering to its caller. Device-to-host -`Vec` copies synchronize before exposing initialized elements to host combinators. +hatch leaves resource lifetime and access ordering to its caller. On return, +including an error, it releases registered owners and access leases without +synchronizing or arranging background retention. The caller must keep every +submitted allocation, module, graph and DMA buffer alive until completion, +including resources discarded by combinators. Device-to-host `Vec` copies +synchronize before exposing initialized elements to host combinators. ## CUDA Graphs @@ -234,3 +234,27 @@ is why it holds a single test and runs as its own process. To exercise the reactor rather than the inline spin, run the GPU tests with `CUDA_ASYNC_SPIN_BUDGET_US=0`; to exercise the host-callback fallback, add `CUDA_ASYNC_HOST_SYNC=spin` (or `block`). + +## Shared host API migration + +`cuda_async::simt` has been removed. Import `device_operation`, `device_future`, +`device_context`, `scheduling_policies`, `device_box` and `launch` directly from +`cuda_async`. `DeviceOperation` aliases `DeviceOp`; `and_then`, `apply` and `arc` +delegate to the shared combinators. `SingleStream::from_stream(stream)` schedules +both programming models on an existing shared `cuda_core::Stream`. + +Raw launches use `AsyncKernelLaunchBuilder`, followed by an explicit unsafe +`finalize_unchecked(config)` to produce an immutable `AsyncKernelLaunch`. +Tile's generated launchers perform this step internally. Kernel scalar arguments +share one `KernelArgument` blanket implementation for `DeviceCopy`. A generic +Tile scalar therefore needs both `DType` and `DeviceCopy`; tensor-only generic +parameters keep their `DType` bound. Over-aligned arguments are supported. +Non-`DeviceCopy` boxed arguments require `push_boxed_arg_unchecked` and an +explicit ABI proof. + +Custom `DeviceOp::execute` implementations must return valid Rust values that +can be dropped immediately: host combinators may discard them before completion. +Retain GPU/DMA resources independently before enqueueing them. Bare borrowed +host slices cannot provide ownership for a forgettable asynchronous transfer; +use an owned retained source or a synchronous transfer. Device-to-host Vec +transfers complete before exposing the Vec to host code. diff --git a/cuda-async/src/cuda_graph.rs b/cuda-async/src/cuda_graph.rs index fc163050dc..ad88fa50dc 100644 --- a/cuda-async/src/cuda_graph.rs +++ b/cuda-async/src/cuda_graph.rs @@ -7,7 +7,7 @@ use crate::device_context::with_default_device_policy; use crate::device_future::DeviceFuture; use crate::device_operation::{DeviceOp, ExecutionContext, GraphNode, ReplayResource}; use crate::error::DeviceError; -use cuda_core::{sys, Device, IntoResult, Stream}; +use cuda_core::{sys, Context, IntoResult, Stream}; use std::collections::HashMap; use std::future::IntoFuture; use std::mem::MaybeUninit; @@ -67,7 +67,7 @@ struct GraphExecHandle { /// Contexts of `update` ops not yet consumed by a replay. pending: Mutex>, /// Bound before destruction: the driver needs the owning context current. - device: Arc, + device: Arc, cu_graph: sys::CUgraph, cu_graph_exec: sys::CUgraphExec, } @@ -85,7 +85,7 @@ impl GraphExecHandle { /// instantiation fails, and owned by the returned handle otherwise /// (whose `Drop` also covers an upload failure). fn instantiate( - device: Arc, + device: Arc, cu_graph: sys::CUgraph, stream: &Stream, ) -> Result { @@ -165,7 +165,7 @@ fn capture_on( stream: &Arc, record: impl FnOnce() -> Result, ) -> Result<(R, Arc), DeviceError> { - let device = stream.device().clone(); + let device = stream.context().clone(); device.bind_to_thread()?; unsafe { @@ -204,7 +204,7 @@ fn capture_on( let exec = GraphExecHandle::instantiate(device, cu_graph, stream)?; // Drain the upload. The recorded work itself has *not* run. - unsafe { stream.synchronize() }?; + stream.synchronize()?; Ok((output, Arc::new(exec))) } @@ -367,6 +367,12 @@ impl DeviceOp for GraphLaunch { type Output = (); unsafe fn execute(self, context: &ExecutionContext) -> Result<(), DeviceError> { + if self.exec.device.as_ref() != context.context().as_ref() { + return Err(DeviceError::Internal( + "graph and launch stream belong to different CUDA contexts".into(), + )); + } + context.context().bind_to_thread()?; context.retain(self.exec.clone())?; // Reacquire every captured buffer for this replay: one lease per // distinct storage, registered on the launch submission so other @@ -377,7 +383,7 @@ impl DeviceOp for GraphLaunch { } } // Pending updates need not accumulate over the graph's lifetime. - // The launch submission now owns them; their contexts wait on the + // The launch submission now owns them; their contexts retain owners on the // update stream if it differs from the launch stream. let pending = std::mem::take(&mut *self.exec.pending.lock().unwrap_or_else(|e| e.into_inner())); diff --git a/cuda-async/src/device_box.rs b/cuda-async/src/device_box.rs new file mode 100644 index 0000000000..e85bfc624a --- /dev/null +++ b/cuda-async/src/device_box.rs @@ -0,0 +1,141 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Owned and borrowed wrappers around CUDA device pointers. +//! +//! [`DevicePointer`] is a thin, `Copy` handle to device memory suitable for +//! passing as a kernel argument. [`DeviceBox`] owns device memory and frees it +//! through the shared allocation owner without waiting on drop. Generated +//! launches retain that owner and its access leases until GPU completion. +//! Raw-pointer callers must register storage with `ExecutionContext::retain_buffer` +//! before enqueueing work, or establish completion themselves. + +pub use crate::device_buffer::DevicePointer; +use cuda_bindings::CUdeviceptr; +use std::{marker::PhantomData, sync::Arc}; + +/// Owning wrapper backed by the shared CUDA allocation owner and access tracker. +/// Generated launch adapters retain storage until GPU completion. Raw-pointer +/// callers must independently retain storage and exclude conflicting accesses. +#[derive(Debug)] +pub struct DeviceBox { + storage: cuda_core::RawDeviceBuffer, + shared_access_writes: bool, + /// Ordinal of the device that owns this allocation. + device_id: usize, + /// Raw CUDA device pointer to the start of the allocation. + cudptr: CUdeviceptr, + /// Number of elements (not bytes) in the allocation. + len: usize, + /// Ties the box to the element type. + _marker: PhantomData, +} + +/// # Safety +/// +/// Device pointers are plain integers. The CUDA driver API is thread-safe for +/// distinct streams, and `DeviceBox` holds no thread-affine state. +unsafe impl Send for DeviceBox {} + +/// # Safety +/// +/// See the [`Send`] impl. Shared references only expose the device pointer +/// value, which is safe to read from any thread. +unsafe impl Sync for DeviceBox {} + +impl DeviceBox<[DType]> { + /// Constructs a `DeviceBox<[DType]>` from a raw device pointer, element + /// count, and allocating stream. + /// + /// # Safety + /// + /// * `dptr` must point to a valid device allocation of at least + /// `len * size_of::()` bytes, allocated asynchronously on `stream`. + /// * Every element must be valid for its device-side type before use. + /// * The caller must order uses after allocation and track all raw accesses. + /// * Ownership of the allocation transfers to the returned `DeviceBox`. + pub unsafe fn from_raw_parts( + dptr: CUdeviceptr, + len: usize, + stream: Arc, + ) -> Self + where + DType: Copy, + { + let bytes = len + .checked_mul(std::mem::size_of::()) + .expect("device allocation size overflow"); + let device_id = stream.context().ordinal(); + let storage = + unsafe { cuda_core::RawDeviceBuffer::from_raw_parts_async(dptr, bytes, stream) }; + Self { + storage, + shared_access_writes: false, + _marker: PhantomData, + cudptr: dptr, + len, + device_id, + } + } + + /// Whether a shared reference may mutate through interior mutability. + #[doc(hidden)] + pub fn shared_access_writes(&self) -> bool { + self.shared_access_writes + } + + pub fn storage(&self) -> &cuda_core::RawDeviceBuffer { + &self.storage + } + + /// Returns `true` if the buffer contains zero elements. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Returns the number of elements in the buffer. + pub fn len(&self) -> usize { + self.len + } + + /// Returns the raw [`CUdeviceptr`]. + pub fn cu_deviceptr(&self) -> CUdeviceptr { + self.cudptr + } + + /// Returns the ordinal of the device that owns this allocation. + pub fn device_id(&self) -> usize { + self.device_id + } + + /// Creates a non-owning [`DevicePointer`] into this allocation. + pub fn device_pointer(&self) -> DevicePointer { + // SAFETY: this box retains the allocation described by its pointer. + unsafe { DevicePointer::from_cu_deviceptr(self.cudptr) } + } +} + +impl DeviceBox<[T]> { + /// Transfers a typed buffer without copying device memory. Storage identity, + /// context, allocation extent and access tracking remain unchanged. + pub fn from_buffer(buffer: cuda_core::DeviceBuffer) -> Self { + Self { + cudptr: buffer.cu_deviceptr(), + len: buffer.len(), + device_id: buffer.context().ordinal(), + shared_access_writes: buffer.shared_access_writes(), + storage: buffer.into_storage(), + _marker: PhantomData, + } + } + + /// Transfers ownership back to the shared typed buffer without a GPU copy. + pub fn into_buffer(self) -> cuda_core::DeviceBuffer { + // SAFETY: this box owns valid T elements; consuming it transfers that + // ownership and the same access tracker without creating a typed alias. + unsafe { cuda_core::DeviceBuffer::from_storage_copy(self.storage, self.len) } + .expect("DeviceBox storage must cover its elements") + } +} diff --git a/cuda-async/src/device_buffer.rs b/cuda-async/src/device_buffer.rs index a6da6f5a75..6245317b34 100644 --- a/cuda-async/src/device_buffer.rs +++ b/cuda-async/src/device_buffer.rs @@ -3,70 +3,29 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Wrappers around CUDA device pointers: owned (freed on drop), foreign (kept -//! alive by an external owner via a liveness token), and raw-borrowed (an -//! `unsafe` escape that does nothing on drop). +//! Shared allocation ownership and typed device addresses. -use crate::device_context::with_deallocator_stream; -use cuda_core::free_async; use cuda_core::sys::CUdeviceptr; +pub use cuda_core::{DeviceAllocation, RawDeviceBuffer as DeviceBuffer}; use std::marker::PhantomData; -use std::sync::Arc; - -/// A CUDA device allocation owned by an external framework (cudarc, torch, a VMM -/// range, ...). -/// -/// A borrowed [`DeviceBuffer`] can hold an `Arc` as a -/// *liveness token*: while the buffer — and anything derived from it — is alive, -/// the owner's refcount is nonzero, so the memory cannot be freed out from under -/// cutile. This turns liveness from an ongoing caller promise into a fact -/// enforced at construction. -/// -/// # Safety -/// Implementors must guarantee that, for as long as the implementor is alive: -/// - `device_ptr()` returns a valid CUDA device pointer on `device_id()`, -/// - the allocation there is at least `len_bytes()` bytes, and -/// - all three getters are **stable**: they return the same values on every -/// call for the implementor's entire lifetime. Consumers read them once at -/// construction and cache the results (`DeviceBuffer::foreign`), so an -/// implementor whose pointer moves would silently diverge from the cached -/// copy. This is the same stability `std`'s allocator traits demand of a -/// currently-allocated block. -/// -/// This positional obligation is discharged **once, in the `impl`** — not at -/// each borrow. The temporal obligation (nobody else touches the bytes while -/// cutile work is in flight) cannot live here: it is per-borrow, and is why -/// `Tensor::from_foreign` remains an `unsafe` call. -pub unsafe trait DeviceAllocation: Send + Sync + 'static { - /// The device pointer of the allocation. - fn device_ptr(&self) -> CUdeviceptr; - /// The size of the allocation in bytes. - fn len_bytes(&self) -> usize; - /// The device ordinal the allocation lives on. - fn device_id(&self) -> usize; -} - -/// How a [`DeviceBuffer`]'s memory is owned. -enum Owner { - /// cutile allocated the memory; free it asynchronously on drop. - Owned, - /// A raw, unowned borrow (via `borrowed_from_raw_parts`); do nothing on drop. - /// Liveness is the caller's unverified obligation. - Borrowed, - /// A foreign allocation kept alive by its owner; do nothing on drop, and - /// hold the owner so the memory provably outlives this buffer. The `Arc` is - /// held purely for that RAII liveness effect — never read directly. - Foreign(#[allow(dead_code)] Arc), -} /// A non-owning, copyable handle to a typed device memory address. -#[derive(Debug, Copy, Clone)] +#[derive(Debug)] +#[repr(transparent)] pub struct DevicePointer { - dtype: PhantomData, + dtype: PhantomData T>, pub dptr: CUdeviceptr, } -unsafe impl Send for DevicePointer {} +impl Copy for DevicePointer {} +impl Clone for DevicePointer { + fn clone(&self) -> Self { + *self + } +} +// SAFETY: this transparent u64 address owns no memory; all address bit patterns, +// including zero, are valid wrapper values. Dereferencing remains unsafe. +unsafe impl cuda_core::DeviceCopy for DevicePointer {} impl DevicePointer { /// Returns the raw CUDA device pointer. @@ -85,141 +44,3 @@ impl DevicePointer { } } } - -/// A type-erased handle to a CUDA device memory allocation. -/// -/// `DeviceBuffer` manages a raw byte buffer on the GPU. It stores only the device pointer, -/// byte length, and device id — all type information lives in higher-level wrappers -/// such as `Tensor`. -/// -/// An *owned* buffer (the default) is freed asynchronously on drop. A *foreign* -/// buffer (see [`foreign`](Self::foreign)) wraps memory owned by an external -/// allocator (cudarc, torch, a VMM range, ...), holds a liveness token, and is a -/// no-op on drop. A raw *borrowed* buffer -/// (see [`borrowed_from_raw_parts`](Self::borrowed_from_raw_parts)) is the same -/// no-op-on-drop but without a token — an `unsafe` escape whose liveness is the -/// caller's responsibility. -pub struct DeviceBuffer { - device_id: usize, - cudptr: CUdeviceptr, - len: usize, - owner: Owner, -} - -impl std::fmt::Debug for DeviceBuffer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let kind = match self.owner { - Owner::Owned => "owned", - Owner::Borrowed => "borrowed", - Owner::Foreign(_) => "foreign", - }; - f.debug_struct("DeviceBuffer") - .field("device_id", &self.device_id) - .field("cudptr", &self.cudptr) - .field("len", &self.len) - .field("owner", &kind) - .finish() - } -} - -unsafe impl Send for DeviceBuffer {} -unsafe impl Sync for DeviceBuffer {} - -impl Drop for DeviceBuffer { - fn drop(&mut self) { - // Only owned buffers are freed here. Borrowed/foreign buffers wrap memory - // owned elsewhere; freeing it would be a double free (or invalid, e.g. - // for a VMM-mapped range). A `Foreign` owner's `Arc` drops with `self`. - if !matches!(self.owner, Owner::Owned) { - return; - } - // A destructor has no way to report failure, and a panic here would - // abort the process if this drop runs during an unwind. Both a missing - // device context and a refused `cuMemFreeAsync` are therefore ignored: - // the allocation leaks until the primary context is released, which - // is the least bad outcome for memory the driver would not free. - // SAFETY: the pointer was allocated stream-ordered by this crate, and - // the driver completes queued work on the deallocator stream before - // the memory is handed out again. - let _ = unsafe { - with_deallocator_stream(self.device_id, |stream| free_async(self.cudptr, stream)) - }; - } -} - -impl DeviceBuffer { - /// Constructs an owned `DeviceBuffer` from a raw device pointer, byte length, and device id. - /// - /// The buffer takes ownership of the allocation and frees it asynchronously on drop. - /// - /// # Safety - /// The caller must ensure `dptr` points to a valid device allocation of at least - /// `len_bytes` bytes that is safe to free via the deallocator stream. - pub unsafe fn from_raw_parts(dptr: CUdeviceptr, len_bytes: usize, device_id: usize) -> Self { - Self { - cudptr: dptr, - len: len_bytes, - device_id, - owner: Owner::Owned, - } - } - - /// Wraps `len_bytes` of a foreign allocation, holding `owner` alive so the - /// memory outlives this buffer. Dropping it does **not** free the memory. - /// - /// `len_bytes` is the borrowing view's **logical** size (what the tensor - /// invariant checks), not the owner's allocation size — the owner may be - /// larger, and the addressable-extent-vs-owner check happens in - /// `Tensor::from_foreign` before this is called. Liveness is enforced by - /// `owner`'s refcount; pointer validity was asserted once when `owner`'s - /// type implemented [`DeviceAllocation`]. - pub fn foreign(owner: Arc, len_bytes: usize) -> Self { - Self { - cudptr: owner.device_ptr(), - len: len_bytes, - device_id: owner.device_id(), - owner: Owner::Foreign(owner), - } - } - - /// Constructs a borrowed `DeviceBuffer` that wraps memory owned by an external - /// allocator. Dropping it does **not** free the allocation. - /// - /// # Safety - /// This borrows, rather than takes ownership of, `dptr`; dropping the buffer - /// never frees it. The caller must ensure `dptr` points to a valid device - /// allocation of at least `len_bytes` bytes, and must keep that allocation - /// mapped — never freed, reallocated, or resized — for the entire lifetime of - /// this buffer and every value derived from it. When wrapped in a tensor used - /// as a mutable output, the owner must additionally not access the same memory - /// concurrently; see `Tensor::borrow_raw_parts` in the `cutile` crate for the - /// full obligations. - pub unsafe fn borrowed_from_raw_parts( - dptr: CUdeviceptr, - len_bytes: usize, - device_id: usize, - ) -> Self { - Self { - cudptr: dptr, - len: len_bytes, - device_id, - owner: Owner::Borrowed, - } - } - /// Returns whether the allocation is empty (zero bytes). - pub fn is_empty(&self) -> bool { - self.len_bytes() == 0 - } - /// Returns the byte length of the allocation. - pub fn len_bytes(&self) -> usize { - self.len - } - /// Returns the raw CUDA device pointer. - pub fn cu_deviceptr(&self) -> CUdeviceptr { - self.cudptr - } - /// Returns the device id this allocation belongs to. - pub fn device_id(&self) -> usize { - self.device_id - } -} diff --git a/cuda-async/src/device_context.rs b/cuda-async/src/device_context.rs index a119beae28..0e418ece95 100644 --- a/cuda-async/src/device_context.rs +++ b/cuda-async/src/device_context.rs @@ -7,7 +7,7 @@ //! //! ## Architecture //! -//! - **Global (process-wide)**: one [`Device`] per device id, shared across all threads +//! - **Global (process-wide)**: one [`Context`] per device id, shared across all threads //! via [`OnceLock`]. //! //! - **Per-thread**: Scheduling policy and deallocator stream remain thread-local, since @@ -18,7 +18,7 @@ use crate::error::{device_assert, device_error, DeviceError}; use crate::scheduling_policies::{SchedulingPolicy, StreamPoolRoundRobin}; -use cuda_core::{Device, Function, MemPool, Module, Stream}; +use cuda_core::{Context, Function, MemPool, Module, Stream}; use std::cell::Cell; use std::collections::HashMap; use std::hash::{DefaultHasher, Hash, Hasher}; @@ -89,28 +89,28 @@ pub struct Validator { pub launch_checks: Vec, } -// ── Global Device (process-wide, per-device singleton) ───────────────────── +// ── Global Context (process-wide, per-device singleton) ───────────────────── /// Global per-device handles. Shared across all threads so that /// `Module`/`Function` loaded against a device can be used from any thread. -static DEVICES: OnceLock>>> = OnceLock::new(); +static DEVICES: OnceLock>>> = OnceLock::new(); -fn devices() -> &'static Mutex>> { +fn devices() -> &'static Mutex>> { DEVICES.get_or_init(|| Mutex::new(HashMap::new())) } -/// Get or create the global [`Device`] for a device ordinal. +/// Get or create the global [`Context`] for a device ordinal. /// /// The first call for a given `device_id` creates the device handle; subsequent -/// calls return the same `Arc`. -fn get_or_init_device(device_id: usize) -> Result, DeviceError> { +/// calls return the same `Arc`. +fn get_or_init_device(device_id: usize) -> Result, DeviceError> { let mut devices = devices() .lock() .map_err(|_| device_error(device_id, "device map lock poisoned"))?; if let Some(device) = devices.get(&device_id) { return Ok(Arc::clone(device)); } - let device = Device::new(device_id)?; + let device = Context::new(device_id)?; devices.insert(device_id, Arc::clone(&device)); Ok(device) } @@ -243,7 +243,7 @@ pub fn init_device( ) -> Result<(), DeviceError> { let device_context = new_device_context(device_id, policy)?; let pred = hashmap.insert(device_id, device_context).is_none(); - device_assert(device_id, pred, "Device is already initialized.") + device_assert(device_id, pred, "Context is already initialized.") } pub fn init_with_default_policy( @@ -260,7 +260,7 @@ pub fn init_with_default_policy( pool: None, }; let pred = hashmap.insert(device_id, device_context).is_none(); - device_assert(device_id, pred, "Device is already initialized.") + device_assert(device_id, pred, "Context is already initialized.") } /// Holds the thread's context map while a closure runs against it, and puts it @@ -371,10 +371,13 @@ where }) } -/// Run a closure with a reference to the [`Device`] for `device_id`. +/// Run a closure with the shared primary CUDA context for a device ordinal. +pub use with_device as with_cuda_context; + +/// Run a closure with a reference to the [`Context`] for `device_id`. pub fn with_device(device_id: usize, f: F) -> Result where - F: FnOnce(&Arc) -> R, + F: FnOnce(&Arc) -> R, { let device = get_or_init_device(device_id)?; Ok(f(&device)) @@ -426,10 +429,10 @@ pub fn set_default_device(default_device_id: usize) { /// /// # Errors /// -/// Returns [`DeviceError::Context`](crate::error::DeviceError::Context) if +/// Returns [`DeviceError::Context`] if /// `pool` was created on a different device than `device_id`. pub fn set_device_pool(device_id: usize, pool: Arc) -> Result<(), DeviceError> { - let pool_device = pool.device().ordinal(); + let pool_device = pool.context().ordinal(); device_assert( device_id, pool_device == device_id, @@ -467,7 +470,10 @@ pub fn get_device_pool(device_id: usize) -> Result>, DeviceE /// Errors from the device-context lookup are downgraded to `None`; this is the /// single choke-point for that decision so callers don't each re-derive it. pub fn pool_for_stream(stream: &Arc) -> Option> { - get_device_pool(stream.device().ordinal()).ok().flatten() + get_device_pool(stream.context().ordinal()) + .ok() + .flatten() + .filter(|pool| pool.context().as_ref() == stream.context().as_ref()) } /// Run a closure with the scheduling policy of the current thread's default device. @@ -495,12 +501,12 @@ pub fn load_module_from_file(filename: &str, device_id: usize) -> Result StreamHealth { return StreamHealth::Capturing; } // SAFETY: see above; not capturing, so a query is permitted. - match unsafe { stream.query() } { + match stream.query() { Ok(true) => StreamHealth::Idle, Ok(false) => StreamHealth::Busy, Err(e) => StreamHealth::Faulted(e), @@ -134,35 +127,15 @@ pub(crate) fn probe_stream(stream: &Stream) -> StreamHealth { /// A future that executes a [`DeviceOp`] on a CUDA stream and resolves upon completion. /// -/// # Cancellation and the in-flight result +/// # Cancellation and retained resources /// -/// Dropping a `DeviceFuture` after its first poll — after `execute` has -/// enqueued GPU work — but before it resolved leaves an *undelivered -/// result*: the operation's output, which owns the buffers the GPU is still -/// writing (a tensor, a `Vec` DMA target), or borrows the caller's. The -/// drop waits for the stream to drain and only then drops that result. On a -/// wait failure the result is leaked — loudly, unless the future already -/// delivered the stream's fault to its caller. -/// -/// cuTile registers owned storage and device access leases with the execution -/// context before enqueueing work. These survive argument recovery, output -/// projection, partial submission failures, and `mem::forget`. Forgetting the -/// future leaks both storage and leases; conflicting cross-stream use remains -/// rejected even after the original Rust borrow ends. Successful completion -/// releases the registered owners before returning the result. -/// -/// The wait is synchronous by necessity, not preference. The alternative — -/// parking the result behind a CUDA event and dropping it later, once -/// `cuEventQuery` passes (the design used by `simt::device_future`) — needs -/// the result to be `'static`, and `DeviceOp::Output` is not: cutile -/// launchers return borrowed inputs (`&'a Tensor`, -/// `Partition<&'a mut Tensor>`) as part of their output. For those the -/// result itself cannot be handed to a background thread. Registered storage -/// owners are independent of these borrows, but arbitrary custom outputs may -/// still require waiting before destruction. Rust cannot specialize on `'static`, so -/// the same rule applies to every output type until `Output: 'static` is a -/// trait-level requirement; at that point the event-gated limbo becomes the -/// default and blocking the last-resort fallback. +/// Dropping a future never waits for GPU work. Its possibly borrowed output is +/// dropped normally, just as an output projection can discard it during enqueue. +/// Every operation must independently register GPU/DMA owners and access leases +/// before submitting work. The submission retires those owners on a CUDA event; +/// error paths retain resources if completion cannot be established. Forgetting +/// a future leaks its submission and therefore preserves ownership and access +/// exclusion. Successful completion releases registered owners before delivery. #[derive(Debug)] pub struct DeviceFuture> { pub(crate) device_operation: Option, @@ -171,10 +144,6 @@ pub struct DeviceFuture> { pub(crate) error: Option, pub(crate) state: DeviceFutureState, pub(crate) callback_state: Option>, - /// Set when `poll` resolved with the stream's fault: the caller has the - /// error, so the drop-time leak of the stored result is the documented - /// consequence of that fault, not news worth reporting again. - pub(crate) fault_delivered: bool, } impl> DeviceFuture { @@ -196,7 +165,6 @@ impl> DeviceFuture { error: None, state: DeviceFutureState::Idle, callback_state: None, - fault_delivered: false, } } @@ -213,7 +181,6 @@ impl> DeviceFuture { callback_state: None, result: None, error: Some(error), - fault_delivered: false, } } @@ -294,89 +261,10 @@ impl> DeviceFuture { .expect("Expected future result to be Some.") } - /// Returns `true` when GPU work was submitted but the stored result has - /// not been handed to the caller. - /// - /// `result` is only populated by a successful [`execute`], so a stored - /// result implies submitted GPU work. The state check excludes futures - /// that never submitted anything (`Idle`, `Failed`); a `Complete` future - /// can still hold a result when completion registration failed after a - /// successful launch, or when the stream faulted before completion. - /// - /// [`execute`]: Self::execute - fn has_undelivered_submission(&self) -> bool { - matches!( - self.state, - DeviceFutureState::Executing | DeviceFutureState::Complete - ) && self.result.is_some() - } - - /// Waits for the submitted work, then drops the undelivered result. - /// - /// A stream that is idle costs one query; a busy one is synchronized. A - /// capturing stream cannot be waited on (querying it would invalidate - /// the capture) and a faulted one cannot prove completion; both fall - /// through to the loud leak in - /// [`release_in_flight_result_with`](Self::release_in_flight_result_with). + /// Release the host output immediately; its GPU owners belong to the + /// submission, which retires independently when the context is dropped. fn release_in_flight_result(&mut self) { - if !self.has_undelivered_submission() { - return; - } - let stream = self - .execution_context - .as_ref() - .map(|ctx| Arc::clone(ctx.get_cuda_stream())); - self.release_in_flight_result_with(move || { - let stream = stream.ok_or_else(|| { - DeviceError::Internal( - "Cannot release an in-flight future without an execution context.".to_string(), - ) - })?; - // The drop may run on an executor thread that never touched - // CUDA; the query and synchronize below need a current context. - stream.device().bind_to_thread()?; - match probe_stream(&stream) { - StreamHealth::Idle => Ok(()), - // SAFETY: the context was bound above; the stream is valid. - StreamHealth::Busy => unsafe { stream.synchronize() }.map_err(DeviceError::Driver), - StreamHealth::Capturing => Err(DeviceError::Internal( - "the future's stream is recording a graph; it cannot be synchronized".into(), - )), - StreamHealth::Faulted(e) => Err(DeviceError::Driver(e)), - } - }); - } - - /// Runs `wait` and drops the stored result on success; on failure the - /// result is leaked loudly, because dropping resources the device may - /// still use is worse than leaking them. - fn release_in_flight_result_with(&mut self, wait: F) - where - F: FnOnce() -> Result<(), DeviceError>, - { - if !self.has_undelivered_submission() { - return; - } - let Some(result) = self.result.take() else { - return; - }; - if let Err(error) = wait() { - if self.fault_delivered { - // The caller already received this stream's fault from - // `poll`; the leak is its documented consequence. A test - // (or program) that handled the error correctly should not - // see an error report for it. - mem::forget(result); - return; - } - crate::leak::report_leak(format_args!( - "cuda-async: leaking the result of a dropped in-flight future; the driver \ - could not prove its GPU work finished: {error}" - )); - mem::forget(result); - return; - } - drop(result); + drop(self.result.take()); } } @@ -395,7 +283,6 @@ impl> Default for DeviceFuture { error: None, state: DeviceFutureState::Idle, callback_state: None, - fault_delivered: false, } } } @@ -470,7 +357,7 @@ impl> Future for DeviceFuture { let deadline = std::time::Instant::now() + std::time::Duration::from_micros(inline_spin_budget_us()); loop { - match unsafe { ctx.get_cuda_stream().query() } { + match ctx.get_cuda_stream().query() { Ok(true) => break 'spin Ok(true), Ok(false) => {} Err(e) => break 'spin Err(e), @@ -488,11 +375,9 @@ impl> Future for DeviceFuture { } Ok(false) => {} Err(e) => { - // The result stays stored; `Drop` decides how to - // release it (it will leak, quietly, on a dead - // context — the caller receives the fault here). + // The output can drop normally. The submission preserves + // owners and access leases when completion is unknown. self.state = DeviceFutureState::Complete; - self.fault_delivered = true; return Poll::Ready(Err(DeviceError::Driver(e))); } } @@ -526,17 +411,16 @@ impl> Future for DeviceFuture { Some(ctx) => { // The poll may run on an executor thread that never // touched CUDA; the probe needs a current context. - let _ = ctx.device().bind_to_thread(); + let _ = ctx.context().bind_to_thread(); probe_stream(ctx.get_cuda_stream()) } None => StreamHealth::Busy, }; match health { StreamHealth::Faulted(e) => { - // The result stays stored for `Drop` (which leaks it, - // quietly — the caller receives the fault here). + // Submission owners and leases remain retained after + // the fault; the output itself can drop normally. self.state = DeviceFutureState::Complete; - self.fault_delivered = true; return Poll::Ready(Err(DeviceError::Driver(e))); } StreamHealth::Idle => { @@ -579,26 +463,6 @@ mod release_tests { use super::*; use crate::device_operation::Value; use std::sync::atomic::AtomicUsize; - use std::sync::Mutex; - - #[derive(Clone)] - struct DropTracker { - events: Arc>>, - } - - impl Drop for DropTracker { - fn drop(&mut self) { - self.events.lock().unwrap().push("drop"); - } - } - - struct CountDrop(Arc); - - impl Drop for CountDrop { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } fn future_in_state( state: DeviceFutureState, @@ -611,129 +475,23 @@ mod release_tests { error: None, state, callback_state: None, - fault_delivered: false, } } - /// An `Executing` future waits before its result drops. #[test] - fn release_waits_before_dropping_the_result() { - let events = Arc::new(Mutex::new(Vec::new())); - let tracker = DropTracker { - events: Arc::clone(&events), - }; - let mut future = future_in_state(DeviceFutureState::Executing, Some(tracker)); - - future.release_in_flight_result_with(|| { - events.lock().unwrap().push("wait"); - Ok(()) - }); - - assert_eq!(events.lock().unwrap().as_slice(), ["wait", "drop"]); - assert!(future.result.is_none()); - assert!(!future.has_undelivered_submission()); - } - - /// If the wait fails the result is leaked, never dropped early. - #[test] - fn release_leaks_when_the_wait_fails() { - let drops = Arc::new(AtomicUsize::new(0)); - let mut future = future_in_state( - DeviceFutureState::Executing, - Some(CountDrop(Arc::clone(&drops))), - ); - - let mut capture = crate::leak::capture::start(); - future.release_in_flight_result_with(|| Err(DeviceError::Internal("boom".to_string()))); - let reports = capture.take(); - - assert_eq!(drops.load(Ordering::Relaxed), 0); - assert!(future.result.is_none()); - assert_eq!(reports.len(), 1, "the leak must be reported: {reports:?}"); - assert!(reports[0].contains("dropped in-flight future")); - } - - /// A future that already delivered the stream's fault to its caller - /// leaks its stored result *quietly*: the caller has the error, the - /// leak is its documented consequence. - #[test] - fn release_leaks_quietly_after_the_fault_was_delivered() { - let drops = Arc::new(AtomicUsize::new(0)); - let mut future = future_in_state( - DeviceFutureState::Complete, - Some(CountDrop(Arc::clone(&drops))), - ); - future.fault_delivered = true; - - let mut capture = crate::leak::capture::start(); - future.release_in_flight_result_with(|| Err(DeviceError::Internal("boom".to_string()))); - let reports = capture.take(); - - assert_eq!( - drops.load(Ordering::Relaxed), - 0, - "must still leak, not drop" - ); - assert!(future.result.is_none()); - assert!( - reports.is_empty(), - "a delivered fault must not be re-reported: {reports:?}" - ); - } - - /// The registration-failure shape: `execute` succeeded (work submitted, - /// result stored) but completion registration failed, so poll flipped - /// the state to `Complete` with the result still stored. Release must - /// treat it exactly like a cancelled `Executing` future. - #[test] - fn release_covers_complete_future_left_by_registration_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let tracker = DropTracker { - events: Arc::clone(&events), - }; - let mut future = future_in_state(DeviceFutureState::Complete, Some(tracker)); - - assert!(future.has_undelivered_submission()); - future.release_in_flight_result_with(|| { - events.lock().unwrap().push("wait"); - Ok(()) - }); - - assert_eq!(events.lock().unwrap().as_slice(), ["wait", "drop"]); - assert!(!future.has_undelivered_submission()); - } - - /// A delivered result leaves nothing to release: no wait happens. - #[test] - fn release_is_noop_after_result_delivery() { - let mut future: DeviceFuture> = - future_in_state(DeviceFutureState::Complete, None); - - assert!(!future.has_undelivered_submission()); - future.release_in_flight_result_with(|| { - panic!("delivered futures must not wait during release") - }); - future.release_in_flight_result(); - } - - /// An `Idle` future never submitted work: no wait, and dropping it drops - /// its contents normally. - #[test] - fn release_is_noop_for_idle_future() { - let drops = Arc::new(AtomicUsize::new(0)); - let mut future = - future_in_state(DeviceFutureState::Idle, Some(CountDrop(Arc::clone(&drops)))); - - future - .release_in_flight_result_with(|| panic!("idle futures must not wait during release")); - assert_eq!(drops.load(Ordering::Relaxed), 0); - assert!(future.result.is_some()); - + fn cancellation_drops_borrowed_outputs_without_waiting() { + struct Borrowed<'a>(&'a AtomicUsize); + impl Drop for Borrowed<'_> { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + let drops = AtomicUsize::new(0); + let future = future_in_state(DeviceFutureState::Executing, Some(Borrowed(&drops))); drop(future); assert_eq!(drops.load(Ordering::Relaxed), 1); } - /// `Drop` on a `Failed` future (a scheduling error) is a plain drop. #[test] fn dropping_failed_future_is_a_noop() { let future: DeviceFuture> = diff --git a/cuda-async/src/device_operation.rs b/cuda-async/src/device_operation.rs index 9c0c51d173..59c171ac52 100644 --- a/cuda-async/src/device_operation.rs +++ b/cuda-async/src/device_operation.rs @@ -9,7 +9,7 @@ use crate::device_context::{pool_for_stream, with_default_device_policy}; use crate::device_future::DeviceFuture; use crate::error::DeviceError; use crate::scheduling_policies::SchedulingPolicy; -use cuda_core::{Device, Event, MemPool, Stream}; +use cuda_core::{Context, Event, MemPool, Stream}; use std::cell::Cell; use std::fmt::Debug; use std::future::IntoFuture; @@ -89,6 +89,9 @@ impl Drop for ExecutionLockRelease { } } +/// Compatibility spelling of the single shared operation trait. +pub use DeviceOp as DeviceOperation; + pub type DeviceOrdinal = usize; pub use crate::submission::{AccessLease, AccessTracked}; @@ -103,11 +106,27 @@ pub trait ReplayResource: Send + Sync { fn replay_identity(&self) -> (usize, bool); } +struct BufferAccess { + storage: cuda_core::RawDeviceBuffer, + write: bool, +} +impl ReplayResource for BufferAccess { + fn retain_for_launch(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + let lease = self + .storage + .acquire(ctx.get_cuda_stream(), self.write, ctx.is_recording())?; + ctx.retain(lease) + } + fn replay_identity(&self) -> (usize, bool) { + (self.storage.allocation_id(), self.write) + } +} + #[derive(Debug, Clone)] pub struct ExecutionContext { ordinal: DeviceOrdinal, cuda_stream: Arc, - device: Arc, + device: Arc, pool: Option>, submission: Arc, recording: bool, @@ -115,7 +134,7 @@ pub struct ExecutionContext { impl ExecutionContext { pub fn new(cuda_stream: Arc) -> Self { - let device = cuda_stream.device().clone(); + let device = cuda_stream.context().clone(); let ordinal = device.ordinal(); let pool = pool_for_stream(&cuda_stream); Self { @@ -135,10 +154,27 @@ impl ExecutionContext { /// operation's output. A forgotten future deliberately leaks these owners. /// /// Clones share the submission; a completed context cannot be reused. + /// [`DeviceOp::async_on`] releases these owners on return instead: its + /// unsafe caller independently guarantees resource lifetimes and ordering. pub fn retain(&self, owner: impl Send + 'static) -> Result<(), DeviceError> { self.submission.retain(owner) } + /// Registers shared allocation storage before submitting a kernel or copy. + /// Capture retains a descriptor and each graph replay reacquires access. + pub fn retain_buffer( + &self, + storage: cuda_core::RawDeviceBuffer, + write: bool, + ) -> Result<(), DeviceError> { + let descriptor = Arc::new(BufferAccess { storage, write }); + descriptor.retain_for_launch(self)?; + if self.recording { + self.record_resource(descriptor); + } + Ok(()) + } + /// Retain an in-flight access lease until this submission completes. /// The allocation-free counterpart of [`retain`](Self::retain) for the /// per-argument hot path. @@ -181,12 +217,12 @@ impl ExecutionContext { self.submission.take_recorded() } - /// The caller has proved completion, or explicitly assumes responsibility - /// for all submitted resources (the unsafe `async_on` escape hatch). + /// The caller has proved completion, or independently guarantees resource + /// lifetimes and ordering (the unsafe `async_on` escape hatch). pub(crate) unsafe fn complete(&self) { self.submission.complete(); } - pub fn device(&self) -> &Arc { + pub fn context(&self) -> &Arc { &self.device } pub fn get_device_id(&self) -> DeviceOrdinal { @@ -313,10 +349,16 @@ pub trait DeviceOp: /// Implementations must register owned resources with [`ExecutionContext::retain`] /// before submitting work. Returning them or borrowing them in the output /// is insufficient: safe callers may discard outputs or forget futures. + /// The output must be a valid Rust value immediately, including safe Drop: + /// combinators and cancelled futures may destroy it before GPU completion. + /// Any host memory still used by DMA needs an independent retained owner; + /// a borrowed host slice alone cannot support a forgettable submission. /// /// # Safety /// The caller must establish completion before host access to the output, - /// and must retain the context until completion or deliberately leak it. + /// and retain the context until completion, deliberately leak it, or + /// independently uphold the resource lifetime and ordering requirements + /// documented on [`async_on`](Self::async_on). unsafe fn execute( self, context: &ExecutionContext, @@ -324,7 +366,7 @@ pub trait DeviceOp: /// Schedule this operation on a specific policy and return a [`DeviceFuture`]. fn schedule( self, - policy: &Arc, + policy: &(impl SchedulingPolicy + ?Sized), ) -> Result::Output, Self>, DeviceError> { let stream = policy.next_stream()?; let mut future = DeviceFuture::new(); @@ -334,9 +376,9 @@ pub trait DeviceOp: } /// Chain a follow-up operation that runs **on the same stream** as `self`. /// - /// Because both operations share a stream, `f` is guaranteed to see `self`'s output - /// fully written. This is the recommended way to express data dependencies without - /// manual synchronization. + /// GPU work enqueued by the returned operation runs after the preceding + /// work. The host closure itself runs during submission, before completion; + /// it must not read device results directly. /// /// The closure must not execute other DeviceOps (e.g., via `sync_on` or `sync`). /// This is enforced at runtime by the thread-local execution lock — attempting @@ -352,6 +394,29 @@ pub trait DeviceOp: closure: f, } } + /// Compatibility spelling of [`then`](Self::then), using the same runtime. + fn and_then(self, f: F) -> AndThen<::Output, Self, O, DO, F> + where + DO: DeviceOp, + F: FnOnce(::Output) -> DO + Send, + { + self.then(f) + } + /// Compatibility spelling of [`then`](Self::then). + fn apply(self, f: F) -> AndThen<::Output, Self, O, DO, F> + where + DO: DeviceOp, + F: FnOnce(::Output) -> DO + Send, + { + self.then(f) + } + /// Wraps the output in an [`Arc`] through the shared map combinator. + fn arc(self) -> impl DeviceOp::Output>> + where + ::Output: Sync, + { + self.map(Arc::new) + } /// Like [`then`](DeviceOp::then), but the closure runs with the /// thread-local execution lock **released**, so it may execute other /// operations: `.sync()`, `.sync_on(&stream)`, or a nested `.await` @@ -492,17 +557,24 @@ pub trait DeviceOp: /// /// # Safety /// - /// The caller must ensure the stream is synchronized before accessing - /// GPU data from the output, and must keep every submitted resource alive - /// and prevent conflicting accesses until completion, including resources - /// discarded by combinators. Unlike safe execution, this method does not - /// retain resources after it returns. + /// The caller must establish GPU completion before host access to the + /// output, keep every submitted resource alive, and prevent conflicting + /// accesses until completion. This includes allocations, kernel modules, + /// graph executables, DMA buffers, and resources discarded by combinators. + /// These obligations also apply to work submitted before an error. + /// + /// Unlike safe future execution, this method releases registered owners + /// and access leases on return, even on error. It does not synchronize or + /// arrange background retention of those resources. unsafe fn async_on( self, stream: &Arc, ) -> Result<::Output, DeviceError> { let ctx = ExecutionContext::new(stream.clone()); let result = unsafe { self.execute(&ctx) }; + // SAFETY: the unsafe caller independently retains all resources and + // orders their accesses, including work submitted before an error. + // Release tracking now so context drop needs no completion probe/event. unsafe { ctx.complete() }; result } @@ -516,7 +588,7 @@ pub trait DeviceOp: let _execution_lock = acquire_execution_lock()?; let ctx = ExecutionContext::new(stream.clone()); let res = unsafe { self.execute(&ctx) }; - let sync_res = unsafe { stream.synchronize() }; + let sync_res = stream.synchronize(); sync_res?; unsafe { ctx.complete() }; res @@ -600,7 +672,7 @@ impl Producer { /// context's stream. fn record(ctx: &ExecutionContext) -> Result { let stream = Arc::clone(ctx.get_cuda_stream()); - let event = stream.device().new_event()?; + let event = stream.context().new_event()?; event.record(&stream)?; Ok(Self { stream, event }) } @@ -942,6 +1014,8 @@ impl>> DeviceOpUnwrapArc for DI // AndThen +// Send follows the stored operation and closure. A manual impl with FnOnce +// bounds over-constrains lifetime-bearing launch outputs in async blocks. pub struct AndThen where DI: DeviceOp, @@ -952,14 +1026,6 @@ where closure: F, } -unsafe impl Send for AndThen -where - DI: DeviceOp, - DO: DeviceOp, - F: FnOnce(I) -> DO + Send, -{ -} - impl DeviceOp for AndThen where DI: DeviceOp, @@ -1743,14 +1809,6 @@ where closure: F, } -unsafe impl Send for AndThenWithContext -where - DI: DeviceOp, - DO: DeviceOp, - F: FnOnce(&ExecutionContext, I) -> DO + Send, -{ -} - impl DeviceOp for AndThenWithContext where DI: DeviceOp, diff --git a/cuda-async/src/launch.rs b/cuda-async/src/launch.rs index d948cb0a85..43492af0d8 100644 --- a/cuda-async/src/launch.rs +++ b/cuda-async/src/launch.rs @@ -4,35 +4,133 @@ */ //! CUDA kernel launch builder with argument marshalling. +//! +//! [`AsyncKernelLaunchBuilder`] safely accumulates a kernel function reference, +//! launch attributes, and type-erased argument pointers. An explicit unsafe +//! call to [`AsyncKernelLaunchBuilder::finalize_unchecked`] supplies the raw +//! launch geometry and returns an immutable [`AsyncKernelLaunch`] that can be +//! scheduled as a [`DeviceOp`]. +//! +//! Arguments use an inline arena with a heap fallback and are freed when +//! the launcher is dropped after submission. This keeps the pointed-to values +//! alive until `cuLaunchKernel` / `cuLaunchKernelEx` has copied the launch +//! parameter values out of the host-side argument array. +//! +//! [`DeviceOp`]: crate::device_operation::DeviceOp use crate::device_context::with_default_device_policy; use crate::device_future::DeviceFuture; use crate::device_operation::{DeviceOp, ExecutionContext}; use crate::error::DeviceError; -use anyhow::{Context, Result}; -use cuda_core::sys::CUdeviceptr; -use cuda_core::{launch_kernel, DType, Function, LaunchConfig, Stream}; + +use cuda_core::CudaFunction; +use cuda_core::LaunchConfig; use std::ffi::c_void; use std::fmt::Debug; use std::future::IntoFuture; +use std::marker::PhantomData; use std::mem::MaybeUninit; use std::sync::Arc; -use std::vec::Vec; -/// A builder for asynchronously launching a CUDA kernel on a stream. +/// Safe, inert builder for an unchecked CUDA kernel launch. +/// +/// Building a launch does not enqueue GPU work. The builder deliberately does +/// not implement [`DeviceOp`] or [`IntoFuture`]: raw launch geometry +/// must first cross the explicit unsafe boundary at +/// [`finalize_unchecked`](Self::finalize_unchecked). +/// +/// ```compile_fail,E0277 +/// use cuda_async::device_operation::DeviceOp; +/// use cuda_async::launch::AsyncKernelLaunchBuilder; +/// +/// fn schedule(_operation: O) {} +/// +/// fn cannot_schedule(builder: AsyncKernelLaunchBuilder<'static>) { +/// schedule(builder); +/// } +/// ``` +/// +/// Builders stay on the thread that assembles their type-erased argument +/// storage. Only the immutable finalized operation is `Send`: +/// +/// ```compile_fail,E0277 +/// use cuda_async::launch::AsyncKernelLaunchBuilder; +/// +/// fn assert_send() {} +/// assert_send::>(); +/// ``` +#[derive(Debug)] +pub struct AsyncKernelLaunchBuilder<'a> { + /// Handle to the compiled device function. + func: Arc, + resources: Vec<(cuda_core::RawDeviceBuffer, bool)>, + /// Owned, type-erased argument storage passed to the CUDA driver. + args: KernelArgStorage, + configured: Option, + /// Optional thread-block cluster dimensions for `cuLaunchKernelEx`. + cluster_dim: Option<(u32, u32, u32)>, + /// When `true`, launch via `cuLaunchKernelEx` with + /// `CU_LAUNCH_ATTRIBUTE_COOPERATIVE` set, which guarantees the whole grid + /// is co-resident on the device (required for `cuda_device::grid::sync()`). + cooperative: bool, + programmatic_dependent_launch: bool, + /// Ties borrowed device buffers to the lazy launch operation. + _borrows: PhantomData<&'a mut ()>, +} + +/// Immutable, runnable CUDA kernel launch. +/// +/// This type can only be created by consuming an +/// [`AsyncKernelLaunchBuilder`]. It exposes no safe geometry or argument +/// setters, so its safety assumptions cannot be invalidated after the raw +/// launch configuration has been accepted. +/// +/// ```compile_fail,E0599 +/// use cuda_async::launch::AsyncKernelLaunch; +/// use cuda_core::LaunchConfig; /// -/// Arguments are copied into an inline arena by [`push_arg`](Self::push_arg) / -/// [`push_device_ptr`](Self::push_device_ptr) and handed to the driver as a -/// `*mut c_void` array at launch; the driver copies the parameter values out -/// before `cuLaunchKernel` returns, so nothing outlives the launch. +/// fn cannot_reconfigure(launch: &mut AsyncKernelLaunch<'_>, config: LaunchConfig) { +/// launch.set_launch_config(config); +/// } +/// ``` +/// Chaining finalized launches remains `Send`, including borrowed launches: +/// +/// ``` +/// use cuda_async::{device_operation::DeviceOp, launch::AsyncKernelLaunch}; +/// fn assert_send(_: impl Send) {} +/// fn check<'a>(first: AsyncKernelLaunch<'a>, second: AsyncKernelLaunch<'a>, third: AsyncKernelLaunch<'a>) { +/// assert_send(async move { +/// first.then(move |_| second).and_then_with_context(move |_, ()| third).await +/// }); +/// } +/// ``` #[derive(Debug)] -pub struct AsyncKernelLaunch { - pub func: Arc, +pub struct AsyncKernelLaunch<'a> { + /// Handle to the compiled device function. + func: Arc, + resources: Vec<(cuda_core::RawDeviceBuffer, bool)>, + /// Owned, type-erased argument storage passed to the CUDA driver. args: KernelArgStorage, - cfg: Option, + /// Complete grid/block dimensions and dynamic shared-memory size. + cfg: LaunchConfig, + /// Optional thread-block cluster dimensions for `cuLaunchKernelEx`. + cluster_dim: Option<(u32, u32, u32)>, + /// Whether this launch requests cooperative grid residency. + cooperative: bool, programmatic_dependent_launch: bool, + /// Ties borrowed device buffers to the lazy launch operation. + _borrows: PhantomData<&'a mut ()>, } +/// # Safety +/// +/// Argument values are owned by the launch. Non-`Copy` values enter only +/// through the explicit boxed argument method, which requires `Send`; +/// copied scalar values have no destructor. The runnable type exposes none of +/// those erased values to safe Rust. The `Arc` is `Send + Sync`, +/// and `_borrows` preserves source lifetimes captured while building. +unsafe impl<'a> Send for AsyncKernelLaunch<'a> {} + /// Inline capacity of the argument arena, in 16-byte value slots and in /// arguments. A tensor parameter takes `1 + 2 * rank` slots (pointer, shape, /// strides), so this covers four rank-3 tensors plus scalars; longer @@ -49,6 +147,11 @@ const INLINE_SLOTS: usize = 32; /// materialized only at launch, after every push, so growth can never /// invalidate a recorded pointer. enum KernelArgStorage { + Mixed { + plain: Box, + boxed: Vec, + order: Vec, + }, Inline { /// 16-byte-aligned value slots; the first `values_len` are occupied. A /// padded value leaves its padding bytes uninitialized. @@ -87,6 +190,7 @@ impl Debug for KernelArgStorage { } => d .field("len", args_len) .field("offsets", &&offsets[..*args_len]), + Self::Mixed { order, .. } => d.field("len", &order.len()), Self::Heap { offsets, .. } => d.field("len", &offsets.len()).field("offsets", offsets), }; d.finish() @@ -96,10 +200,17 @@ impl Debug for KernelArgStorage { impl KernelArgStorage { /// Copies `arg` into the arena. `Copy` makes the missing destructor /// bookkeeping sound: dropping the arena drops nothing. - fn push(&mut self, arg: T) { + fn push(&mut self, arg: T) { const SLOT: usize = std::mem::size_of::(); - const { - assert!(std::mem::align_of::() <= SLOT); + if std::mem::align_of::() > SLOT { + self.push_box(Box::new(arg)); + return; + } + if let Self::Mixed { plain, order, .. } = self { + let index = plain.len(); + plain.push(arg); + order.push(ArgumentIndex::Plain(index)); + return; } let slots = std::mem::size_of::().div_ceil(SLOT).max(1); if let Self::Inline { @@ -117,7 +228,7 @@ impl KernelArgStorage { // SAFETY: the pointer derives from the whole array, so a value // spanning several slots stays inside the borrowed range; the // span is at least `size_of::()` bytes and 16-byte aligned, - // which the const assertion above bounds `T`'s alignment by; + // which the alignment guard above bounds `T`'s alignment by; // `T: Copy` so overwriting drops nothing. unsafe { std::ptr::write(values.as_mut_ptr().add(offset).cast::(), arg) }; offsets[*args_len] = offset as u8; @@ -153,6 +264,21 @@ impl KernelArgStorage { /// nothing. fn with_param_ptrs(&mut self, f: impl FnOnce(&mut [*mut c_void]) -> R) -> R { match self { + Self::Mixed { + plain, + boxed, + order, + } => { + let ptrs = plain.param_vec(); + let mut all: Vec<*mut c_void> = order + .iter() + .map(|index| match index { + ArgumentIndex::Plain(i) => ptrs[*i], + ArgumentIndex::Boxed(i) => boxed[*i].ptr, + }) + .collect(); + f(&mut all) + } Self::Inline { values, offsets, @@ -194,81 +320,281 @@ impl KernelArgStorage { } } -impl AsyncKernelLaunch { - /// Creates a new kernel launch builder for the given CUDA function. - pub fn new(func: Arc) -> AsyncKernelLaunch { - AsyncKernelLaunch { +#[derive(Debug)] +struct ErasedArgument { + ptr: *mut c_void, + drop_value: unsafe fn(*mut c_void), +} +impl Drop for ErasedArgument { + fn drop(&mut self) { + unsafe { (self.drop_value)(self.ptr) }; + } +} +enum ArgumentIndex { + Plain(usize), + Boxed(usize), +} +impl KernelArgStorage { + fn param_vec(&mut self) -> Vec<*mut c_void> { + self.with_param_ptrs(|p| p.to_vec()) + } + fn len(&self) -> usize { + match self { + Self::Inline { args_len, .. } => *args_len, + Self::Heap { offsets, .. } => offsets.len(), + Self::Mixed { order, .. } => order.len(), + } + } + fn push_box(&mut self, value: Box) { + unsafe fn drop_box(ptr: *mut c_void) { + drop(unsafe { Box::from_raw(ptr.cast::()) }); + } + if !matches!(self, Self::Mixed { .. }) { + let order = (0..self.len()).map(ArgumentIndex::Plain).collect(); + let plain = Box::new(std::mem::take(self)); + *self = Self::Mixed { + plain, + boxed: Vec::new(), + order, + }; + } + let Self::Mixed { boxed, order, .. } = self else { + unreachable!() + }; + order.push(ArgumentIndex::Boxed(boxed.len())); + boxed.push(ErasedArgument { + ptr: Box::into_raw(value).cast(), + drop_value: drop_box::, + }); + } +} + +impl<'a> AsyncKernelLaunchBuilder<'a> { + /// Creates an inert builder for `func` with no arguments or launch + /// attributes. + pub fn new(func: Arc) -> Self { + Self { func, + resources: Vec::new(), args: KernelArgStorage::default(), - cfg: None, + cluster_dim: None, + cooperative: false, programmatic_dependent_launch: false, + configured: None, + _borrows: PhantomData, } } - /// Pushes a kernel argument by value. - #[inline(always)] - pub fn push_arg(&mut self, arg: T) -> &mut Self { - arg.push_arg(self); + /// Push a raw CUDA address. Its pointee must remain valid through completion. + /// # Safety + /// The address, extent, element type and permissions must match the kernel. + pub unsafe fn push_device_ptr(&mut self, ptr: cuda_core::sys::CUdeviceptr) -> &mut Self { + self.push_scalar_arg(ptr) + } + pub fn push_arg_arc(&mut self, arg: &Arc) -> &mut Self { + arg.push_arg_arc(self); + self + } + /// Sets the geometry on this inert builder. It cannot be scheduled. + pub fn set_launch_config(&mut self, config: LaunchConfig) -> &mut Self { + self.configured = Some(config); + self + } + /// Executes a configured builder using the same sealed launch operation. + /// # Safety + /// All requirements of `finalize_unchecked` apply to the configured launch. + pub unsafe fn execute(self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + let config = self + .configured + .ok_or_else(|| DeviceError::Launch("missing launch configuration".into()))?; + unsafe { self.finalize_unchecked(config).execute(ctx) } + } + /// Enable programmatic dependent launch. + /// # Safety + /// The kernel must perform the required dependency wait before dependent + /// accesses; neither kernel may require overlap. Retain all used resources. + pub unsafe fn programmatic_dependent_launch(&mut self) -> &mut Self { + self.programmatic_dependent_launch = true; self } + /// Pushes an explicitly typed, heap-owned argument without a DeviceCopy bound. + /// # Safety + /// The object's layout and initialized representation must match the kernel + /// parameter. It must not carry host ownership the device can dereference. + pub unsafe fn push_boxed_arg_unchecked(&mut self, arg: Box) -> &mut Self { + self.args.push_box(arg); + self + } + + /// Retains a shared allocation and records its access mode for submission. + pub fn retain_buffer(&mut self, storage: cuda_core::RawDeviceBuffer, write: bool) { + self.resources.push((storage, write)); + } - /// Pushes a kernel argument from an `Arc` reference. + /// Appends a by-value `Copy` argument to the kernel launch packet. + /// + /// This is the typed-module path used for scalar, raw-pointer, custom + /// `Copy` structs, and `Copy` closure arguments. #[inline(always)] - pub fn push_arg_arc(&mut self, arg: &Arc) -> &mut Self { - arg.push_arg_arc(self); + pub fn push_scalar_arg(&mut self, arg: T) -> &mut Self { + self.args.push(arg); self } - /// Pushes a device pointer as a kernel argument. + /// Appends a kernel argument to owned storage for the driver call. + /// Common scalar arguments stay in the inline arena; large argument lists + /// and over-aligned types use heap storage. /// - /// # Safety - /// `ptr` must stay a valid device allocation, on the device that owns the - /// stream this launch is executed on, until the kernel has completed — not - /// merely until the launch is submitted. The kernel signature must expect a - /// pointer at this position. - pub unsafe fn push_device_ptr(&mut self, ptr: CUdeviceptr) -> &mut Self { - self.push_arg_raw(ptr) + /// The allocated storage remains alive until launch submission finishes or + /// the builder is dropped without launching. + /// + /// Values with thread-affine ownership cannot be hidden in the erased + /// storage that will later become a `Send` operation: + /// + /// ```compile_fail,E0277 + /// use cuda_async::launch::AsyncKernelLaunchBuilder; + /// use std::rc::Rc; + /// + /// fn reject_rc(mut builder: AsyncKernelLaunchBuilder<'static>) { + /// builder.push_arg(Box::new(Rc::new(1_u32))); + /// } + /// ``` + #[inline(always)] + pub fn push_arg(&mut self, arg: T) -> &mut Self { + arg.push_arg(self); + self } - /// Pushes a raw argument to the kernel parameter list. + /// Appends multiple kernel arguments at once from a tuple. /// - /// # Safety - /// `T` must match the size and alignment of the kernel's formal parameter - /// at this position; the driver copies `size_of::()` bytes from the - /// stored value. - unsafe fn push_arg_raw(&mut self, arg: T) -> &mut Self { - self.args.push(arg); + /// Equivalent to chained [`push_arg`](Self::push_arg) calls but allows + /// grouping all arguments in a single expression: + /// + /// ```ignore + /// launch.push_args((m, n, k, alpha, a_ptr, a_len, b_ptr, b_len, beta, c_ptr, c_len)); + /// ``` + /// + /// Supports tuples up to 32 elements. + #[inline(always)] + pub fn push_args(&mut self, args: T) -> &mut Self { + args.push_args(self); self } - /// Sets the grid/block dimensions and shared memory configuration for the launch. - pub fn set_launch_config(&mut self, cfg: LaunchConfig) -> &mut Self { - self.cfg = Some(cfg); + /// Sets thread-block cluster dimensions for a cluster launch. + pub fn set_cluster_dim(&mut self, cluster_dim: (u32, u32, u32)) -> &mut Self { + self.cluster_dim = Some(cluster_dim); self } - /// Enable programmatic dependent launch for this submission. Off by default. + /// Marks this launch as cooperative (`CU_LAUNCH_ATTRIBUTE_COOPERATIVE`). /// - /// # Safety - /// The kernel must wait for predecessor completion before dependent memory - /// accesses (token-ordered after `gdc_wait_tko` for Tile kernels). Work before - /// the wait must not race predecessor work. Neither kernel may depend on - /// overlap, and all resources must remain valid through their last use. - pub unsafe fn programmatic_dependent_launch(&mut self) -> &mut Self { - self.programmatic_dependent_launch = true; + /// A cooperative launch guarantees every block in the grid is co-resident + /// on the device, which is the precondition for grid-wide barriers like + /// `cuda_device::grid::sync()`. May be combined with + /// [`set_cluster_dim`](Self::set_cluster_dim); both attributes are then + /// passed to `cuLaunchKernelEx` in the same call. + pub fn set_cooperative(&mut self, cooperative: bool) -> &mut Self { + self.cooperative = cooperative; self } - /// Launches the kernel on the given CUDA stream. + /// Accepts an unverified raw launch configuration and seals this builder + /// into an immutable, schedulable operation. + /// + /// Prefer the generated typed module APIs for application code. They check + /// a kernel's declared launch contract before constructing the runnable + /// operation. + /// + /// # Safety + /// + /// The caller must prove all facts that a typed prepared launch normally + /// checks: + /// + /// - `cfg` must match the kernel's indexing dimensionality, launch bounds, + /// block shape, and dynamic shared-memory requirements; + /// - the cluster and cooperative attributes must be legal for the kernel; + /// - the function signature must exactly match the arguments, in order, + /// type, size, and alignment; + /// - referenced device allocations must remain valid for the complete GPU + /// operation and satisfy Rust's aliasing requirements; and + /// - the scheduling policy must select a stream from the CUDA context that + /// owns `func`. + /// + /// Violating these requirements can cause memory corruption, data races, + /// or CUDA driver errors. + /// + /// ```compile_fail,E0133 + /// use cuda_async::launch::AsyncKernelLaunchBuilder; + /// use cuda_core::LaunchConfig; + /// + /// fn raw_finalize_requires_unsafe( + /// builder: AsyncKernelLaunchBuilder<'static>, + /// config: LaunchConfig, + /// ) { + /// let _ = builder.finalize_unchecked(config); + /// } + /// ``` + pub unsafe fn finalize_unchecked(self, cfg: LaunchConfig) -> AsyncKernelLaunch<'a> { + AsyncKernelLaunch { + func: self.func, + resources: self.resources, + args: self.args, + cfg, + cluster_dim: self.cluster_dim, + cooperative: self.cooperative, + programmatic_dependent_launch: self.programmatic_dependent_launch, + _borrows: self._borrows, + } + } +} + +impl<'a> AsyncKernelLaunch<'a> { + /// Submits the kernel to `stream` via `cuLaunchKernel`, or via + /// `cuLaunchKernelEx` when cluster dimensions or a cooperative launch + /// were requested. /// /// # Safety - /// The caller must ensure the kernel arguments and launch config are valid. - unsafe fn launch(mut self, stream: &Arc) -> Result<(), DeviceError> { - let cfg = self.cfg.ok_or_else(|| { - DeviceError::Launch("Await called before launching the kernel.".to_string()) - })?; - let launch = if self.programmatic_dependent_launch { - let architecture = cuda_core::get_device_sm_name(stream.device().cu_device())?; + /// + /// - `self.func` must refer to a kernel loaded from the same CUDA context + /// that owns `stream`. + /// - All argument pointers in `self.args` must point to correctly typed and + /// aligned host-side values for the kernel's formal parameters. + /// - The pointed-to argument values must remain valid until launch + /// submission returns. The stream-aware launch helper binds the correct + /// context before calling into the CUDA driver. + /// + /// # Errors + /// + /// Returns [`DeviceError::Driver`] if context binding or launch submission + /// fails. + unsafe fn launch(mut self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + let stream = ctx.get_cuda_stream(); + if self.func.context().as_ref() != stream.context().as_ref() { + return Err(DeviceError::Driver(cuda_core::DriverError( + cuda_core::sys::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + ))); + } + if self.programmatic_dependent_launch && (self.cluster_dim.is_some() || self.cooperative) { + return Err(DeviceError::Launch( + "programmatic dependent launch cannot be combined with cluster/cooperative mode" + .into(), + )); + } + stream.context().bind_to_thread()?; + // Register before enqueue: partial failures, unwind and forgotten futures + // must not shorten the lifetime of any argument or module. + ctx.retain(self.func.clone())?; + for (storage, write) in &self.resources { + ctx.retain_buffer(storage.clone(), *write)?; + } + let cfg = self.cfg; + let function = self.func.clone(); + let cluster = self.cluster_dim; + let cooperative = self.cooperative; + let pdl = self.programmatic_dependent_launch; + if pdl { + let architecture = cuda_core::get_device_sm_name(stream.context().cu_device())?; let sm: u32 = architecture .strip_prefix("sm_") .and_then(|s| s.trim_end_matches(['a', 'f']).parse().ok()) @@ -278,91 +604,209 @@ impl AsyncKernelLaunch { "programmatic dependent launch requires sm_90 or newer; target {architecture}" ))); } - cuda_core::launch_kernel_pdl - } else { - launch_kernel - }; - let func = self.func.cu_function(); + } self.args - .with_param_ptrs(|params| { - launch( - func, - cfg.grid_dim, - cfg.block_dim, - cfg.shared_mem_bytes, - stream.cu_stream(), - params, - ) + .with_param_ptrs(|params| unsafe { + match (cluster, cooperative, pdl) { + (_, _, true) => cuda_core::launch_kernel_pdl( + function.cu_function(), + cfg.grid_dim, + cfg.block_dim, + cfg.shared_mem_bytes, + stream.cu_stream(), + params, + ), + (Some(cluster), true, false) => { + cuda_core::launch_kernel_ex_cooperative_on_stream( + &function, + cfg.grid_dim, + cfg.block_dim, + cfg.shared_mem_bytes, + cluster, + stream, + params, + ) + } + (Some(cluster), false, false) => cuda_core::launch_kernel_ex_on_stream( + &function, + cfg.grid_dim, + cfg.block_dim, + cfg.shared_mem_bytes, + cluster, + stream, + params, + ), + (None, true, false) => cuda_core::launch_kernel_cooperative_on_stream( + &function, + cfg.grid_dim, + cfg.block_dim, + cfg.shared_mem_bytes, + stream, + params, + ), + (None, false, false) => cuda_core::launch_kernel_on_stream( + &function, + cfg.grid_dim, + cfg.block_dim, + cfg.shared_mem_bytes, + stream, + params, + ), + } }) - .with_context(|| { - format!( - r#" - Failed to launch kernel. - args: {:#?} - cfg: {:#?}"#, - self.args, cfg - ) - })?; - Ok(()) + .map_err(DeviceError::Driver) } } -/// A kernel argument that can be pushed from an `Arc` reference. -pub trait ArcKernelArgument { - // #[inline(always)] Dont think this is necessary. This will be deprecated for required trait methods - fn push_arg_arc(self: &Arc, launcher: &mut AsyncKernelLaunch); +/// Returns caller-owned values from a lazy async kernel launch. GPU-visible +/// storage must also be retained in the launch's descriptors: host combinators +/// and cancellation may drop the returned values before GPU completion. +#[derive(Debug)] +pub struct OwnedAsyncKernelLaunch { + launch: AsyncKernelLaunch<'static>, + resources: R, +} + +impl OwnedAsyncKernelLaunch { + /// Associates returned values with a launch that already retains its GPU resources. + pub fn new(launch: AsyncKernelLaunch<'static>, resources: R) -> Self { + Self { launch, resources } + } } -/// A kernel argument that can be pushed by value into an `AsyncKernelLaunch`. +/// A by-value kernel argument encoder. Geometry and signature validation belong +/// to the finalized launch, not to this inert builder. pub trait KernelArgument { - // #[inline(always)] Dont think this is necessary. This will be deprecated for required trait methods - fn push_arg(self, launcher: &mut AsyncKernelLaunch); + fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>); +} +/// All plain device-copyable values use the same marshalling contract in both +/// programming models. Tile dtype validation is performed by the Tile frontend. +impl KernelArgument for T { + fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { + launcher.args.push(self); + } +} +pub trait ArcKernelArgument { + fn push_arg_arc(self: &Arc, launcher: &mut AsyncKernelLaunchBuilder<'_>); } -/// Safe implementation for scalar types. Values implementing `DType` are copied -/// into the kernel's parameter space during launch — the kernel reads the value, -/// not a device pointer, so no `unsafe` is required. -impl KernelArgument for T { - fn push_arg(self, launcher: &mut AsyncKernelLaunch) { - // SAFETY: a `DType` scalar is a plain `Copy` value with the layout the - // compiled kernel declares for that scalar type; the launcher's - // signature validation (in cutile) checks the position, and the value - // is copied out by the driver at launch, so nothing outlives the box. - unsafe { - launcher.push_arg_raw(self); +// --------------------------------------------------------------------------- +// KernelArguments — push multiple heterogeneous args in a single call +// --------------------------------------------------------------------------- + +/// Trait for a group of kernel arguments that can be pushed together. +/// +/// Implemented for tuples of [`KernelArgument`] types up to arity 32, enabling +/// `launch.push_args((dim_m, dim_n, alpha, ptr, len))` as an alternative to +/// chained `.push_arg()` calls. +#[diagnostic::on_unimplemented( + message = "cannot push `{Self}` as kernel arguments", + note = "KernelArguments is implemented for tuples of KernelArgument types up to 32 elements" +)] +pub trait KernelArguments { + /// Pushes every element into `launcher` in order. + fn push_args(self, launcher: &mut AsyncKernelLaunchBuilder<'_>); +} + +macro_rules! impl_kernel_args_tuple { + // Base case: empty tuple + () => { + impl KernelArguments for () { + #[inline(always)] + fn push_args(self, _launcher: &mut AsyncKernelLaunchBuilder<'_>) {} } - } + }; + // Recursive case: (A, B, C, ...) where each element is a KernelArgument + ($($idx:tt : $T:ident),+) => { + impl<$($T: KernelArgument),+> KernelArguments for ($($T,)+) { + #[inline(always)] + fn push_args(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { + $(launcher.push_arg(self.$idx);)+ + } + } + }; } -impl DeviceOp for AsyncKernelLaunch { +impl_kernel_args_tuple!(); +impl_kernel_args_tuple!(0: A); +impl_kernel_args_tuple!(0: A, 1: B); +impl_kernel_args_tuple!(0: A, 1: B, 2: C); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD, 30: AE); +impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD, 30: AE, 31: AF); + +/// Launches the kernel on the stream bound to the execution context. +impl<'a> DeviceOp for AsyncKernelLaunch<'a> { type Output = (); - unsafe fn execute( - self, - ctx: &ExecutionContext, - ) -> Result<::Output, DeviceError> { - self.launch(ctx.get_cuda_stream()) + unsafe fn execute(self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + unsafe { self.launch(ctx) } } } -impl IntoFuture for AsyncKernelLaunch { +/// Schedules the kernel launch via the thread-local default scheduling policy. +impl<'a> IntoFuture for AsyncKernelLaunch<'a> { type Output = Result<(), DeviceError>; - type IntoFuture = DeviceFuture<(), AsyncKernelLaunch>; + type IntoFuture = DeviceFuture<(), AsyncKernelLaunch<'a>>; + fn into_future(self) -> Self::IntoFuture { + match with_default_device_policy(|policy| self.schedule(policy)) { + Ok(Ok(future)) => future, + Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), + } + } +} + +impl DeviceOp for OwnedAsyncKernelLaunch { + type Output = R; + + unsafe fn execute(self, ctx: &ExecutionContext) -> Result { + let Self { launch, resources } = self; + unsafe { launch.launch(ctx) }?; + Ok(resources) + } +} + +impl IntoFuture for OwnedAsyncKernelLaunch { + type Output = Result; + type IntoFuture = DeviceFuture>; + fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| { - let stream = policy.next_stream()?; - let mut f = DeviceFuture::new(); - f.device_operation = Some(self); - f.execution_context = Some(ExecutionContext::new(stream)); - Ok(f) - }) { + match with_default_device_policy(|policy| self.schedule(policy)) { Ok(Ok(future)) => future, - Ok(Err(e)) => DeviceFuture::failed(e), - Err(e) => DeviceFuture::failed(e), + Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), } } } +impl crate::device_operation::GraphNode for AsyncKernelLaunch<'_> {} +impl crate::device_operation::GraphNode for OwnedAsyncKernelLaunch {} + #[cfg(test)] mod arg_storage_tests { //! Host-only: the storage never touches the driver. @@ -477,8 +921,7 @@ mod arg_storage_tests { } /// Every parameter pointer is aligned for its slot (16 bytes), which - /// bounds all accepted argument types; the `const` assertion in `push` - /// rejects wider alignments at compile time. + /// covers ordinary arguments. Wider alignments use typed heap storage. #[test] fn slots_are_sixteen_byte_aligned() { let mut storage = KernelArgStorage::default(); @@ -490,3 +933,44 @@ mod arg_storage_tests { }); } } + +#[cfg(test)] +mod extra_argument_tests { + use super::*; + #[test] + fn over_aligned_values_keep_order_across_plain_arena_growth() { + #[derive(Clone, Copy)] + #[repr(C, align(256))] + struct Aligned([u8; 256]); + let mut storage = KernelArgStorage::default(); + storage.push(7u32); + storage.push(Aligned([19; 256])); + for i in 0..64u64 { + storage.push(i); + } + storage.with_param_ptrs(|args| unsafe { + assert_eq!(args.len(), 66); + assert_eq!(*(args[0] as *const u32), 7); + assert_eq!(args[1] as usize % 256, 0); + assert_eq!((*(args[1] as *const Aligned)).0, [19; 256]); + for i in 0..64 { + assert_eq!(*(args[i + 2] as *const u64), i as u64); + } + }); + } + #[test] + fn erased_arguments_drop_with_their_original_type() { + use std::sync::atomic::{AtomicUsize, Ordering}; + struct Probe(Arc); + impl Drop for Probe { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + let drops = Arc::new(AtomicUsize::new(0)); + let mut storage = KernelArgStorage::default(); + storage.push_box(Box::new(Probe(drops.clone()))); + drop(storage); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } +} diff --git a/cuda-async/src/leak.rs b/cuda-async/src/leak.rs index 7dda508d21..4464980fa4 100644 --- a/cuda-async/src/leak.rs +++ b/cuda-async/src/leak.rs @@ -13,58 +13,8 @@ use std::io::{self, Write}; -/// Reports a leak to stderr — or to the active test capture, if one is -/// installed. +/// Reports a leak to stderr. pub(crate) fn report_leak(message: std::fmt::Arguments<'_>) { - #[cfg(test)] - if capture::push(&message) { - return; - } let mut stderr = io::stderr().lock(); let _ = writeln!(stderr, "{message}"); } - -/// Test-only capture of leak reports, so tests that intentionally exercise -/// the loud-leak paths can assert on the message instead of printing it. -#[cfg(test)] -pub(crate) mod capture { - use std::sync::{Mutex, MutexGuard}; - - static ACTIVE: Mutex<()> = Mutex::new(()); - static CAPTURED: Mutex>> = Mutex::new(None); - - /// Serializes capturing tests and scopes the capture window; captured - /// messages are retrieved (and the window closed) by [`Capture::take`]. - pub(crate) struct Capture { - _serial: MutexGuard<'static, ()>, - } - - pub(crate) fn start() -> Capture { - let serial = ACTIVE.lock().unwrap_or_else(|e| e.into_inner()); - *CAPTURED.lock().unwrap() = Some(Vec::new()); - Capture { _serial: serial } - } - - impl Capture { - pub(crate) fn take(&mut self) -> Vec { - CAPTURED.lock().unwrap().take().unwrap_or_default() - } - } - - impl Drop for Capture { - fn drop(&mut self) { - let _ = CAPTURED.lock().map(|mut c| c.take()); - } - } - - pub(crate) fn push(message: &std::fmt::Arguments<'_>) -> bool { - let mut captured = CAPTURED.lock().unwrap(); - match captured.as_mut() { - Some(messages) => { - messages.push(message.to_string()); - true - } - None => false, - } - } -} diff --git a/cuda-async/src/lib.rs b/cuda-async/src/lib.rs index 8f61f510a6..b22c482f57 100644 --- a/cuda-async/src/lib.rs +++ b/cuda-async/src/lib.rs @@ -7,6 +7,7 @@ //! and device memory management. pub mod cuda_graph; +pub mod device_box; pub mod device_buffer; pub mod device_context; pub mod device_future; @@ -23,9 +24,6 @@ pub mod prelude; #[cfg(not(loom))] mod reactor; pub mod scheduling_policies; -/// SIMT-model async surface, copied from cuda-oxide for the shared -/// host-crate migration. Not re-exported at the root; see the module docs. -pub mod simt; mod slot_table; mod submission; diff --git a/cuda-async/src/reactor.rs b/cuda-async/src/reactor.rs index 3a29ab208b..30c9501b9c 100644 --- a/cuda-async/src/reactor.rs +++ b/cuda-async/src/reactor.rs @@ -216,7 +216,7 @@ fn probe_stale_slots( /// slot's flag write can no longer land — the error is the context's sticky /// fault — so recycling the slot cannot be clobbered by a late device write. fn stream_is_faulted(stream: &Stream) -> bool { - if stream.device().bind_to_thread().is_err() { + if stream.context().bind_to_thread().is_err() { return false; } matches!(probe_stream(stream), StreamHealth::Faulted(_)) diff --git a/cuda-async/src/scheduling_policies.rs b/cuda-async/src/scheduling_policies.rs index 5bcc0cac79..51d4ce243e 100644 --- a/cuda-async/src/scheduling_policies.rs +++ b/cuda-async/src/scheduling_policies.rs @@ -6,7 +6,7 @@ //! Stream scheduling policies that control how operations are assigned to CUDA streams. use crate::error::DeviceError; -use cuda_core::{Device, Stream}; +use cuda_core::{Context, Stream}; use std::sync::atomic::AtomicUsize; use std::sync::Arc; @@ -27,6 +27,12 @@ pub trait SchedulingPolicy: Send + Sync { fn next_stream(&self) -> Result, DeviceError>; } +impl SchedulingPolicy for Arc

{ + fn next_stream(&self) -> Result, DeviceError> { + self.as_ref().next_stream() + } +} + /// Distributes operations across a fixed-size pool of CUDA streams using round-robin assignment. /// /// This is the **default scheduling policy**. Each call to `next_stream()` picks the @@ -39,7 +45,7 @@ pub struct StreamPoolRoundRobin { impl StreamPoolRoundRobin { /// Creates a round-robin pool with `num_streams` streams on the given device. - pub fn new(device: &Arc, num_streams: usize) -> Result { + pub fn new(device: &Arc, num_streams: usize) -> Result { let mut stream_pool = Vec::with_capacity(num_streams); for _ in 0..num_streams { stream_pool.push(device.new_stream()?); @@ -67,8 +73,13 @@ pub struct SingleStream { } impl SingleStream { + /// Uses an existing shared CUDA stream. + pub fn from_stream(stream: Arc) -> Self { + Self { stream } + } + /// Creates a single-stream policy on the given device. - pub fn new(device: &Arc) -> Result { + pub fn new(device: &Arc) -> Result { Ok(Self { stream: device.new_stream()?, }) diff --git a/cuda-async/src/simt/device_box.rs b/cuda-async/src/simt/device_box.rs deleted file mode 100644 index 941f201778..0000000000 --- a/cuda-async/src/simt/device_box.rs +++ /dev/null @@ -1,164 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Owned and borrowed wrappers around CUDA device pointers. -//! -//! [`DevicePointer`] is a thin, `Copy` handle to device memory suitable for -//! passing as a kernel argument. [`DeviceBox`] owns device memory and frees it -//! asynchronously through a dedicated deallocator stream on drop. -//! -//! # Async deallocation -//! -//! [`DeviceBox`] enqueues `cuMemFreeAsync` on a per-device deallocator stream -//! rather than blocking with `cuMemFree`. Callers **must** synchronize all -//! streams that reference the allocation before dropping the box. - -use crate::simt::device_context::with_deallocator_stream; -use crate::simt::error::DeviceError; -use crate::simt::launch::{AsyncKernelLaunchBuilder, KernelArgument}; -use cuda_bindings::CUdeviceptr; -use cuda_core::simt::memory::free_async; -use std::io::{self, Write}; -use std::marker::PhantomData; - -/// Non-owning, `Copy` handle to a typed device allocation. -/// -/// Does not manage lifetime -- the underlying memory must outlive all uses. -#[derive(Debug, Copy, Clone)] -pub struct DevicePointer { - /// Ties the pointer to element type `T` without owning one. - _marker: PhantomData, - /// Raw CUDA device pointer. - pub dptr: CUdeviceptr, -} - -/// # Safety -/// -/// `DevicePointer` is a plain integer handle with no thread-affine state. -/// Sending it across threads does not violate CUDA's driver API contract. -unsafe impl Send for DevicePointer {} - -impl DevicePointer { - /// Returns the raw [`CUdeviceptr`]. - pub fn cu_deviceptr(&self) -> CUdeviceptr { - self.dptr - } -} - -/// Pushes the underlying device address as a kernel argument. -impl KernelArgument for DevicePointer { - fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { - launcher.push_arg(self.cu_deviceptr()); - } -} - -/// Owning wrapper around a device allocation that frees memory asynchronously -/// on drop. -/// -/// `DeviceBox<[DType]>` represents a contiguous device buffer of `len` -/// elements. The allocation is freed via `cuMemFreeAsync` on the device's -/// deallocator stream when the box is dropped. -/// -/// # Safety contract -/// -/// All streams that reference this allocation **must** be synchronized before -/// the box is dropped. Violating this causes use-after-free on the device. -#[derive(Debug)] -pub struct DeviceBox { - /// Ordinal of the device that owns this allocation. - device_id: usize, - /// Raw CUDA device pointer to the start of the allocation. - cudptr: CUdeviceptr, - /// Number of elements (not bytes) in the allocation. - len: usize, - /// Ties the box to the element type. - _marker: PhantomData, -} - -/// # Safety -/// -/// Device pointers are plain integers. The CUDA driver API is thread-safe for -/// distinct streams, and `DeviceBox` holds no thread-affine state. -unsafe impl Send for DeviceBox {} - -/// # Safety -/// -/// See the [`Send`] impl. Shared references only expose the device pointer -/// value, which is safe to read from any thread. -unsafe impl Sync for DeviceBox {} - -/// Best-effort enqueue of an asynchronous free on the device's deallocator -/// stream. -/// -/// Drop cannot return errors, so failures to access the deallocator stream or -/// enqueue the free are reported to stderr instead of panicking. -impl Drop for DeviceBox { - fn drop(&mut self) { - let result = unsafe { - with_deallocator_stream(self.device_id, |stream| { - free_async(self.cudptr, stream.cu_stream()).map_err(DeviceError::Driver) - }) - }; - - match result { - Ok(Ok(())) => {} - Ok(Err(err)) | Err(err) => { - let mut stderr = io::stderr().lock(); - let _ = writeln!( - stderr, - "cuda-async: failed to enqueue async free for device pointer on device_id={}: {}", - self.device_id, err - ); - } - } - } -} - -impl DeviceBox<[DType]> { - /// Constructs a `DeviceBox<[DType]>` from a raw device pointer, element - /// count, and device ordinal. - /// - /// # Safety - /// - /// * `dptr` must point to a valid device allocation of at least - /// `len * size_of::()` bytes on device `device_id`. - /// * Ownership of the allocation transfers to the returned `DeviceBox`. - pub unsafe fn from_raw_parts(dptr: CUdeviceptr, len: usize, device_id: usize) -> Self { - Self { - _marker: PhantomData, - cudptr: dptr, - len, - device_id, - } - } - - /// Returns `true` if the buffer contains zero elements. - pub fn is_empty(&self) -> bool { - self.len == 0 - } - - /// Returns the number of elements in the buffer. - pub fn len(&self) -> usize { - self.len - } - - /// Returns the raw [`CUdeviceptr`]. - pub fn cu_deviceptr(&self) -> CUdeviceptr { - self.cudptr - } - - /// Returns the ordinal of the device that owns this allocation. - pub fn device_id(&self) -> usize { - self.device_id - } - - /// Creates a non-owning [`DevicePointer`] into this allocation. - pub fn device_pointer(&self) -> DevicePointer { - DevicePointer { - _marker: PhantomData, - dptr: self.cudptr, - } - } -} diff --git a/cuda-async/src/simt/device_context.rs b/cuda-async/src/simt/device_context.rs deleted file mode 100644 index 29bf86b615..0000000000 --- a/cuda-async/src/simt/device_context.rs +++ /dev/null @@ -1,361 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Thread-local GPU device state, kernel cache, and scheduling policy management. -//! -//! Each thread maintains a set of [`AsyncDeviceContext`]s (one per device) via -//! the `DEVICE_CONTEXTS` thread-local. A context bundles: -//! -//! * A [`CudaContext`] for driver API calls. -//! * A [`GlobalSchedulingPolicy`] for stream selection. -//! * A dedicated deallocator stream for async `cuMemFreeAsync`. -//! * A kernel function cache keyed by [`FunctionKey`] hashes. -//! -//! Most users interact through the convenience functions -//! ([`with_default_device_policy`], [`with_cuda_context`], etc.) rather than -//! touching the thread-local directly. -//! -//! [`CudaContext`]: cuda_core::CudaContext - -use crate::simt::error::{device_assert, device_error, DeviceError}; -use crate::simt::scheduling_policies::{ - GlobalSchedulingPolicy, SchedulingPolicy, StreamPoolRoundRobin, -}; -use cuda_core::{CudaContext, CudaFunction, CudaModule, CudaStream}; -use rustc_hash::FxHashMap; -use std::cell::Cell; -use std::hash::{DefaultHasher, Hash, Hasher}; -use std::sync::Arc; - -/// Default CUDA device ordinal used when no explicit device is specified. -pub const DEFAULT_DEVICE_ID: usize = 0; - -/// Default number of devices initialized by `init_device_contexts_default`. -pub const DEFAULT_NUM_DEVICES: usize = 1; - -/// Default number of streams in the [`StreamPoolRoundRobin`] policy. -pub const DEFAULT_ROUND_ROBIN_STREAM_POOL_SIZE: usize = 4; - -/// Trait for types that uniquely identify a compiled kernel function. -/// -/// The hash is used as the lookup key in the per-device function cache. -pub trait FunctionKey: Hash { - /// Returns a hex-encoded hash string suitable for use as a cache key. - fn get_hash_string(&self) -> String { - let mut hasher = DefaultHasher::new(); - self.hash(&mut hasher); - format!("{:x}", hasher.finish()) - } -} - -/// Cached mapping from function hash keys to loaded `(module, function)` pairs. -type DeviceFunctions = FxHashMap, Arc)>; - -/// Per-device state: CUDA context, scheduling policy, deallocator stream, and -/// compiled-kernel cache. -pub struct AsyncDeviceContext { - /// Device ordinal. - #[allow(dead_code)] - device_id: usize, - /// Primary CUDA context for this device. - context: Arc, - /// Dedicated stream for async memory deallocation (see [`DeviceBox`]). - /// - /// [`DeviceBox`]: crate::simt::device_box::DeviceBox - deallocator_stream: Arc, - /// Scheduling policy shared (via `Arc`) with all operations on this device. - policy: Arc, - /// Cache of loaded kernel functions, keyed by [`FunctionKey`] hashes. - functions: DeviceFunctions, -} - -/// Container for all per-device contexts on a single thread. -/// -/// Uses interior mutability ([`Cell`]) because it lives inside a -/// `thread_local!` and must be borrowed without `&mut`. -pub struct AsyncDeviceContexts { - /// Currently selected default device ordinal. - default_device: Cell, - /// Lazily initialized map of device ordinal to context. - devices: Cell>>, -} - -// Thread-local storage for per-device CUDA state. -// Lazily initialized on first access. Each thread gets its own independent -// set of device contexts, streams, and kernel caches. -thread_local!(static DEVICE_CONTEXTS: AsyncDeviceContexts = const { - AsyncDeviceContexts { - default_device: Cell::new(DEFAULT_DEVICE_ID), - devices: Cell::new(None), - } -}); - -/// Returns the default device ordinal for the current thread. -pub fn get_default_device() -> usize { - DEVICE_CONTEXTS.with(|ctx| ctx.default_device.get()) -} - -/// Initializes the thread-local device context map with capacity for -/// `num_devices` devices and sets `default_device_id` as the default. -/// -/// Must be called at most once per thread. Returns an error if already -/// initialized. -pub fn init_device_contexts( - default_device_id: usize, - num_devices: usize, -) -> Result<(), DeviceError> { - DEVICE_CONTEXTS.with(|ctx| { - let devices = ctx.devices.take(); - let is_uninitialized = devices.is_none(); - ctx.devices.set(devices); - device_assert( - default_device_id, - is_uninitialized, - "Context already initialized.", - ) - })?; - let devices = FxHashMap::with_capacity_and_hasher(num_devices, Default::default()); - DEVICE_CONTEXTS.with(|ctx| { - ctx.default_device.set(default_device_id); - ctx.devices.set(Some(devices)); - }); - Ok(()) -} - -/// Initializes the thread-local device contexts with default parameters. -fn init_device_contexts_default() -> Result<(), DeviceError> { - let default_device = get_default_device(); - init_device_contexts(default_device, DEFAULT_NUM_DEVICES) -} - -/// Creates a new [`AsyncDeviceContext`] for `device_id` with the given policy. -/// -/// Initializes the CUDA context, the scheduling policy's stream pool, and a -/// dedicated deallocator stream. -pub fn new_device_context( - device_id: usize, - mut policy: GlobalSchedulingPolicy, -) -> Result { - let context = CudaContext::new(device_id)?; - policy.init(&context)?; - let deallocator_stream = context.new_stream()?; - Ok(AsyncDeviceContext { - device_id, - context, - deallocator_stream, - policy: Arc::new(policy), - functions: FxHashMap::default(), - }) -} - -/// Inserts a new device context into `hashmap`. Errors if the device is -/// already present. -fn init_device( - hashmap: &mut FxHashMap, - device_id: usize, - policy: GlobalSchedulingPolicy, -) -> Result<(), DeviceError> { - let device_context = new_device_context(device_id, policy)?; - let pred = hashmap.insert(device_id, device_context).is_none(); - device_assert(device_id, pred, "Device is already initialized.") -} - -/// Initializes a device with the default round-robin policy. -fn init_with_default_policy( - hashmap: &mut FxHashMap, - device_id: usize, -) -> Result<(), DeviceError> { - let policy = - unsafe { StreamPoolRoundRobin::new(device_id, DEFAULT_ROUND_ROBIN_STREAM_POOL_SIZE) }; - init_device( - hashmap, - device_id, - GlobalSchedulingPolicy::RoundRobin(policy), - ) -} - -/// Borrows the thread-local [`AsyncDeviceContext`] for `device_id` immutably. -/// -/// Lazily initializes the context map and the specific device if needed. -fn with_global_device_context(device_id: usize, f: F) -> Result -where - F: FnOnce(&AsyncDeviceContext) -> R, -{ - DEVICE_CONTEXTS.with(|ctx| { - let mut hashmap = match ctx.devices.take() { - Some(hashmap) => hashmap, - None => { - init_device_contexts_default()?; - ctx.devices - .take() - .ok_or_else(|| device_error(device_id, "Failed to initialize context"))? - } - }; - if !hashmap.contains_key(&device_id) { - init_with_default_policy(&mut hashmap, device_id)?; - } - let device_context = hashmap - .get(&device_id) - .ok_or_else(|| device_error(device_id, "Failed to get context"))?; - let r = f(device_context); - ctx.devices.replace(Some(hashmap)); - Ok(r) - }) -} - -/// Borrows the thread-local [`AsyncDeviceContext`] for `device_id` mutably. -/// -/// Lazily initializes the context map and the specific device if needed. -fn with_global_device_context_mut(device_id: usize, f: F) -> Result -where - F: FnOnce(&mut AsyncDeviceContext) -> R, -{ - DEVICE_CONTEXTS.with(|ctx| { - let mut hashmap = match ctx.devices.take() { - Some(hashmap) => hashmap, - None => { - init_device_contexts_default()?; - ctx.devices - .take() - .ok_or_else(|| device_error(device_id, "Failed to initialize context"))? - } - }; - if !hashmap.contains_key(&device_id) { - init_with_default_policy(&mut hashmap, device_id)?; - } - let device_context = hashmap - .get_mut(&device_id) - .ok_or_else(|| device_error(device_id, "Failed to get context"))?; - let r = f(device_context); - ctx.devices.replace(Some(hashmap)); - Ok(r) - }) -} - -/// Provides the default device's scheduling policy to `f`. -/// -/// This is the primary entry point used by [`IntoFuture`] impls on -/// `DeviceOperation` combinators. -/// -/// [`IntoFuture`]: std::future::IntoFuture -pub fn with_default_device_policy(f: F) -> Result -where - F: FnOnce(&Arc) -> R, -{ - let default_device = get_default_device(); - with_global_device_context(default_device, |dc| f(&dc.policy)) -} - -/// Provides the deallocator stream for `device_id` to `f`. -/// -/// # Safety -/// -/// The stream must only be used for `cuMemFreeAsync` calls. Enqueuing other -/// work on this stream may interfere with async deallocation ordering. -pub unsafe fn with_deallocator_stream(device_id: usize, f: F) -> Result -where - F: FnOnce(&Arc) -> R, -{ - with_global_device_context(device_id, |dc| f(&dc.deallocator_stream)) -} - -/// Provides the CUDA context for `device_id` to `f`. -pub fn with_cuda_context(device_id: usize, f: F) -> Result -where - F: FnOnce(&Arc) -> R, -{ - with_global_device_context(device_id, |dc| f(&dc.context)) -} - -/// Sets the default device ordinal for the current thread. -pub fn set_default_device(default_device_id: usize) { - DEVICE_CONTEXTS.with(|ctx| { - ctx.default_device.set(default_device_id); - }) -} - -/// Loads a CUDA module from a file (`.cubin`, `.fatbin`, or `.ptx`) on -/// `device_id`. -pub fn load_module_from_file( - filename: &str, - device_id: usize, -) -> Result, DeviceError> { - with_cuda_context(device_id, |cuda_ctx| { - let module = cuda_ctx.load_module_from_file(filename)?; - Ok(module) - })? -} - -/// Loads a CUDA module from PTX source text on `device_id`. -pub fn load_module_from_ptx( - ptx_src: &str, - device_id: usize, -) -> Result, DeviceError> { - with_cuda_context(device_id, |cuda_ctx| { - let module = cuda_ctx.load_module_from_ptx_src(ptx_src)?; - Ok(module) - })? -} - -/// Inserts a compiled function into the per-device kernel cache. -/// -/// The function is keyed by the hash of `func_key`. Returns an error if a -/// different function was already registered under the same key (hash -/// collision). -pub fn insert_cuda_function( - device_id: usize, - func_key: &impl FunctionKey, - value: (Arc, Arc), -) -> Result<(), DeviceError> { - with_global_device_context_mut(device_id, |dc| { - let key = func_key.get_hash_string(); - let res = dc.functions.insert(key, value); - device_assert(device_id, res.is_none(), "Unexpected cache key collision.") - })? -} - -/// Retrieves a compiled [`CudaFunction`] from the per-device kernel cache. -/// -/// Returns an error if no function has been registered under `func_key`. -pub fn get_cuda_function( - device_id: usize, - func_key: &impl FunctionKey, -) -> Result, DeviceError> { - with_global_device_context(device_id, |dc| { - let key = func_key.get_hash_string(); - let (_module, function) = dc - .functions - .get(&key) - .ok_or_else(|| device_error(device_id, "Failed to get cuda function."))?; - Ok(Arc::clone(function)) - })? -} - -#[cfg(test)] -mod tests { - use super::*; - - fn assert_duplicate_init_error(result: Result<(), DeviceError>) { - assert!(matches!( - result, - Err(DeviceError::Context { - device_id: 0, - message, - }) if message == "Context already initialized." - )); - } - - #[test] - fn duplicate_init_preserves_existing_device_contexts() { - std::thread::spawn(|| { - init_device_contexts(0, 1).expect("initial context initialization should succeed"); - - assert_duplicate_init_error(init_device_contexts(0, 1)); - assert_duplicate_init_error(init_device_contexts(0, 1)); - }) - .join() - .expect("test thread should not panic"); - } -} diff --git a/cuda-async/src/simt/device_future.rs b/cuda-async/src/simt/device_future.rs deleted file mode 100644 index 82cfdcb79b..0000000000 --- a/cuda-async/src/simt/device_future.rs +++ /dev/null @@ -1,508 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Future type that bridges CUDA stream callbacks with Rust's async executor. -//! -//! [`DeviceFuture`] wraps a [`DeviceOperation`] and drives it through a -//! three-state machine: -//! -//! ```text -//! Idle ──poll()──> Executing ──callback fires──> Complete -//! │ │ -//! (enqueue work (return result) -//! + host callback) -//! ``` -//! -//! On the first poll the operation is executed on its assigned stream and a -//! host callback is registered via `cuLaunchHostFunc`. When the GPU reaches -//! the callback, it wakes the future through an [`AtomicWaker`], avoiding -//! busy-waits. -//! -//! [`DeviceOperation`]: crate::simt::device_operation::DeviceOperation -//! [`AtomicWaker`]: futures::task::AtomicWaker - -use crate::simt::device_operation::{DeviceOperation, ExecutionContext}; -use crate::simt::error::DeviceError; -use crate::simt::reclaim; -use futures::task::AtomicWaker; -use std::future::Future; -use std::mem; -use std::pin::Pin; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; -use std::task::{Context, Poll}; - -/// Lifecycle state of a [`DeviceFuture`]. -#[derive(Debug, Default, Eq, PartialEq, Copy, Clone)] -pub enum DeviceFutureState { - /// The future was constructed in a failed state (e.g. scheduling error). - Failed, - /// Initial state: the operation has not yet been submitted to the GPU. - #[default] - Idle, - /// The operation has been submitted; waiting for the stream callback. - Executing, - /// The result has been produced. Polling again will panic. - Complete, -} - -/// Shared state between a [`DeviceFuture`] and its `cuLaunchHostFunc` callback. -/// -/// The callback sets `complete` and wakes the stored waker, allowing the -/// executor to re-poll the future without busy-waiting. -#[derive(Debug, Default)] -pub struct StreamCallbackState { - /// Waker registered by the executor during [`Future::poll`]. - pub(crate) waker: AtomicWaker, - /// Set to `true` by the host callback when the GPU reaches the callback - /// point in the stream. - pub(crate) complete: AtomicBool, -} - -impl StreamCallbackState { - /// Creates a new callback state with no registered waker and - /// `complete = false`. - pub fn new() -> Self { - Self::default() - } - - /// Marks the stream callback as complete and wakes the associated future. - /// - /// Called from the `cuLaunchHostFunc` host-side callback. - pub fn signal(&self) { - self.complete.store(true, Ordering::Relaxed); - self.waker.wake(); - } -} - -/// A [`Future`] that executes a [`DeviceOperation`] on a CUDA stream and -/// resolves when the GPU signals completion via a host callback. -/// -/// Constructed by [`SchedulingPolicy::schedule`] or by the [`IntoFuture`] impl -/// on any `DeviceOperation`. -/// -/// # Cancellation -/// -/// Dropping the future never cancels submitted GPU work; kernels always run -/// to completion. Dropping an in-flight future records a CUDA event on the -/// assigned stream and parks the stored result in the [`reclaim`] limbo, -/// deferring host-side reclamation until the device timeline passes that -/// event. The drop itself never blocks on GPU progress. -/// -/// [`reclaim`]: crate::simt::reclaim -/// -/// [`SchedulingPolicy::schedule`]: crate::simt::scheduling_policies::SchedulingPolicy::schedule -/// [`IntoFuture`]: std::future::IntoFuture -#[derive(Debug)] -pub struct DeviceFuture> { - /// The operation to execute. Consumed on first poll. - pub(crate) device_operation: Option, - /// Stream and context for execution. Set by the scheduling policy. - pub(crate) execution_context: Option, - /// Holds the result between execution and final poll resolution. - pub(crate) result: Option, - /// Holds an error when the future is in the `Failed` state. - pub(crate) error: Option, - /// Current lifecycle state. - pub(crate) state: DeviceFutureState, - /// Shared state with the `cuLaunchHostFunc` callback. - pub(crate) callback_state: Option>, -} - -impl> DeviceFuture { - /// Creates an idle future with no operation or context attached. - pub fn new() -> Self { - Self::default() - } - - /// Creates a future that will immediately yield `Err(error)` on first poll. - pub fn failed(error: DeviceError) -> Self { - Self { - execution_context: None, - device_operation: None, - state: DeviceFutureState::Failed, - callback_state: None, - result: None, - error: Some(error), - } - } - - /// Registers a `cuLaunchHostFunc` callback that will signal `waker_state` - /// when the GPU reaches this point in the stream. - /// - /// # Safety - /// - /// The execution context must hold a valid, non-destroyed CUDA stream. - unsafe fn register_callback( - &self, - waker_state: Arc, - ) -> Result<(), DeviceError> { - let ctx = self.execution_context.as_ref().ok_or_else(|| { - DeviceError::Internal("Cannot execute future without an execution context.".to_string()) - })?; - ctx.get_cuda_stream().launch_host_function(move || { - waker_state.signal(); - })?; - Ok(()) - } - - /// Takes the stored operation, executes it on the bound stream, and stashes - /// the result. Called exactly once during the `Idle -> Executing` transition. - fn execute(&mut self) -> Result<(), DeviceError> { - let ctx = self.execution_context.as_ref().ok_or_else(|| { - DeviceError::Internal("Cannot execute future without an execution context.".to_string()) - })?; - let operation = self - .device_operation - .take() - .ok_or_else(|| DeviceError::Internal("No operation has been set.".to_string()))?; - let out = unsafe { operation.execute(ctx) }?; - self.result = Some(out); - Ok(()) - } - - /// Returns `true` when GPU work was submitted but the stored result has - /// not been handed to the caller. - /// - /// `result` is only populated by a successful [`execute`], so a stored - /// result implies submitted GPU work. The state check excludes futures - /// that never submitted anything (`Idle`, `Failed`); a `Complete` future - /// can still hold a result when callback registration failed after a - /// successful launch. - /// - /// [`execute`]: Self::execute - fn has_undelivered_submission(&self) -> bool { - matches!( - self.state, - DeviceFutureState::Executing | DeviceFutureState::Complete - ) && self.result.is_some() - } - - /// Hands a submitted-but-undelivered result to deferred reclamation. - /// - /// Records a completion event on the assigned stream and parks the - /// result in the [`reclaim`] limbo; the result is dropped by a later - /// sweep once the device timeline passes the event. Never blocks in - /// this path. Only when the event cannot be recorded does it fall back - /// to [`cleanup_executing_result_with`], which synchronizes the stream - /// and, if even that fails, leaks the result loudly. - /// - /// [`cleanup_executing_result_with`]: Self::cleanup_executing_result_with - fn reclaim_in_flight_result(&mut self) { - if !self.has_undelivered_submission() { - return; - } - let stream = self - .execution_context - .as_ref() - .map(|ctx| Arc::clone(ctx.get_cuda_stream())); - if let Some(stream) = &stream { - if let Ok(event) = stream.record_event(None) { - let result = self.result.take().expect("Checked above."); - reclaim::park(event, Box::new(result)); - return; - } - } - // No execution context, or the driver rejected the event record. - // Fall back to the blocking path before resorting to a loud leak. - self.cleanup_executing_result_with(move || { - let stream = stream.ok_or_else(|| { - DeviceError::Internal( - "Cannot clean up an in-flight future without an execution context.".to_string(), - ) - })?; - stream.synchronize().map_err(DeviceError::Driver) - }); - } - - /// Blocking cleanup fallback: runs `synchronize` and drops the stored - /// result on success; on failure the result is leaked loudly, because - /// dropping resources the device may still use is worse than leaking. - /// - /// Deferred (non-blocking) reclamation in - /// [`reclaim_in_flight_result`](Self::reclaim_in_flight_result) is - /// always preferred; this path only runs when no completion event could - /// be recorded. - fn cleanup_executing_result_with(&mut self, synchronize: F) - where - F: FnOnce() -> Result<(), DeviceError>, - { - if !self.has_undelivered_submission() { - return; - } - - let Some(result) = self.result.take() else { - return; - }; - - if let Err(error) = synchronize() { - crate::leak::report_leak(format_args!( - "cuda-async: leaking in-flight future result after cleanup failure: {error}" - )); - // If cleanup cannot prove the stream is idle, leaking the owned - // result is safer than dropping buffers that device work may still - // be using. - mem::forget(result); - return; - } - - drop(result); - } -} - -impl> Default for DeviceFuture { - fn default() -> Self { - Self { - device_operation: Default::default(), - execution_context: Default::default(), - result: Default::default(), - error: Default::default(), - state: Default::default(), - callback_state: Default::default(), - } - } -} - -/// `DeviceFuture` does not contain self-referential pointers, so it is safe -/// to move. -impl> Unpin for DeviceFuture {} - -impl> Drop for DeviceFuture { - fn drop(&mut self) { - // Reclaim earlier cancellations whose GPU work has since finished, - // then defer this future's own in-flight result (if any). - reclaim::sweep(); - self.reclaim_in_flight_result(); - } -} - -/// State-machine implementation of [`Future`] for CUDA device work. -/// -/// | State | Action on poll | -/// |-------------|---------------------------------------------------------| -/// | `Failed` | Immediately returns `Err`. | -/// | `Idle` | Executes the operation, registers callback, -> Pending. | -/// | `Executing` | Checks callback flag; returns result if done. | -/// | `Complete` | Panics -- must not poll after completion. | -impl> Future for DeviceFuture { - type Output = Result; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // Opportunistically reclaim cancelled results whose GPU work has - // completed. Cheap when nothing is parked (one atomic load). - reclaim::sweep(); - - if self.state == DeviceFutureState::Failed { - self.state = DeviceFutureState::Complete; - let error = self - .error - .take() - .expect("Failed state must carry an error."); - return Poll::Ready(Err(error)); - } - - if self.callback_state.is_none() { - self.callback_state = Some(Arc::new(StreamCallbackState::new())); - } - let waker_state = self - .callback_state - .as_ref() - .map(Arc::clone) - .expect("Impossible."); - - match self.state { - DeviceFutureState::Idle => { - waker_state.waker.register(cx.waker()); - if let Err(e) = self.execute() { - self.state = DeviceFutureState::Complete; - return Poll::Ready(Err(e)); - } - if let Err(e) = unsafe { self.register_callback(Arc::clone(&waker_state)) } { - self.state = DeviceFutureState::Complete; - // The launch already succeeded, so GPU work is in flight - // and `result` holds the owned resources. Route them - // through deferred reclamation instead of leaving them - // for an immediate drop while the kernel may still run - // (issue #99, callback-registration-failure path). - self.reclaim_in_flight_result(); - return Poll::Ready(Err(e)); - } - self.state = DeviceFutureState::Executing; - Poll::Pending - } - DeviceFutureState::Executing => { - if waker_state.complete.load(Ordering::Relaxed) { - self.state = DeviceFutureState::Complete; - return Poll::Ready(Ok(self.result.take().expect("Expected result."))); - } - waker_state.waker.register(cx.waker()); - if waker_state.complete.load(Ordering::Relaxed) { - self.state = DeviceFutureState::Complete; - Poll::Ready(Ok(self.result.take().expect("Expected result."))) - } else { - Poll::Pending - } - } - DeviceFutureState::Complete => panic!("Poll called after completion."), - DeviceFutureState::Failed => unreachable!(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::simt::device_operation::Value; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - use std::sync::Mutex; - - #[derive(Clone)] - struct DropTracker { - events: Arc>>, - } - - impl Drop for DropTracker { - fn drop(&mut self) { - self.events.lock().unwrap().push("drop"); - } - } - - #[test] - fn cleanup_executing_result_synchronizes_before_drop() { - let events = Arc::new(Mutex::new(Vec::new())); - let tracker = DropTracker { - events: Arc::clone(&events), - }; - let mut future: DeviceFuture> = DeviceFuture { - device_operation: None, - execution_context: None, - result: Some(tracker), - error: None, - state: DeviceFutureState::Executing, - callback_state: None, - }; - - future.cleanup_executing_result_with(|| { - events.lock().unwrap().push("sync"); - Ok(()) - }); - - assert_eq!(events.lock().unwrap().as_slice(), ["sync", "drop"]); - assert!(future.result.is_none()); - } - - #[test] - fn cleanup_executing_result_leaks_when_synchronize_fails() { - let drops = Arc::new(AtomicUsize::new(0)); - - struct CountDrop(Arc); - - impl Drop for CountDrop { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - let mut future: DeviceFuture> = DeviceFuture { - device_operation: None, - execution_context: None, - result: Some(CountDrop(Arc::clone(&drops))), - error: None, - state: DeviceFutureState::Executing, - callback_state: None, - }; - - let mut capture = crate::leak::capture::start(); - future.cleanup_executing_result_with(|| Err(DeviceError::Internal("boom".to_string()))); - let reports = capture.take(); - - assert_eq!(drops.load(Ordering::Relaxed), 0); - assert!(future.result.is_none()); - assert_eq!(reports.len(), 1, "the leak must be reported: {reports:?}"); - } - - /// The shape left behind by issue #99's second path: `execute()` - /// succeeded (GPU work submitted, `result` populated) but - /// `register_callback()` failed, so poll flipped the state to - /// `Complete` with the result still stored. Cleanup must treat that - /// exactly like a cancelled `Executing` future. - #[test] - fn cleanup_covers_complete_future_left_by_callback_registration_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let tracker = DropTracker { - events: Arc::clone(&events), - }; - let mut future: DeviceFuture> = DeviceFuture { - device_operation: None, - execution_context: None, - result: Some(tracker), - error: None, - state: DeviceFutureState::Complete, - callback_state: None, - }; - - assert!(future.has_undelivered_submission()); - future.cleanup_executing_result_with(|| { - events.lock().unwrap().push("sync"); - Ok(()) - }); - - assert_eq!(events.lock().unwrap().as_slice(), ["sync", "drop"]); - assert!(future.result.is_none()); - assert!(!future.has_undelivered_submission()); - } - - /// A `Complete` future whose result was already delivered to the caller - /// has nothing to reclaim. - #[test] - fn cleanup_is_noop_after_result_delivery() { - let mut future: DeviceFuture> = DeviceFuture { - device_operation: None, - execution_context: None, - result: None, - error: None, - state: DeviceFutureState::Complete, - callback_state: None, - }; - - assert!(!future.has_undelivered_submission()); - future.cleanup_executing_result_with(|| { - panic!("delivered futures must not synchronize during cleanup") - }); - future.reclaim_in_flight_result(); - } - - #[test] - fn cleanup_is_noop_for_idle_future() { - let drops = Arc::new(AtomicUsize::new(0)); - - struct CountDrop(Arc); - - impl Drop for CountDrop { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - let mut future: DeviceFuture> = DeviceFuture { - device_operation: None, - execution_context: None, - result: Some(CountDrop(Arc::clone(&drops))), - error: None, - state: DeviceFutureState::Idle, - callback_state: None, - }; - - future.cleanup_executing_result_with(|| { - panic!("idle futures should not synchronize during cleanup") - }); - - assert_eq!(drops.load(Ordering::Relaxed), 0); - assert!(future.result.is_some()); - - drop(future); - assert_eq!(drops.load(Ordering::Relaxed), 1); - } -} diff --git a/cuda-async/src/simt/device_operation.rs b/cuda-async/src/simt/device_operation.rs deleted file mode 100644 index 21d336c6de..0000000000 --- a/cuda-async/src/simt/device_operation.rs +++ /dev/null @@ -1,981 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Lazy, composable GPU operations and combinator types. -//! -//! The [`DeviceOperation`] trait is the core abstraction. Each operation -//! describes GPU work without binding to a stream. Combinators (`and_then`, -//! `zip`, `apply`, `with_context`) compose operations into dataflow graphs -//! that remain stream-agnostic until scheduling time. -//! -//! # Scheduling model -//! -//! | Method | What it does | -//! |--------------|-------------------------------------------------------------------| -//! | [`schedule`] | Pairs the operation with a stream, returns a [`DeviceFuture`]. | -//! | [`sync`] | Shorthand: schedule + execute + synchronize on the default device.| -//! | [`sync_on`] | Execute and synchronize on a specific stream. | -//! | [`async_on`] | `unsafe`: execute on a specific stream **without** synchronizing. | -//! -//! [`schedule`]: DeviceOperation::schedule -//! [`sync`]: DeviceOperation::sync -//! [`sync_on`]: DeviceOperation::sync_on -//! [`async_on`]: DeviceOperation::async_on -//! [`DeviceFuture`]: crate::simt::device_future::DeviceFuture - -use crate::simt::device_context::with_default_device_policy; -use crate::simt::device_future::DeviceFuture; -use crate::simt::error::DeviceError; -use crate::simt::scheduling_policies::SchedulingPolicy; -use cuda_core::{CudaContext, CudaStream}; -use std::cell::UnsafeCell; -use std::future::IntoFuture; -use std::marker::PhantomData; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; - -/// CUDA device ordinal. Type alias for readability. -pub type Device = usize; - -/// Binds a [`DeviceOperation`] to a concrete CUDA stream and context for -/// execution. -/// -/// Created by the scheduling policy when an operation is scheduled. Passed to -/// [`DeviceOperation::execute`] to provide the stream and context. -#[derive(Debug, Clone)] -pub struct ExecutionContext { - /// Device ordinal derived from the CUDA context. - device: Device, - /// Stream on which GPU work will be enqueued. - cuda_stream: Arc, - /// CUDA context that owns the stream. - cuda_context: Arc, -} - -impl ExecutionContext { - /// Constructs a context from a stream, deriving the device and CUDA context - /// from the stream's owning context. - pub fn new(cuda_stream: Arc) -> Self { - let cuda_context = Arc::clone(cuda_stream.context()); - let device = cuda_context.ordinal(); - Self { - cuda_stream, - cuda_context, - device, - } - } - - /// Returns the CUDA stream. - pub fn get_cuda_stream(&self) -> &Arc { - &self.cuda_stream - } - - /// Returns the CUDA context. - pub fn get_cuda_context(&self) -> &Arc { - &self.cuda_context - } - - /// Returns the device ordinal. - pub fn get_device_id(&self) -> Device { - self.device - } -} - -/// A lazy, composable GPU operation that may be executed synchronously or -/// asynchronously. -/// -/// `DeviceOperation` is the core trait of the `cuda-async` crate. It -/// represents a unit of GPU work that is **stream-agnostic**: the concrete -/// CUDA stream is chosen only at scheduling time, not at construction time. -/// -/// # Composing operations -/// -/// Combinators build complex dataflow graphs without touching streams: -/// -/// | Combinator | Effect | -/// |---------------------------|-------------------------------------------------| -/// | [`and_then`] | Sequence: `A` then `f(result_a)`. | -/// | [`and_then_with_context`] | Like `and_then` but the closure sees the stream.| -/// | [`apply`] | Alias for `and_then`. | -/// | [`arc`] | Wraps the output in `Arc`. | -/// | `zip!` | Runs two or three operations, returns a tuple. | -/// | `unzip!` | Splits a tuple-producing operation. | -/// -/// # Executing operations -/// -/// | Method | Picks stream via | Blocks? | Async? | `unsafe`? | -/// |--------------|--------------------------|---------|--------|-----------| -/// | [`schedule`] | `SchedulingPolicy` | No | Yes | No | -/// | `.await` | Default policy | No | Yes | No | -/// | [`sync`] | Default policy | Yes | No | No | -/// | [`sync_on`] | Caller-provided stream | Yes | No | No | -/// | [`async_on`] | Caller-provided stream | No | No | **`unsafe`** | -/// -/// [`async_on`] is the only one that is `unsafe`, and it is the only one that -/// returns while GPU work may still be in flight: the caller must synchronize -/// the stream before consuming device-side outputs. The other four either block -/// or hand back a future that does the waiting. -/// -/// # Implementors -/// -/// Implement [`execute`] to describe the GPU work. The blanket [`IntoFuture`] -/// impl must also be provided (typically via the same boilerplate that -/// delegates to `with_default_device_policy`). -/// -/// [`and_then`]: DeviceOperation::and_then -/// [`and_then_with_context`]: DeviceOperation::and_then_with_context -/// [`apply`]: DeviceOperation::apply -/// [`arc`]: DeviceOperation::arc -/// [`schedule`]: DeviceOperation::schedule -/// [`sync`]: DeviceOperation::sync -/// [`sync_on`]: DeviceOperation::sync_on -/// [`async_on`]: DeviceOperation::async_on -/// [`execute`]: DeviceOperation::execute -pub trait DeviceOperation: - Send + Sized + IntoFuture::Output, DeviceError>> -{ - /// The value produced when the operation completes successfully. - /// - /// The `'static` bound exists because a cancelled in-flight result is - /// parked in the [`reclaim`](crate::simt::reclaim) limbo and dropped only - /// after the GPU work completes, which can be after any non-`'static` - /// borrow it carries would have expired. - type Output: Send + 'static; - - /// Submits GPU work to the stream in `context` and returns the result. - /// - /// # Safety - /// - /// GPU work may still be in flight when this returns. The caller must - /// synchronize the stream before reading device-side outputs. - unsafe fn execute( - self, - context: &ExecutionContext, - ) -> Result<::Output, DeviceError>; - - /// Pairs this operation with a stream chosen by `policy` and returns a - /// [`DeviceFuture`] that can be `.await`-ed. - fn schedule( - self, - policy: &P, - ) -> Result::Output, Self>, DeviceError> { - policy.schedule(self) - } - - /// Chains a dependent operation: executes `self`, then passes its output - /// to `f` to produce the next operation. - fn and_then( - self, - f: F, - ) -> AndThen<::Output, Self, O, DO, F> - where - DO: DeviceOperation, - F: FnOnce(::Output) -> DO, - { - AndThen { - op: self, - closure: f, - } - } - - /// Like [`and_then`](Self::and_then), but the closure also receives the - /// [`ExecutionContext`] so it can inspect the stream or device. - fn and_then_with_context( - self, - f: F, - ) -> AndThenWithContext<::Output, Self, O, DO, F> - where - DO: DeviceOperation, - F: FnOnce(&ExecutionContext, ::Output) -> DO, - { - AndThenWithContext { - op: self, - closure: f, - } - } - - /// Wraps the output in an [`Arc`], useful when the result must be shared - /// across multiple consumers. - fn arc(self) -> DeviceOperationArc<::Output, Self> - where - ::Output: Sync, - { - DeviceOperationArc { op: self } - } - - /// Alias for [`and_then`](Self::and_then). - fn apply( - self, - f: F, - ) -> AndThen<::Output, Self, O, DO, F> - where - DO: DeviceOperation, - F: FnOnce(::Output) -> DO, - { - self.and_then(f) - } - - /// Executes the operation synchronously on the default device using the - /// thread-local scheduling policy. Blocks until the stream is idle. - fn sync(self) -> Result<::Output, DeviceError> { - with_default_device_policy(|policy| policy.sync(self))? - } - - /// Executes the operation on `stream` **without** synchronizing. - /// - /// # Safety - /// - /// GPU work may still be in flight when this returns. The caller must - /// synchronize `stream` before consuming device-side outputs. - unsafe fn async_on( - self, - stream: &Arc, - ) -> Result<::Output, DeviceError> { - let ctx = ExecutionContext::new(Arc::clone(stream)); - unsafe { self.execute(&ctx) } - } - - /// Executes the operation on `stream` and synchronizes before returning. - fn sync_on( - self, - stream: &Arc, - ) -> Result<::Output, DeviceError> { - let ctx = ExecutionContext::new(Arc::clone(stream)); - let res = unsafe { self.execute(&ctx) }; - finish_sync(res, stream.synchronize()) - } -} - -fn finish_sync( - operation_result: Result, - synchronize_result: Result<(), cuda_core::DriverError>, -) -> Result { - let output = operation_result?; - synchronize_result.map_err(DeviceError::Driver)?; - Ok(output) -} - -// --- Combinators --- - -/// Wraps a [`DeviceOperation`] whose output is `Arc`-wrapped. -/// -/// Produced by [`DeviceOperation::arc`]. -pub struct DeviceOperationArc> { - /// The inner operation whose result will be wrapped in [`Arc`]. - op: DI, -} - -/// # Safety -/// -/// `DI` is `Send` (required by `DeviceOperation`), and `I: Send + Sync` -/// ensures the `Arc` output is safe to transfer. -unsafe impl> Send for DeviceOperationArc {} - -/// Executes the inner operation and wraps the result in [`Arc`]. -impl> DeviceOperation - for DeviceOperationArc -{ - type Output = Arc; - - unsafe fn execute(self, context: &ExecutionContext) -> Result, DeviceError> { - unsafe { - let val = self.op.execute(context)?; - Ok(Arc::new(val)) - } - } -} - -/// Schedules via the thread-local default policy. -impl> IntoFuture - for DeviceOperationArc -{ - type Output = Result, DeviceError>; - type IntoFuture = DeviceFuture, DeviceOperationArc>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Sequential composition: execute `DI`, then pass its output through `F` to -/// produce a second operation `DO`, and execute that. -/// -/// Produced by [`DeviceOperation::and_then`]. -pub struct AndThen -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(I) -> DO, -{ - /// First operation. - op: DI, - /// Closure mapping the first operation's output to the second operation. - closure: F, -} - -/// # Safety -/// -/// Both `DI` and `F` are `Send`. The struct owns them exclusively, so -/// transferring across threads is safe. -unsafe impl Send for AndThen -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(I) -> DO + Send, -{ -} - -/// Executes the first operation, feeds its result to the closure, then -/// executes the resulting second operation on the same stream. -impl DeviceOperation for AndThen -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(I) -> DO + Send, -{ - type Output = O; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - let input = self.op.execute(context)?; - let output_op = (self.closure)(input); - output_op.execute(context) - } - } -} - -/// Schedules via the thread-local default policy. -impl IntoFuture for AndThen -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(I) -> DO + Send, -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Like [`AndThen`] but the closure additionally receives the -/// [`ExecutionContext`], giving access to the stream and device. -/// -/// Produced by [`DeviceOperation::and_then_with_context`]. -pub struct AndThenWithContext -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext, I) -> DO, -{ - /// First operation. - op: DI, - /// Closure mapping `(context, result)` to the second operation. - closure: F, -} - -/// # Safety -/// -/// Both `DI` and `F` are `Send`. The struct owns them exclusively. -unsafe impl Send for AndThenWithContext -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext, I) -> DO + Send, -{ -} - -/// Executes the first operation, then passes `(context, result)` to the -/// closure and executes the resulting second operation. -impl DeviceOperation for AndThenWithContext -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext, I) -> DO + Send, -{ - type Output = O; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - let input = self.op.execute(context)?; - let output_op = (self.closure)(context, input); - output_op.execute(context) - } - } -} - -/// Schedules via the thread-local default policy. -impl IntoFuture for AndThenWithContext -where - DI: DeviceOperation, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext, I) -> DO + Send, -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// A [`DeviceOperation`] that immediately returns a pre-computed value without -/// touching the GPU. -pub struct Value(T); - -/// # Safety -/// -/// `Value` holds only the inner `T`. If `T` is `Send` the wrapper is too; -/// the bound is enforced by the `DeviceOperation` impl. -unsafe impl Send for Value {} - -/// Returns the wrapped value directly -- no GPU work is performed. -impl DeviceOperation for Value { - type Output = T; - - unsafe fn execute(self, _context: &ExecutionContext) -> Result { - Ok(self.0) - } -} - -/// Schedules via the thread-local default policy. -impl IntoFuture for Value { - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Wraps `x` in a [`Value`] operation that returns it immediately. -pub fn value(x: T) -> Value { - Value(x) -} - -/// Converts any `Send` value into a no-op [`DeviceOperation`] via [`value`]. -pub trait IntoDeviceOperation { - /// Wraps `self` into a [`Value`] device operation. - fn device_operation(self) -> Value; -} - -impl IntoDeviceOperation for T { - fn device_operation(self) -> Value { - value(self) - } -} - -/// Deferred-closure operation: the closure produces the real operation at -/// execution time rather than at construction time. -/// -/// Useful when building the inner operation requires state only available -/// after scheduling (though it does not receive the [`ExecutionContext`] -- -/// see [`StreamOperation`] for that). -pub struct Empty, F: FnOnce() -> DO> { - /// Closure that produces the inner operation. - closure: F, -} - -/// Wraps a closure in an [`Empty`] deferred operation. -pub fn empty, F: FnOnce() -> DO>( - closure: F, -) -> Empty { - Empty { closure } -} - -/// # Safety -/// -/// The closure `F` and its produced operation `DO` are both `Send`. The -/// struct owns the closure exclusively. -unsafe impl, F: FnOnce() -> DO> Send for Empty {} - -/// Invokes the closure to produce the inner operation, then executes it. -impl, F: FnOnce() -> DO> DeviceOperation - for Empty -{ - type Output = O; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - let op = (self.closure)(); - op.execute(context) - } - } -} - -/// Schedules via the thread-local default policy. -impl, F: FnOnce() -> DO> IntoFuture - for Empty -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Pair combinator: executes two operations sequentially on the same stream -/// and returns both results as a tuple. -/// -/// Constructed via `_zip` or the `zip!` macro. -pub struct Zip, B: DeviceOperation> -{ - phantom: PhantomData<(T1, T2)>, - /// First operation. - a: A, - /// Second operation. - b: B, -} - -/// # Safety -/// -/// Both `A` and `B` are `Send` (required by `DeviceOperation`). -unsafe impl, B: DeviceOperation> - Send for Zip -{ -} - -/// Constructs a [`Zip`] from two operations. -fn _zip, B: DeviceOperation>( - a: A, - b: B, -) -> Zip { - Zip { - phantom: PhantomData, - a, - b, - } -} - -/// Executes `a` then `b` on the same stream, returning `(T1, T2)`. -impl< - T1: Send + 'static, - T2: Send + 'static, - A: DeviceOperation, - B: DeviceOperation, - > DeviceOperation for Zip -{ - type Output = (T1, T2); - - unsafe fn execute(self, context: &ExecutionContext) -> Result<(T1, T2), DeviceError> { - unsafe { - let a = self.a.execute(context)?; - let b = self.b.execute(context)?; - Ok((a, b)) - } - } -} - -/// Schedules via the thread-local default policy. -impl< - T1: Send + 'static, - T2: Send + 'static, - A: DeviceOperation, - B: DeviceOperation, - > IntoFuture for Zip -{ - type Output = Result<(T1, T2), DeviceError>; - type IntoFuture = DeviceFuture<(T1, T2), Zip>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Trait enabling `.zip()` on tuples of [`DeviceOperation`]s. -/// -/// Implemented for 2-tuples and 3-tuples. -pub trait Zippable { - /// Combines the operations into a single operation returning a tuple of - /// results. - fn zip(self) -> impl DeviceOperation; -} - -/// Zips two operations into a pair. -impl< - T0: Send + 'static, - T1: Send + 'static, - DI0: DeviceOperation, - DI1: DeviceOperation, - > Zippable<(DI0, DI1), (T0, T1)> for (DI0, DI1) -{ - fn zip(self) -> impl DeviceOperation { - _zip(self.0, self.1) - } -} - -/// Zips three operations into a triple by nesting two binary zips. -impl< - T0: Send + 'static, - T1: Send + 'static, - T2: Send + 'static, - DI0: DeviceOperation, - DI1: DeviceOperation, - DI2: DeviceOperation, - > Zippable<(DI0, DI1, DI2), (T0, T1, T2)> for (DI0, DI1, DI2) -{ - fn zip(self) -> impl DeviceOperation { - let cons = _zip(self.1, self.2); - let cons = _zip(self.0, cons); - cons.and_then(|(arg0, (arg1, arg2))| value((arg0, arg1, arg2))) - } -} - -/// Zips one, two, or three [`DeviceOperation`]s into a single operation -/// returning a tuple of results. -/// -/// ```ignore -/// let (a, b) = zip!(op_a, op_b).sync()?; -/// let (x, y, z) = zip!(op_x, op_y, op_z).sync()?; -/// ``` -#[allow(unused_macros)] // collides with this crate's exported zip!; reconcile at merge -macro_rules! zip { - ($arg0:expr) => { - $arg0 - }; - ($arg0:expr, $arg1:expr) => { - ($arg0, $arg1).zip() - }; - ($arg0:expr, $arg1:expr, $arg2:expr) => { - ($arg0, $arg1, $arg2).zip() - }; -} -#[allow(unused_imports)] // kept for parity with oxide's exported form -pub(crate) use zip; - -/// Deferred operation that receives the [`ExecutionContext`] before producing -/// the inner operation. -/// -/// Unlike [`Empty`], the closure has access to the stream and device, making -/// it possible to build context-dependent operations at execution time. -pub struct StreamOperation< - O: Send, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext) -> DO + Send, -> { - /// Closure that receives the execution context and produces the inner op. - f: F, -} - -/// Calls the closure with the context, then executes the resulting operation. -impl< - O: Send + 'static, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext) -> DO + Send, - > DeviceOperation for StreamOperation -{ - type Output = O; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - let op = (self.f)(context); - op.execute(context) - } - } -} - -/// Wraps a closure that needs the [`ExecutionContext`] into a -/// [`DeviceOperation`]. -/// -/// The closure is invoked at execution time with the stream and context, -/// and must return a `DeviceOperation` that will be immediately executed. -pub fn with_context< - O: Send + 'static, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext) -> DO + Send, ->( - f: F, -) -> impl DeviceOperation { - StreamOperation { f } -} - -/// Schedules via the thread-local default policy. -impl< - O: Send + 'static, - DO: DeviceOperation, - F: FnOnce(&ExecutionContext) -> DO + Send, - > IntoFuture for StreamOperation -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Shared state backing the [`SelectLeft`] / [`SelectRight`] split. -/// -/// Holds the source operation and its memoized results. The first selector -/// to execute triggers the source; subsequent selectors read the cached -/// values. Interior mutability is used via [`UnsafeCell`] because execution -/// is always sequential within a single stream. -pub struct Select> { - /// Guards one-shot execution of the source operation. - computed: AtomicBool, - /// Source operation. Consumed on the first call to [`execute`](Self::execute). - input: UnsafeCell>, - /// Cached left result. - left: UnsafeCell>, - /// Cached right result. - right: UnsafeCell>, -} - -impl> Select { - /// Executes the source operation if it has not been executed yet, caching - /// both halves of the tuple. - /// - /// # Safety - /// - /// Must only be called from within a single-stream execution context. - /// Concurrent calls from different threads would race on the `UnsafeCell`s. - unsafe fn execute(self: &Arc, context: &ExecutionContext) -> Result<(), DeviceError> { - unsafe { - if !self.computed.load(Ordering::Acquire) { - let input = self.input.get(); - let input = input.as_mut(); - let input = input.unwrap().take().ok_or_else(|| { - crate::simt::error::device_error( - context.get_device_id(), - "Select operation failed.", - ) - })?; - let (left, right) = input.execute(context)?; - *self.left.get() = Some(left); - *self.right.get() = Some(right); - self.computed.store(true, Ordering::Release); - } - Ok(()) - } - } - - /// Takes the cached left value. Must be called after [`execute`](Self::execute). - /// - /// # Safety - /// - /// Caller must ensure `execute` has completed and this is called at most once. - unsafe fn left(&self) -> T1 { - let cell = self.left.get(); - let cell = unsafe { cell.as_mut() }; - cell.unwrap().take().unwrap() - } - - /// Takes the cached right value. Must be called after [`execute`](Self::execute). - /// - /// # Safety - /// - /// Caller must ensure `execute` has completed and this is called at most once. - unsafe fn right(&self) -> T2 { - let cell = self.right.get(); - let cell = unsafe { cell.as_mut() }; - cell.unwrap().take().unwrap() - } -} - -/// Operation that extracts the **left** element of an unzipped pair. -/// -/// Shares a [`Select`] with its corresponding [`SelectRight`] so the source -/// operation is executed at most once. -pub struct SelectLeft> { - /// Shared memoization state. - select: Arc>, -} - -/// # Safety -/// -/// The `Arc>` is safe to send because `Select` is only mutated -/// through `UnsafeCell` during single-stream execution, never concurrently. -unsafe impl> Send - for SelectLeft -{ -} - -/// Triggers the shared source operation (if not yet done) and returns the -/// left element. -impl> DeviceOperation - for SelectLeft -{ - type Output = T1; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - self.select.execute(context)?; - Ok(self.select.left()) - } - } -} - -/// Schedules via the thread-local default policy. -impl> IntoFuture - for SelectLeft -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Operation that extracts the **right** element of an unzipped pair. -/// -/// Shares a [`Select`] with its corresponding [`SelectLeft`] so the source -/// operation is executed at most once. -pub struct SelectRight> { - /// Shared memoization state. - select: Arc>, -} - -/// # Safety -/// -/// See [`SelectLeft`]'s `Send` impl. -unsafe impl> Send - for SelectRight -{ -} - -/// Triggers the shared source operation (if not yet done) and returns the -/// right element. -impl> DeviceOperation - for SelectRight -{ - type Output = T2; - - unsafe fn execute(self, context: &ExecutionContext) -> Result { - unsafe { - self.select.execute(context)?; - Ok(self.select.right()) - } - } -} - -/// Schedules via the thread-local default policy. -impl> IntoFuture - for SelectRight -{ - type Output = Result; - type IntoFuture = DeviceFuture>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -/// Splits a tuple-producing operation into two independent operations that -/// share execution: the source runs at most once, and each selector extracts -/// one element. -fn _unzip>( - input: DI, -) -> (SelectLeft, SelectRight) { - let select = Select { - computed: AtomicBool::new(false), - input: UnsafeCell::new(Some(input)), - left: UnsafeCell::new(None), - right: UnsafeCell::new(None), - }; - let select = Arc::new(select); - let out1 = SelectLeft { - select: Arc::clone(&select), - }; - let out2 = SelectRight { select }; - (out1, out2) -} - -/// Trait enabling `.unzip()` on any [`DeviceOperation`] that produces a -/// 2-tuple. -pub trait Unzippable2 -where - Self: DeviceOperation, -{ - /// Splits this operation into two independent operations, one for each - /// tuple element. The source executes at most once. - fn unzip( - self, - ) -> ( - impl DeviceOperation, - impl DeviceOperation, - ) { - _unzip(self) - } -} - -/// Blanket impl: any operation producing `(T0, T1)` is unzippable. -impl> - Unzippable2 for DI -{ -} - -/// Splits a tuple-producing [`DeviceOperation`] into per-element operations. -/// -/// ```ignore -/// let (left, right) = unzip!(pair_op); -/// ``` -#[allow(unused_macros)] // collides with this crate's exported unzip!; reconcile at merge -macro_rules! unzip { - ($arg0:expr) => { - $arg0.unzip() - }; -} -#[allow(unused_imports)] // kept for parity with oxide's exported form -pub(crate) use unzip; - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn finish_sync_returns_operation_result_after_successful_synchronize() { - let result = finish_sync::(Ok(7), Ok(())); - - assert_eq!(result, Ok(7)); - } - - #[test] - fn finish_sync_preserves_operation_error_after_successful_synchronize() { - let operation_error = DeviceError::Launch("launch failed".to_string()); - let result = finish_sync::(Err(operation_error.clone()), Ok(())); - - assert_eq!(result, Err(operation_error)); - } - - #[test] - fn finish_sync_propagates_synchronize_error_instead_of_panicking() { - let driver_error = - cuda_core::DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); - let result = finish_sync::(Ok(7), Err(driver_error)); - - assert_eq!(result, Err(DeviceError::Driver(driver_error))); - } - - #[test] - fn finish_sync_preserves_operation_error_when_synchronize_also_fails() { - let operation_error = DeviceError::Launch("launch failed".to_string()); - let driver_error = - cuda_core::DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); - let result = finish_sync::(Err(operation_error.clone()), Err(driver_error)); - - assert_eq!(result, Err(operation_error)); - } -} diff --git a/cuda-async/src/simt/error.rs b/cuda-async/src/simt/error.rs deleted file mode 100644 index 60c9296d1d..0000000000 --- a/cuda-async/src/simt/error.rs +++ /dev/null @@ -1,94 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Error types and assertion helpers for CUDA device operations. -//! -//! [`DeviceError`] is the unified error type for the `cuda-async` crate. It -//! captures failures from the CUDA driver, device context management, kernel -//! cache lookups, scheduling, and kernel launches. Assertion helpers provide a -//! concise way to produce contextual errors in device-facing code. - -use cuda_core::DriverError; - -/// Unified error type for all `cuda-async` operations. -/// -/// Each variant identifies a distinct failure domain so callers can -/// pattern-match on the error source without parsing messages. -#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)] -pub enum DeviceError { - /// Propagated from the CUDA driver API (see [`DriverError`]). - #[error("CUDA driver error: {0}")] - Driver(#[from] DriverError), - - /// Device context initialization or access failure for a specific device. - #[error("device context error (device_id={device_id}): {message}")] - Context { - /// Ordinal of the device that reported the error. - device_id: usize, - /// Human-readable description of the failure. - message: String, - }, - - /// Failure during kernel module loading or function cache lookup. - #[error("kernel cache error: {0}")] - KernelCache(String), - - /// Stream scheduling or policy configuration error. - #[error("scheduling error: {0}")] - Scheduling(String), - - /// Kernel launch parameter validation or driver-level launch failure. - #[error("kernel launch error: {0}")] - Launch(String), - - /// Logic error internal to the async runtime. - #[error("internal error: {0}")] - Internal(String), - - /// Catch-all for errors converted from [`anyhow::Error`]. - #[error("{0}")] - Anyhow(String), -} - -/// Converts an [`anyhow::Error`] by capturing its debug representation. -impl From for DeviceError { - fn from(error: anyhow::Error) -> Self { - DeviceError::Anyhow(format!("{:?}", error)) - } -} - -/// Returns `Err(DeviceError::Launch)` when `pred` is `false`. -/// -/// Use for precondition checks on kernel launch parameters (grid/block dims, -/// shared memory size, argument counts, etc.). -pub fn kernel_launch_assert(pred: bool, message: &str) -> Result<(), DeviceError> { - if !pred { - Err(DeviceError::Launch(message.to_string())) - } else { - Ok(()) - } -} - -/// Returns `Err(DeviceError::Context)` when `pred` is `false`. -/// -/// Attaches `device_id` to the error for per-device diagnostics. -pub fn device_assert(device_id: usize, pred: bool, message: &str) -> Result<(), DeviceError> { - if !pred { - Err(DeviceError::Context { - device_id, - message: message.to_string(), - }) - } else { - Ok(()) - } -} - -/// Constructs a [`DeviceError::Context`] for the given device and message. -pub fn device_error(device_id: usize, message: &str) -> DeviceError { - DeviceError::Context { - device_id, - message: message.to_string(), - } -} diff --git a/cuda-async/src/simt/launch.rs b/cuda-async/src/simt/launch.rs deleted file mode 100644 index de151fe3cd..0000000000 --- a/cuda-async/src/simt/launch.rs +++ /dev/null @@ -1,570 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! CUDA kernel launch builder with argument marshalling. -//! -//! [`AsyncKernelLaunchBuilder`] safely accumulates a kernel function reference, -//! launch attributes, and type-erased argument pointers. An explicit unsafe -//! call to [`AsyncKernelLaunchBuilder::finalize_unchecked`] supplies the raw -//! launch geometry and returns an immutable [`AsyncKernelLaunch`] that can be -//! scheduled as a [`DeviceOperation`]. -//! -//! Arguments are heap-allocated via [`KernelArgument::push_arg`] and freed when -//! the launcher is dropped after submission. This keeps the pointed-to values -//! alive until `cuLaunchKernel` / `cuLaunchKernelEx` has copied the launch -//! parameter values out of the host-side argument array. -//! -//! [`DeviceOperation`]: crate::simt::device_operation::DeviceOperation - -use crate::simt::device_context::with_default_device_policy; -use crate::simt::device_future::DeviceFuture; -use crate::simt::device_operation::{DeviceOperation, ExecutionContext}; -use crate::simt::error::DeviceError; -use crate::simt::scheduling_policies::SchedulingPolicy; -use cuda_core::simt::LaunchConfig; -use cuda_core::{CudaFunction, CudaStream}; -use std::ffi::c_void; -use std::future::IntoFuture; -use std::marker::PhantomData; -use std::sync::Arc; - -/// Safe, inert builder for an unchecked CUDA kernel launch. -/// -/// Building a launch does not enqueue GPU work. The builder deliberately does -/// not implement [`DeviceOperation`] or [`IntoFuture`]: raw launch geometry -/// must first cross the explicit unsafe boundary at -/// [`finalize_unchecked`](Self::finalize_unchecked). -/// -/// ```compile_fail,E0277 -/// use cuda_async::simt::device_operation::DeviceOperation; -/// use cuda_async::simt::launch::AsyncKernelLaunchBuilder; -/// -/// fn schedule(_operation: O) {} -/// -/// fn cannot_schedule(builder: AsyncKernelLaunchBuilder<'static>) { -/// schedule(builder); -/// } -/// ``` -/// -/// Builders stay on the thread that assembles their type-erased argument -/// storage. Only the immutable finalized operation is `Send`: -/// -/// ```compile_fail,E0277 -/// use cuda_async::simt::launch::AsyncKernelLaunchBuilder; -/// -/// fn assert_send() {} -/// assert_send::>(); -/// ``` -#[derive(Debug)] -pub struct AsyncKernelLaunchBuilder<'a> { - /// Handle to the compiled device function. - func: Arc, - /// Heap-allocated, type-erased argument storage passed to the CUDA driver. - args: KernelArgStorage, - /// Optional thread-block cluster dimensions for `cuLaunchKernelEx`. - cluster_dim: Option<(u32, u32, u32)>, - /// When `true`, launch via `cuLaunchKernelEx` with - /// `CU_LAUNCH_ATTRIBUTE_COOPERATIVE` set, which guarantees the whole grid - /// is co-resident on the device (required for `cuda_device::grid::sync()`). - cooperative: bool, - /// Ties borrowed device buffers to the lazy launch operation. - _borrows: PhantomData<&'a mut ()>, -} - -/// Immutable, runnable CUDA kernel launch. -/// -/// This type can only be created by consuming an -/// [`AsyncKernelLaunchBuilder`]. It exposes no safe geometry or argument -/// setters, so its safety assumptions cannot be invalidated after the raw -/// launch configuration has been accepted. -/// -/// ```compile_fail,E0599 -/// use cuda_async::simt::launch::AsyncKernelLaunch; -/// use cuda_core::simt::LaunchConfig; -/// -/// fn cannot_reconfigure(launch: &mut AsyncKernelLaunch<'_>, config: LaunchConfig) { -/// launch.set_launch_config(config); -/// } -/// ``` -#[derive(Debug)] -pub struct AsyncKernelLaunch<'a> { - /// Handle to the compiled device function. - func: Arc, - /// Heap-allocated, type-erased argument storage passed to the CUDA driver. - args: KernelArgStorage, - /// Complete grid/block dimensions and dynamic shared-memory size. - cfg: LaunchConfig, - /// Optional thread-block cluster dimensions for `cuLaunchKernelEx`. - cluster_dim: Option<(u32, u32, u32)>, - /// Whether this launch requests cooperative grid residency. - cooperative: bool, - /// Ties borrowed device buffers to the lazy launch operation. - _borrows: PhantomData<&'a mut ()>, -} - -/// # Safety -/// -/// The argument pointers refer to heap allocations owned by the launch. -/// Non-`Copy` values enter that storage only through `KernelArgument: Send`; -/// copied scalar values have no destructor. The runnable type exposes none of -/// those erased values to safe Rust. The `Arc` is `Send + Sync`, -/// and `_borrows` preserves source lifetimes captured while building. -unsafe impl<'a> Send for AsyncKernelLaunch<'a> {} - -#[derive(Debug, Default)] -struct KernelArgStorage { - ptrs: Vec<*mut c_void>, - drops: Vec, -} - -impl Drop for KernelArgStorage { - fn drop(&mut self) { - for (arg, drop_arg) in self.ptrs.drain(..).zip(self.drops.drain(..)) { - unsafe { drop_arg(arg) }; - } - } -} - -impl KernelArgStorage { - fn push_send_arg(&mut self, arg: Box) { - unsafe fn drop_box(arg: *mut c_void) { - let _ = unsafe { Box::from_raw(arg as *mut T) }; - } - - self.ptrs.push(Box::into_raw(arg) as *mut c_void); - self.drops.push(drop_box::); - } - - fn push_scalar_arg(&mut self, arg: T) { - unsafe fn drop_copy_box(arg: *mut c_void) { - let _ = unsafe { Box::from_raw(arg as *mut T) }; - } - - self.ptrs.push(Box::into_raw(Box::new(arg)) as *mut c_void); - self.drops.push(drop_copy_box::); - } - - fn as_mut_slice(&mut self) -> &mut [*mut c_void] { - &mut self.ptrs - } -} - -impl<'a> AsyncKernelLaunchBuilder<'a> { - /// Creates an inert builder for `func` with no arguments or launch - /// attributes. - pub fn new(func: Arc) -> Self { - Self { - func, - args: KernelArgStorage::default(), - cluster_dim: None, - cooperative: false, - _borrows: PhantomData, - } - } - - /// Appends a by-value `Copy` argument to the kernel launch packet. - /// - /// This is the typed-module path used for scalar, raw-pointer, custom - /// `Copy` structs, and `Copy` closure arguments. - #[inline(always)] - pub fn push_scalar_arg(&mut self, arg: T) -> &mut Self { - self.args.push_scalar_arg(arg); - self - } - - /// Appends a kernel argument. The value is heap-allocated and its pointer - /// stored for the driver call. - /// - /// Scalars like `u32`, `f32`, `u64` etc. are auto-boxed -- no need to - /// wrap them in `Box::new`. - /// - /// The allocated storage remains alive until launch submission finishes or - /// the builder is dropped without launching. - /// - /// Values with thread-affine ownership cannot be hidden in the erased - /// storage that will later become a `Send` operation: - /// - /// ```compile_fail,E0277 - /// use cuda_async::simt::launch::AsyncKernelLaunchBuilder; - /// use std::rc::Rc; - /// - /// fn reject_rc(mut builder: AsyncKernelLaunchBuilder<'static>) { - /// builder.push_arg(Box::new(Rc::new(1_u32))); - /// } - /// ``` - #[inline(always)] - pub fn push_arg(&mut self, arg: T) -> &mut Self { - arg.push_arg(self); - self - } - - /// Appends multiple kernel arguments at once from a tuple. - /// - /// Equivalent to chained [`push_arg`](Self::push_arg) calls but allows - /// grouping all arguments in a single expression: - /// - /// ```ignore - /// launch.push_args((m, n, k, alpha, a_ptr, a_len, b_ptr, b_len, beta, c_ptr, c_len)); - /// ``` - /// - /// Supports tuples up to 32 elements. - #[inline(always)] - pub fn push_args(&mut self, args: T) -> &mut Self { - args.push_args(self); - self - } - - /// Sets thread-block cluster dimensions for a cluster launch. - pub fn set_cluster_dim(&mut self, cluster_dim: (u32, u32, u32)) -> &mut Self { - self.cluster_dim = Some(cluster_dim); - self - } - - /// Marks this launch as cooperative (`CU_LAUNCH_ATTRIBUTE_COOPERATIVE`). - /// - /// A cooperative launch guarantees every block in the grid is co-resident - /// on the device, which is the precondition for grid-wide barriers like - /// `cuda_device::grid::sync()`. May be combined with - /// [`set_cluster_dim`](Self::set_cluster_dim); both attributes are then - /// passed to `cuLaunchKernelEx` in the same call. - pub fn set_cooperative(&mut self, cooperative: bool) -> &mut Self { - self.cooperative = cooperative; - self - } - - /// Accepts an unverified raw launch configuration and seals this builder - /// into an immutable, schedulable operation. - /// - /// Prefer the generated typed module APIs for application code. They check - /// a kernel's declared launch contract before constructing the runnable - /// operation. - /// - /// # Safety - /// - /// The caller must prove all facts that a typed prepared launch normally - /// checks: - /// - /// - `cfg` must match the kernel's indexing dimensionality, launch bounds, - /// block shape, and dynamic shared-memory requirements; - /// - the cluster and cooperative attributes must be legal for the kernel; - /// - the function signature must exactly match the arguments, in order, - /// type, size, and alignment; - /// - referenced device allocations must remain valid for the complete GPU - /// operation and satisfy Rust's aliasing requirements; and - /// - the scheduling policy must select a stream from the CUDA context that - /// owns `func`. - /// - /// Violating these requirements can cause memory corruption, data races, - /// or CUDA driver errors. - /// - /// ```compile_fail,E0133 - /// use cuda_async::simt::launch::AsyncKernelLaunchBuilder; - /// use cuda_core::simt::LaunchConfig; - /// - /// fn raw_finalize_requires_unsafe( - /// builder: AsyncKernelLaunchBuilder<'static>, - /// config: LaunchConfig, - /// ) { - /// let _ = builder.finalize_unchecked(config); - /// } - /// ``` - pub unsafe fn finalize_unchecked(self, cfg: LaunchConfig) -> AsyncKernelLaunch<'a> { - AsyncKernelLaunch { - func: self.func, - args: self.args, - cfg, - cluster_dim: self.cluster_dim, - cooperative: self.cooperative, - _borrows: self._borrows, - } - } -} - -impl<'a> AsyncKernelLaunch<'a> { - /// Submits the kernel to `stream` via `cuLaunchKernel`, or via - /// `cuLaunchKernelEx` when cluster dimensions or a cooperative launch - /// were requested. - /// - /// # Safety - /// - /// - `self.func` must refer to a kernel loaded from the same CUDA context - /// that owns `stream`. - /// - All argument pointers in `self.args` must point to correctly typed and - /// aligned host-side values for the kernel's formal parameters. - /// - The pointed-to argument values must remain valid until launch - /// submission returns. The stream-aware launch helper binds the correct - /// context before calling into the CUDA driver. - /// - /// # Errors - /// - /// Returns [`DeviceError::Driver`] if context binding or launch submission - /// fails. - unsafe fn launch(mut self, stream: &Arc) -> Result<(), DeviceError> { - let cfg = self.cfg; - let result = match (self.cluster_dim, self.cooperative) { - (Some(cluster_dim), true) => unsafe { - cuda_core::launch_kernel_ex_cooperative_on_stream( - self.func.as_ref(), - cfg.grid_dim, - cfg.block_dim, - cfg.shared_mem_bytes, - cluster_dim, - stream.as_ref(), - self.args.as_mut_slice(), - ) - }, - (Some(cluster_dim), false) => unsafe { - cuda_core::launch_kernel_ex_on_stream( - self.func.as_ref(), - cfg.grid_dim, - cfg.block_dim, - cfg.shared_mem_bytes, - cluster_dim, - stream.as_ref(), - self.args.as_mut_slice(), - ) - }, - (None, true) => unsafe { - cuda_core::launch_kernel_cooperative_on_stream( - self.func.as_ref(), - cfg.grid_dim, - cfg.block_dim, - cfg.shared_mem_bytes, - stream.as_ref(), - self.args.as_mut_slice(), - ) - }, - (None, false) => unsafe { - cuda_core::launch_kernel_on_stream( - self.func.as_ref(), - cfg.grid_dim, - cfg.block_dim, - cfg.shared_mem_bytes, - stream.as_ref(), - self.args.as_mut_slice(), - ) - }, - }; - result.map_err(DeviceError::Driver)?; - Ok(()) - } -} - -/// Owns resources for a lazy async kernel launch and returns them when the GPU -/// work has completed. -#[derive(Debug)] -pub struct OwnedAsyncKernelLaunch { - launch: AsyncKernelLaunch<'static>, - resources: R, -} - -unsafe impl Send for OwnedAsyncKernelLaunch {} - -impl OwnedAsyncKernelLaunch { - /// Creates an owned async kernel operation from a prepared launch and the - /// resources that must stay alive until completion. - pub fn new(launch: AsyncKernelLaunch<'static>, resources: R) -> Self { - Self { launch, resources } - } -} - -/// Trait for types that can be marshalled into a CUDA kernel argument list. -/// -/// Implementors heap-allocate the value and push a `*mut c_void` into the -/// launcher's argument vector. -/// -/// Implemented for all common scalar primitives (`u8`–`u64`, `i8`–`i64`, -/// `f32`, `f64`, `usize`, `isize`, `bool`) so callers can write -/// `.push_arg(42u32)` without manual `Box::new`. -pub trait KernelArgument: Send { - /// Heap-allocates `self` and appends the pointer to `launcher.args`. - fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>); -} - -/// Passes the box's raw pointer directly as a kernel argument. -/// -/// This is the low-level escape hatch. Prefer passing scalars directly -- they -/// are auto-boxed via the blanket scalar impls. -impl KernelArgument for Box { - fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { - launcher.args.push_send_arg(self); - } -} - -macro_rules! impl_scalar_kernel_arg { - ($($t:ty),*) => { - $( - impl KernelArgument for $t { - #[inline(always)] - fn push_arg(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { - launcher.push_scalar_arg(self); - } - } - )* - }; -} - -impl_scalar_kernel_arg!(u8, u16, u32, u64, usize, i8, i16, i32, i64, isize, f32, f64, bool); - -// --------------------------------------------------------------------------- -// KernelArguments — push multiple heterogeneous args in a single call -// --------------------------------------------------------------------------- - -/// Trait for a group of kernel arguments that can be pushed together. -/// -/// Implemented for tuples of [`KernelArgument`] types up to arity 32, enabling -/// `launch.push_args((dim_m, dim_n, alpha, ptr, len))` as an alternative to -/// chained `.push_arg()` calls. -#[diagnostic::on_unimplemented( - message = "cannot push `{Self}` as kernel arguments", - note = "KernelArguments is implemented for tuples of KernelArgument types up to 32 elements" -)] -pub trait KernelArguments { - /// Pushes every element into `launcher` in order. - fn push_args(self, launcher: &mut AsyncKernelLaunchBuilder<'_>); -} - -macro_rules! impl_kernel_args_tuple { - // Base case: empty tuple - () => { - impl KernelArguments for () { - #[inline(always)] - fn push_args(self, _launcher: &mut AsyncKernelLaunchBuilder<'_>) {} - } - }; - // Recursive case: (A, B, C, ...) where each element is a KernelArgument - ($($idx:tt : $T:ident),+) => { - impl<$($T: KernelArgument),+> KernelArguments for ($($T,)+) { - #[inline(always)] - fn push_args(self, launcher: &mut AsyncKernelLaunchBuilder<'_>) { - $(launcher.push_arg(self.$idx);)+ - } - } - }; -} - -impl_kernel_args_tuple!(); -impl_kernel_args_tuple!(0: A); -impl_kernel_args_tuple!(0: A, 1: B); -impl_kernel_args_tuple!(0: A, 1: B, 2: C); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD, 30: AE); -impl_kernel_args_tuple!(0: A, 1: B, 2: C, 3: D, 4: E, 5: F, 6: G, 7: H, 8: I, 9: J, 10: K, 11: L, 12: M, 13: N, 14: O, 15: P, 16: Q, 17: R, 18: S, 19: T, 20: U, 21: V, 22: W, 23: X, 24: Y, 25: Z, 26: AA, 27: AB, 28: AC, 29: AD, 30: AE, 31: AF); - -/// Launches the kernel on the stream bound to the execution context. -impl<'a> DeviceOperation for AsyncKernelLaunch<'a> { - type Output = (); - - unsafe fn execute(self, ctx: &ExecutionContext) -> Result<(), DeviceError> { - unsafe { self.launch(ctx.get_cuda_stream()) } - } -} - -/// Schedules the kernel launch via the thread-local default scheduling policy. -impl<'a> IntoFuture for AsyncKernelLaunch<'a> { - type Output = Result<(), DeviceError>; - type IntoFuture = DeviceFuture<(), AsyncKernelLaunch<'a>>; - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -impl DeviceOperation for OwnedAsyncKernelLaunch { - type Output = R; - - unsafe fn execute(self, ctx: &ExecutionContext) -> Result { - let Self { launch, resources } = self; - unsafe { launch.launch(ctx.get_cuda_stream()) }?; - Ok(resources) - } -} - -impl IntoFuture for OwnedAsyncKernelLaunch { - type Output = Result; - type IntoFuture = DeviceFuture>; - - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| policy.schedule(self)) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - #[repr(C)] - #[derive(Clone, Copy, Debug, PartialEq)] - struct TestParams { - scale: f32, - bias: i32, - } - - #[test] - fn scalar_arg_storage_accepts_custom_copy_value() { - let mut storage = KernelArgStorage::default(); - let params = TestParams { - scale: 2.0, - bias: 3, - }; - - storage.push_scalar_arg(params); - - assert_eq!(storage.ptrs.len(), 1); - assert_eq!(unsafe { *(storage.ptrs[0] as *const TestParams) }, params); - } - - #[test] - fn arg_storage_drops_values_with_their_original_type() { - struct DropCounter(Arc); - - impl Drop for DropCounter { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - let drops = Arc::new(AtomicUsize::new(0)); - let mut storage = KernelArgStorage::default(); - - storage.push_send_arg(Box::new(DropCounter(Arc::clone(&drops)))); - assert_eq!(drops.load(Ordering::Relaxed), 0); - - drop(storage); - assert_eq!(drops.load(Ordering::Relaxed), 1); - } -} diff --git a/cuda-async/src/simt/mod.rs b/cuda-async/src/simt/mod.rs deleted file mode 100644 index e8e8f62860..0000000000 --- a/cuda-async/src/simt/mod.rs +++ /dev/null @@ -1,27 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! The cuda-oxide async surface. -//! -//! Everything under this module is the cuda-oxide `cuda-async` API, brought -//! over as-is for the shared host-crate migration. The two async stacks fork -//! the same ancestor, so they share most type names with diverged -//! definitions; nothing here is re-exported at the crate root, and consumers -//! address it as `cuda_async::simt::::`, mirroring oxide's -//! `cuda_async::::`. Reconciliation happens after the -//! repository migration, and this module is deletable as a unit. -//! -//! One departure, recorded for the reconciliation: oxide's `zip!` and -//! `unzip!` macros are present but not exported, since this crate already -//! exports macros with those names. - -pub mod device_box; -pub mod device_context; -pub mod device_future; -pub mod device_operation; -pub mod error; -pub mod launch; -pub mod reclaim; -pub mod scheduling_policies; diff --git a/cuda-async/src/simt/reclaim.rs b/cuda-async/src/simt/reclaim.rs deleted file mode 100644 index 24f04923b4..0000000000 --- a/cuda-async/src/simt/reclaim.rs +++ /dev/null @@ -1,308 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Deferred reclamation of results cancelled while GPU work is in flight. -//! -//! Cancelling a [`DeviceFuture`] cannot cancel GPU work that has already -//! been submitted: kernels run to completion regardless of what the host -//! does. What cancellation *can* do is decide when the host releases the -//! resources the kernel is still using. Dropping them immediately is -//! unsound (the stream-ordered allocator may hand the memory to the next -//! allocation while the kernel still writes to it), and synchronizing the -//! stream inside `Drop` blocks executor threads for an unbounded time. -//! -//! Instead, cancelled in-flight results are *parked* here together with a -//! completion gate (a CUDA event recorded on the assigned stream after the -//! submitted work): -//! -//! ```text -//! drop(future) later sweep (poll / drop / drain) -//! record event on stream event passed? -> drop the result -//! park (event, result) still running? -> keep it parked -//! ``` -//! -//! Parked results are swept opportunistically on subsequent future polls -//! and drops, or deterministically via [`drain`]. Payloads are never -//! dropped inside a CUDA host callback (callbacks must not call CUDA APIs, -//! and dropping a device buffer enqueues an async free), and a sweep never -//! blocks on GPU progress. -//! -//! Entries that are still parked at process exit are leaked; the driver -//! reclaims device memory at process teardown, while dropping them early -//! could corrupt memory that is still in use. -//! -//! [`DeviceFuture`]: crate::simt::device_future::DeviceFuture - -use cuda_core::{CudaEvent, DriverError}; -use std::mem; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Mutex, MutexGuard}; - -/// Completion gate guarding a parked result. -/// -/// Production code uses a [`CudaEvent`] recorded on the stream that runs -/// the cancelled work. Tests substitute host-side mocks. -pub(crate) trait ReclaimGate: Send { - /// Returns `true` when the device timeline has passed the gate. - /// Must never block. - fn passed(&self) -> Result; - - /// Blocks until the device timeline has passed the gate. - fn wait(&self) -> Result<(), DriverError>; -} - -impl ReclaimGate for CudaEvent { - fn passed(&self) -> Result { - self.query() - } - - fn wait(&self) -> Result<(), DriverError> { - self.synchronize() - } -} - -/// A parked result and the gate that decides when it may be dropped. -struct LimboEntry { - /// Completion gate; the payload may only drop after it has passed. - gate: Box, - /// Type-erased result kept alive on behalf of the cancelled future. - payload: Box, -} - -/// Number of parked entries. Mirrors `LIMBO.len()`; updated only while the -/// limbo lock is held, read lock-free as the sweep fast path. -static PENDING: AtomicUsize = AtomicUsize::new(0); - -/// All currently parked entries, across every device and thread. -static LIMBO: Mutex> = Mutex::new(Vec::new()); - -/// Locks the limbo list, recovering from poisoning. -/// -/// A poisoned lock only means another thread panicked while holding it; -/// the list itself stays consistent because payloads are always dropped -/// outside the lock. -fn limbo() -> MutexGuard<'static, Vec> { - LIMBO - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -/// Parks `payload` until `gate` reports completion. -pub(crate) fn park(gate: impl ReclaimGate + 'static, payload: Box) { - let mut entries = limbo(); - entries.push(LimboEntry { - gate: Box::new(gate), - payload, - }); - PENDING.store(entries.len(), Ordering::Release); -} - -/// Number of results currently parked awaiting reclamation. -pub fn pending() -> usize { - PENDING.load(Ordering::Acquire) -} - -/// Drops every parked result whose gate has passed. Never blocks on GPU -/// progress. Returns the number of results reclaimed. -/// -/// Called automatically on every [`DeviceFuture`] poll and drop; safe to -/// call manually at any time. -/// -/// [`DeviceFuture`]: crate::simt::device_future::DeviceFuture -pub fn sweep() -> usize { - if PENDING.load(Ordering::Acquire) == 0 { - return 0; - } - let completed: Vec = { - let mut entries = limbo(); - let mut kept = Vec::with_capacity(entries.len()); - let mut done = Vec::new(); - for entry in entries.drain(..) { - match entry.gate.passed() { - Ok(true) => done.push(entry), - // Still in flight, or the query failed. Dropping early - // would release memory the device may still be using, so - // the entry stays parked. - Ok(false) | Err(_) => kept.push(entry), - } - } - *entries = kept; - PENDING.store(entries.len(), Ordering::Release); - done - }; - // Dropped outside the lock: a payload's own drop may park or sweep. - let reclaimed = completed.len(); - drop(completed); - reclaimed -} - -/// Blocking drain: waits for every parked gate to pass, then drops the -/// payload. Returns the number of results reclaimed. -/// -/// When even the blocking wait fails, the payload is deliberately leaked -/// with a message on stderr: if the driver cannot prove the GPU work -/// finished, leaking is safer than freeing memory the device may still -/// write to. -/// -/// Call this before tearing down a process or test when deterministic -/// reclamation is required. It must not be called from a CUDA host -/// callback or while holding a thread-local device-context borrow. -pub fn drain() -> usize { - let entries: Vec = { - let mut entries = limbo(); - let drained = mem::take(&mut *entries); - PENDING.store(0, Ordering::Release); - drained - }; - let mut reclaimed = 0; - for entry in entries { - let waited = match entry.gate.passed() { - Ok(true) => Ok(()), - _ => entry.gate.wait(), - }; - match waited { - Ok(()) => { - drop(entry.payload); - reclaimed += 1; - } - Err(error) => { - crate::leak::report_leak(format_args!( - "cuda-async: leaking a cancelled in-flight result; the driver \ - could not prove the GPU work finished: {error}" - )); - mem::forget(entry.payload); - } - } - } - reclaimed -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::AtomicBool; - use std::sync::Arc; - - /// Serializes the tests in this module: they share the global limbo, - /// and `drain` consumes (and waits on) every parked entry. - static TEST_LOCK: Mutex<()> = Mutex::new(()); - - fn test_lock() -> MutexGuard<'static, ()> { - TEST_LOCK - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - } - - /// Host-side stand-in for a recorded CUDA event. - struct MockGate { - /// Mirrors "the device timeline has passed the recorded event". - passed: Arc, - /// When `true`, the blocking wait reports a driver failure. - wait_fails: bool, - } - - impl ReclaimGate for MockGate { - fn passed(&self) -> Result { - Ok(self.passed.load(Ordering::Relaxed)) - } - - fn wait(&self) -> Result<(), DriverError> { - if self.wait_fails { - Err(DriverError( - cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, - )) - } else { - self.passed.store(true, Ordering::Relaxed); - Ok(()) - } - } - } - - struct CountDrop(Arc); - - impl Drop for CountDrop { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - #[test] - fn sweep_keeps_payload_parked_until_gate_passes() { - let _guard = test_lock(); - let drops = Arc::new(AtomicUsize::new(0)); - let gate_passed = Arc::new(AtomicBool::new(false)); - - park( - MockGate { - passed: Arc::clone(&gate_passed), - wait_fails: false, - }, - Box::new(CountDrop(Arc::clone(&drops))), - ); - - sweep(); - assert_eq!( - drops.load(Ordering::Relaxed), - 0, - "payload must stay parked while the GPU work is in flight" - ); - - gate_passed.store(true, Ordering::Relaxed); - sweep(); - assert_eq!( - drops.load(Ordering::Relaxed), - 1, - "payload must drop once the gate reports completion" - ); - } - - #[test] - fn drain_waits_for_gate_then_drops_payload() { - let _guard = test_lock(); - let drops = Arc::new(AtomicUsize::new(0)); - let gate_passed = Arc::new(AtomicBool::new(false)); - - park( - MockGate { - passed: Arc::clone(&gate_passed), - wait_fails: false, - }, - Box::new(CountDrop(Arc::clone(&drops))), - ); - - let reclaimed = drain(); - assert_eq!(reclaimed, 1); - assert_eq!(drops.load(Ordering::Relaxed), 1); - assert!( - gate_passed.load(Ordering::Relaxed), - "drain must have waited on the gate before dropping" - ); - } - - #[test] - fn drain_leaks_payload_when_wait_fails() { - let _guard = test_lock(); - let drops = Arc::new(AtomicUsize::new(0)); - - park( - MockGate { - passed: Arc::new(AtomicBool::new(false)), - wait_fails: true, - }, - Box::new(CountDrop(Arc::clone(&drops))), - ); - - let mut capture = crate::leak::capture::start(); - let reclaimed = drain(); - let reports = capture.take(); - assert_eq!(reclaimed, 0); - assert_eq!( - drops.load(Ordering::Relaxed), - 0, - "an unprovable gate must leak the payload, never drop it early" - ); - assert_eq!(reports.len(), 1, "the leak must be reported: {reports:?}"); - } -} diff --git a/cuda-async/src/simt/scheduling_policies.rs b/cuda-async/src/simt/scheduling_policies.rs deleted file mode 100644 index a32d1ee885..0000000000 --- a/cuda-async/src/simt/scheduling_policies.rs +++ /dev/null @@ -1,251 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Stream scheduling policies that control how [`DeviceOperation`]s are -//! assigned to CUDA streams. -//! -//! A [`SchedulingPolicy`] decouples operation construction from stream -//! selection. The policy chooses a stream when an operation is scheduled or -//! executed synchronously, enabling transparent overlap of independent work. -//! -//! | Policy | Behaviour | -//! |--------------------------|----------------------------------------------| -//! | [`StreamPoolRoundRobin`] | Rotates across *N* streams for HW overlap | -//! | [`SingleStream`] | Serialises all work onto one stream | -//! -//! [`DeviceOperation`]: crate::simt::device_operation::DeviceOperation - -use crate::simt::device_future::DeviceFuture; -use crate::simt::device_operation::{DeviceOperation, ExecutionContext}; -use crate::simt::error::{device_error, DeviceError}; -use cuda_core::{CudaContext, CudaStream}; -use std::sync::atomic::AtomicUsize; -use std::sync::Arc; - -/// Top-level enum wrapping all supported scheduling policy implementations. -/// -/// Used as the concrete policy type stored in -/// [`AsyncDeviceContext`](crate::simt::device_context::AsyncDeviceContext). -pub enum GlobalSchedulingPolicy { - /// Round-robin across a pool of CUDA streams. - RoundRobin(StreamPoolRoundRobin), -} - -impl GlobalSchedulingPolicy { - /// Returns a reference to the inner policy as a trait object. - pub fn as_scheduling_policy(&self) -> Result<&impl SchedulingPolicy, DeviceError> { - match self { - GlobalSchedulingPolicy::RoundRobin(rr) => Ok(rr), - } - } -} - -/// Delegates to the inner policy variant. -impl SchedulingPolicy for GlobalSchedulingPolicy { - fn init(&mut self, ctx: &Arc) -> Result<(), DeviceError> { - match self { - GlobalSchedulingPolicy::RoundRobin(rr) => rr.init(ctx), - } - } - fn schedule>( - &self, - op: O, - ) -> Result, DeviceError> { - match self { - GlobalSchedulingPolicy::RoundRobin(rr) => rr.schedule(op), - } - } - fn sync>(&self, op: O) -> Result { - match self { - GlobalSchedulingPolicy::RoundRobin(rr) => rr.sync(op), - } - } -} - -/// `Arc`-wrapped policy delegates scheduling but rejects re-initialization, -/// since the inner data cannot be mutated through an `Arc`. -impl SchedulingPolicy for Arc { - fn init(&mut self, _ctx: &Arc) -> Result<(), DeviceError> { - Err(DeviceError::Scheduling( - "Cannot initialize scheduling policy inside an Arc.".to_string(), - )) - } - fn schedule>( - &self, - op: O, - ) -> Result, DeviceError> { - match self.as_ref() { - GlobalSchedulingPolicy::RoundRobin(rr) => rr.schedule(op), - } - } - fn sync>(&self, op: O) -> Result { - match self.as_ref() { - GlobalSchedulingPolicy::RoundRobin(rr) => rr.sync(op), - } - } -} - -/// Strategy for assigning [`DeviceOperation`]s to CUDA streams. -/// -/// Implementations must be `Sync` because the policy is shared across all -/// operations on a device context. -/// -/// [`DeviceOperation`]: crate::simt::device_operation::DeviceOperation -pub trait SchedulingPolicy: Sync { - /// One-time initialization with the CUDA context (creates streams, etc.). - fn init(&mut self, ctx: &Arc) -> Result<(), DeviceError>; - - /// Assigns a stream to `op` and returns a [`DeviceFuture`] ready to be - /// polled. - fn schedule>( - &self, - op: O, - ) -> Result, DeviceError>; - - /// Executes `op` synchronously on a policy-chosen stream and blocks until - /// the stream is idle. - fn sync>(&self, op: O) -> Result; -} - -/// Round-robin scheduler over a fixed-size pool of CUDA streams. -/// -/// Each call to [`schedule`](SchedulingPolicy::schedule) or -/// [`sync`](SchedulingPolicy::sync) atomically advances the stream index, -/// distributing work across streams to enable hardware-level overlap of -/// independent kernels and memory transfers. -#[derive(Debug)] -pub struct StreamPoolRoundRobin { - /// Device ordinal this pool belongs to. - device_id: usize, - /// Monotonically increasing counter; wrapped with `% num_streams`. - next_stream_idx: AtomicUsize, - /// Number of streams in the pool. - pub(crate) num_streams: usize, - /// `None` until [`init`](SchedulingPolicy::init) is called. - pub(crate) stream_pool: Option>>, -} - -impl StreamPoolRoundRobin { - /// Creates an un-initialized round-robin policy for `device_id` with - /// `num_streams` streams. - /// - /// # Safety - /// - /// The caller must call [`SchedulingPolicy::init`] before scheduling any - /// operations. Using the policy before initialization produces an error. - pub unsafe fn new(device_id: usize, num_streams: usize) -> Self { - Self { - device_id, - num_streams, - stream_pool: None, - next_stream_idx: AtomicUsize::new(0), - } - } -} - -impl SchedulingPolicy for StreamPoolRoundRobin { - /// Allocates `num_streams` CUDA streams on the provided context. - fn init(&mut self, ctx: &Arc) -> Result<(), DeviceError> { - let mut pool = vec![]; - for _ in 0..self.num_streams { - pool.push(ctx.new_stream()?); - } - self.stream_pool = Some(pool); - Ok(()) - } - - /// Picks the next stream (round-robin), executes `op`, and synchronizes. - fn sync>(&self, op: O) -> Result { - let idx = self - .next_stream_idx - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) - % self.num_streams; - let pool = self - .stream_pool - .as_ref() - .ok_or_else(|| device_error(self.device_id, "Stream pool not initialized."))?; - op.sync_on(&pool[idx]) - } - - /// Picks the next stream (round-robin) and wraps `op` in a - /// [`DeviceFuture`]. - fn schedule>( - &self, - op: O, - ) -> Result, DeviceError> { - let idx = self - .next_stream_idx - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) - % self.num_streams; - let pool = self - .stream_pool - .as_ref() - .ok_or_else(|| device_error(self.device_id, "Stream pool not initialized."))?; - // `DeviceFuture` implements `Drop`, so functional-update syntax - // (`..Default::default()`) would move fields out of a droppable - // value (E0509). Spell every field out instead. - Ok(DeviceFuture { - device_operation: Some(op), - execution_context: Some(ExecutionContext::new(Arc::clone(&pool[idx]))), - result: None, - error: None, - state: Default::default(), - callback_state: None, - }) - } -} - -/// Single-stream scheduler that serialises all work onto one CUDA stream. -/// -/// Useful for debugging or when ordering guarantees are required between -/// all operations on a device. -#[derive(Debug)] -pub struct SingleStream { - /// The single stream. `None` until [`init`](SchedulingPolicy::init). - pub stream: Option>, -} - -impl SingleStream { - /// Creates an un-initialized single-stream policy. - /// - /// # Safety - /// - /// The caller must call [`SchedulingPolicy::init`] before use. - pub unsafe fn new() -> Self { - Self { stream: None } - } -} - -impl SchedulingPolicy for SingleStream { - /// Allocates one CUDA stream on `ctx`. - fn init(&mut self, ctx: &Arc) -> Result<(), DeviceError> { - self.stream = Some(ctx.new_stream()?); - Ok(()) - } - - /// Wraps `op` in a [`DeviceFuture`] bound to the single stream. - fn schedule>( - &self, - op: O, - ) -> Result, DeviceError> { - let stream = self.stream.as_ref().unwrap(); - // See `StreamPoolRoundRobin::schedule` for why the fields are - // spelled out instead of using `..Default::default()`. - Ok(DeviceFuture { - device_operation: Some(op), - execution_context: Some(ExecutionContext::new(Arc::clone(stream))), - result: None, - error: None, - state: Default::default(), - callback_state: None, - }) - } - - /// Executes `op` synchronously on the single stream. - fn sync>(&self, op: O) -> Result { - let stream = self.stream.as_ref().unwrap(); - op.sync_on(stream) - } -} diff --git a/cuda-async/src/submission.rs b/cuda-async/src/submission.rs index 8a2e3c890a..8899838325 100644 --- a/cuda-async/src/submission.rs +++ b/cuda-async/src/submission.rs @@ -76,13 +76,8 @@ impl Owners { Ok(()) } - fn release(&mut self, wait: impl FnOnce() -> Result<(), DeviceError>) { - let Some(owners) = self.0.take() else { return }; - if !owners.is_empty() && wait().is_err() { - // This includes access leases, not just allocations. Unblocking an - // access when completion is unknown would permit a device data race. - std::mem::forget(owners); - } + fn complete(&mut self) { + self.0.take(); } } @@ -134,43 +129,52 @@ impl Submission { std::mem::take(&mut *self.recorded.lock().unwrap_or_else(|e| e.into_inner())) } - /// No further work may be submitted using this submission. + /// Releases tracking; no further work may use this submission. + /// + /// # Safety + /// GPU work must be complete, or the caller must independently keep all + /// resources alive and prevent conflicting accesses until completion. pub(crate) unsafe fn complete(&self) { - let owners = self - .owners + self.owners .lock() .unwrap_or_else(|e| e.into_inner()) - .0 - .take(); - drop(owners); + .complete(); } } impl Drop for Submission { fn drop(&mut self) { - self.owners + let Some(owners) = self + .owners .get_mut() .unwrap_or_else(|e| e.into_inner()) - .release(|| { - let result = (|| { - self.stream.device().bind_to_thread()?; - match probe_stream(&self.stream) { - StreamHealth::Idle => Ok(()), - StreamHealth::Busy => { - unsafe { self.stream.synchronize() }.map_err(DeviceError::Driver) - } - StreamHealth::Faulted(e) => Err(DeviceError::Driver(e)), - StreamHealth::Capturing => Err(DeviceError::Internal( - "submission is still being captured".into(), - )), - } - })(); - if result.is_err() { - // Keep the stream identity valid for leaked access leases too. - std::mem::forget(self.stream.clone()); - } - result - }); + .0 + .take() + else { + return; + }; + if owners.is_empty() { + return; + } + let stream = self.stream.clone(); + let health = match stream.context().bind_to_thread() { + Ok(()) => probe_stream(&stream), + Err(error) => StreamHealth::Faulted(error), + }; + match health { + StreamHealth::Idle => drop(owners), + StreamHealth::Busy => { + // CUDA events retain context; explicitly retain the stream too, + // so access tracker stream identities remain valid until release. + let _ = + cuda_core::completion::retain_until_complete((owners, stream.clone()), &stream); + } + StreamHealth::Faulted(_) | StreamHealth::Capturing => { + // Completion cannot be established. Preserve both allocations + // and exclusion instead of freeing storage or allowing a race. + std::mem::forget((owners, stream)); + } + } } } @@ -185,10 +189,7 @@ mod tests { let mut submission = Owners(Some(OwnerSet::default())); submission.retain(owner).unwrap(); assert!(weak.upgrade().is_some()); - submission.release(|| { - assert!(weak.upgrade().is_some()); - Ok(()) - }); + submission.complete(); assert!(weak.upgrade().is_none()); assert!(submission.retain(()).is_err()); } @@ -202,14 +203,4 @@ mod tests { std::mem::forget(submission); assert!(weak.upgrade().is_some()); } - - #[test] - fn failed_wait_keeps_owners_and_access_leases() { - let owner = Arc::new(()); - let weak = Arc::downgrade(&owner); - let mut submission = Owners(Some(OwnerSet::default())); - submission.retain(owner).unwrap(); - submission.release(|| Err(DeviceError::Internal("fault".into()))); - assert!(weak.upgrade().is_some()); - } } diff --git a/cuda-async/tests/concurrent_capture.rs b/cuda-async/tests/concurrent_capture.rs index e574ac2e8d..93c6403c5d 100644 --- a/cuda-async/tests/concurrent_capture.rs +++ b/cuda-async/tests/concurrent_capture.rs @@ -10,7 +10,7 @@ use cuda_async::device_operation::value; use cuda_async::error::DeviceError; fn has_gpu() -> bool { - cuda_core::Device::device_count() + cuda_core::Context::device_count() .map(|n| n > 0) .unwrap_or(false) } @@ -23,7 +23,7 @@ fn concurrent_capture_same_context() { return; } - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream_a = device.new_stream().unwrap(); let stream_b = device.new_stream().unwrap(); @@ -66,7 +66,7 @@ fn new_stream_during_capture_on_another_thread() { return; } - let device_a = cuda_core::Device::new(0).unwrap(); + let device_a = cuda_core::Context::new(0).unwrap(); let stream_a = device_a.new_stream().unwrap(); let barrier = std::sync::Arc::new(std::sync::Barrier::new(2)); @@ -86,7 +86,7 @@ fn new_stream_during_capture_on_another_thread() { let handle_b = std::thread::spawn(move || -> Result<(), DeviceError> { barrier_b.wait(); - let device_b = cuda_core::Device::new(0).unwrap(); + let device_b = cuda_core::Context::new(0).unwrap(); let _stream = device_b.new_stream()?; Ok(()) }); diff --git a/cuda-async/tests/cuda_graph.rs b/cuda-async/tests/cuda_graph.rs index 04aafe04ed..ff7544dcdf 100644 --- a/cuda-async/tests/cuda_graph.rs +++ b/cuda-async/tests/cuda_graph.rs @@ -13,7 +13,7 @@ use cuda_async::error::DeviceError; use std::future::IntoFuture; fn has_gpu() -> bool { - cuda_core::Device::device_count() + cuda_core::Context::device_count() .map(|n| n > 0) .unwrap_or(false) } @@ -28,7 +28,7 @@ fn scope_empty_closure() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let graph = CudaGraph::scope(&stream, |_s| Ok(())).unwrap(); @@ -42,7 +42,7 @@ fn scope_records_value_ops() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let mut recorded = Vec::new(); @@ -66,7 +66,7 @@ fn scope_error_propagation() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let result = CudaGraph::scope(&stream, |_s| { @@ -93,7 +93,7 @@ fn scope_panic_safety() { return; } let result = std::thread::spawn(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { @@ -103,7 +103,7 @@ fn scope_panic_safety() { })); // Stream should still be usable after the panic. - unsafe { stream.synchronize() }.unwrap(); + stream.synchronize().unwrap(); }) .join(); @@ -119,7 +119,7 @@ fn scope_multiple_launches() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let graph = CudaGraph::scope(&stream, |_s| Ok(())).unwrap(); @@ -140,7 +140,7 @@ fn scope_nested_execution_rejected() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let other_stream = device.new_stream().unwrap(); @@ -179,7 +179,7 @@ fn launch_outlives_its_graph() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let graph = CudaGraph::scope(&stream, |s| { @@ -214,7 +214,7 @@ fn capture_rejects_nested_execution() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let mut graph = CudaGraph::capture( @@ -242,7 +242,7 @@ fn capture_panic_ends_capture_and_releases_lock() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { @@ -253,7 +253,9 @@ fn capture_panic_ends_capture_and_releases_lock() { })); assert!(panicked.is_err(), "the panic must propagate"); - unsafe { stream.synchronize() }.expect("stream must not be left in capture mode"); + stream + .synchronize() + .expect("stream must not be left in capture mode"); assert_eq!(value(2).sync_on(&stream).expect("lock must be free"), 2); let mut graph = CudaGraph::capture(stream.clone(), value(4)).expect("recapture"); assert_eq!(graph.take_output(), Some(4)); @@ -284,7 +286,7 @@ fn capture_error_ends_capture_and_releases_lock() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let err = match CudaGraph::capture(stream.clone(), FailingOp) { @@ -296,7 +298,9 @@ fn capture_error_ends_capture_and_releases_lock() { "op error must propagate unchanged, got {err:?}" ); - unsafe { stream.synchronize() }.expect("stream must not be left in capture mode"); + stream + .synchronize() + .expect("stream must not be left in capture mode"); assert_eq!(value(2).sync_on(&stream).expect("lock must be free"), 2); CudaGraph::capture(stream.clone(), value(4)).expect("recapture"); }); @@ -310,7 +314,7 @@ fn graph_combinators_capture_and_replay() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let mut graph = value(5).graph_on(stream.clone()).expect("graph_on"); @@ -327,7 +331,9 @@ fn graph_combinators_capture_and_replay() { .graph_on(stream.clone()) })); assert!(panicked.is_err()); - unsafe { stream.synchronize() }.expect("stream must not be left in capture mode"); + stream + .synchronize() + .expect("stream must not be left in capture mode"); assert_eq!(value(7).sync_on(&stream).expect("lock must be free"), 7); }); } @@ -345,7 +351,7 @@ fn update_runs_unit_graph_nodes() { return; } on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let graph = CudaGraph::scope(&stream, |_s| Ok(())).unwrap(); diff --git a/cuda-async/tests/device_fault.rs b/cuda-async/tests/device_fault.rs index 13e1f20277..f4d926cd69 100644 --- a/cuda-async/tests/device_fault.rs +++ b/cuda-async/tests/device_fault.rs @@ -18,7 +18,7 @@ use cuda_async::device_context::{global_policy, init_device_contexts, load_module_from_ptx}; use cuda_async::device_operation::{value, DeviceOp, ExecutionContext}; use cuda_async::error::DeviceError; -use cuda_async::launch::AsyncKernelLaunch; +use cuda_async::launch::AsyncKernelLaunchBuilder; use cuda_core::{Function, LaunchConfig}; use std::future::{Future, IntoFuture}; use std::pin::Pin; @@ -86,11 +86,13 @@ struct FaultingOp { bytes: usize, passes: usize, trap: Arc, + owner: Arc<()>, } impl DeviceOp for FaultingOp { type Output = (); unsafe fn execute(self, context: &ExecutionContext) -> Result<(), DeviceError> { + context.retain(self.owner)?; let stream = context.get_cuda_stream().cu_stream(); for _ in 0..self.passes { let code = cuda_bindings::cuMemsetD8Async(self.dptr, 0x11, self.bytes, stream); @@ -100,7 +102,7 @@ impl DeviceOp for FaultingOp { ))); } } - let mut launch = AsyncKernelLaunch::new(self.trap); + let mut launch = AsyncKernelLaunchBuilder::new(self.trap); launch.set_launch_config(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), @@ -134,11 +136,13 @@ fn device_fault_resolves_the_awaiting_future_with_the_driver_error() { // 1. Fault after registration. Before the fix this never resolved // (the flag write never lands; nothing wakes the future) and the // reactor spun at 100% CPU on the armed slot. + let owner = Arc::new(()); let op = FaultingOp { dptr, bytes, passes: 32, trap, + owner: owner.clone(), }; let started = Instant::now(); let result = block_on_with_deadline(op.into_future(), Duration::from_secs(30)); @@ -148,6 +152,12 @@ fn device_fault_resolves_the_awaiting_future_with_the_driver_error() { started.elapsed() ); + assert_eq!( + Arc::strong_count(&owner), + 2, + "a faulted submission must preserve its independently registered owners" + ); + // 2. The context is now dead. A fresh future's inline spin sees the // sticky error from `cuStreamQuery` and must propagate it (before // the fix the spin swallowed it and fell through to registration). diff --git a/cuda-async/tests/drop_in_flight.rs b/cuda-async/tests/drop_in_flight.rs index 55464399f8..00f4233904 100644 --- a/cuda-async/tests/drop_in_flight.rs +++ b/cuda-async/tests/drop_in_flight.rs @@ -4,7 +4,7 @@ */ //! Dropping a `DeviceFuture` while its GPU work is in flight must not release -//! the operation's output before the device has finished with it. +//! independently retained owners before the device has finished with them. //! //! The op enqueues a chain of memsets and then records a CUDA event on the //! same stream, handing the event back inside its output. The output's @@ -17,7 +17,7 @@ //! nothing behind the gate can run, whatever the device's memset throughput //! or the inline-spin budget (`CUDA_ASYNC_SPIN_BUDGET_US`), so the first poll //! is guaranteed to find the work in flight and the drop is guaranteed to -//! have to wait for it. Requires a GPU. +//! return while the retained owner stays alive. Requires a GPU. use cuda_async::device_context::{global_policy, init_device_contexts, with_device}; use cuda_async::device_operation::{DeviceOp, ExecutionContext}; @@ -177,10 +177,11 @@ struct SlowOp { } impl DeviceOp for SlowOp { - type Output = Tracked; - unsafe fn execute(self, context: &ExecutionContext) -> Result { + type Output = Arc; + unsafe fn execute(self, context: &ExecutionContext) -> Result, DeviceError> { let stream = context.get_cuda_stream(); if let Some(gate) = &self.gate { + context.retain(gate.clone())?; gate.arm(stream.cu_stream())?; } for _ in 0..self.passes { @@ -192,18 +193,20 @@ impl DeviceOp for SlowOp { ))); } } - let event = stream.device().new_event()?; + let event = stream.context().new_event()?; event.record(stream)?; - Ok(Tracked { + let tracked = Arc::new(Tracked { event, log: self.log, - }) + }); + context.retain(tracked.clone())?; + Ok(tracked) } } impl IntoFuture for SlowOp { - type Output = Result; - type IntoFuture = cuda_async::device_future::DeviceFuture; + type Output = Result, DeviceError>; + type IntoFuture = cuda_async::device_future::DeviceFuture, SlowOp>; fn into_future(self) -> Self::IntoFuture { let policy = global_policy(0).expect("global policy"); match self.schedule(&policy) { @@ -266,7 +269,7 @@ fn block_on_with_deadline(mut future: F, deadline: Duration) /// waits cannot return before then, and one that does not is caught twice: /// it returns early, and its output's event reports the work still in flight. #[test] -fn dropping_in_flight_future_releases_output_after_the_device_finished() { +fn dropping_in_flight_future_returns_promptly_and_retires_owners() { on_fresh_thread(|| { init_device_contexts(0, 1).expect("init failed (requires GPU)"); let dptr = alloc_device(BUF); @@ -308,16 +311,20 @@ fn dropping_in_flight_future_releases_output_after_the_device_finished() { panic!("the first poll failed instead of leaving the op in flight: {error}") } } + let drop_started = Instant::now(); drop(future); - let elapsed = started.elapsed(); + let elapsed = drop_started.elapsed(); assert!( - elapsed >= GATE_DELAY, - "the drop returned after {elapsed:?}, before the gate opened at {GATE_DELAY:?}: \ - it did not wait for the in-flight work" + elapsed < GATE_DELAY / 2, + "dropping async work waited for GPU completion: {elapsed:?}" ); opener.join().expect("gate opener thread panicked"); } + let deadline = Instant::now() + Duration::from_secs(5); + while log.lock().unwrap().len() != 4 && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(5)); + } let log = log.lock().unwrap(); assert_eq!(log.len(), 4, "every dropped future must release its output"); assert!( diff --git a/cuda-async/tests/error_handling.rs b/cuda-async/tests/error_handling.rs index 0513bec421..e1618dcda9 100644 --- a/cuda-async/tests/error_handling.rs +++ b/cuda-async/tests/error_handling.rs @@ -176,7 +176,7 @@ fn new_device_context_with_invalid_device_returns_driver_error() { // Device ordinal 9999 should not exist on any reasonable system. // Device::new(9999) will fail with a driver error, which propagates // through new_device_context. - let result = cuda_core::Device::new(9999); + let result = cuda_core::Context::new(9999); let result = result.map_err(DeviceError::Driver); match result { Err(DeviceError::Driver(_)) => { /* expected */ } diff --git a/cuda-async/tests/execution_lock.rs b/cuda-async/tests/execution_lock.rs index 0a28cc44b4..81f507f1d1 100644 --- a/cuda-async/tests/execution_lock.rs +++ b/cuda-async/tests/execution_lock.rs @@ -48,7 +48,7 @@ fn then_closure_cannot_nest_execution() { fn then_unchecked_closure_can_nest_sync_sync_on_and_await() { on_fresh_thread(|| { init_device_contexts(0, 1).expect("init failed (requires GPU)"); - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let other_stream = device.new_stream().unwrap(); let chain = unsafe { @@ -140,7 +140,7 @@ fn panic_inside_future_poll_releases_lock() { #[test] fn panic_inside_scope_releases_lock() { on_fresh_thread(|| { - let device = cuda_core::Device::new(0).unwrap(); + let device = cuda_core::Context::new(0).unwrap(); let stream = device.new_stream().unwrap(); let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { CudaGraph::scope(&stream, |_s| panic!("intentional panic in scope")) diff --git a/cuda-async/tests/pool_allocation.rs b/cuda-async/tests/pool_allocation.rs index 37b8fdb228..c1bb9ea932 100644 --- a/cuda-async/tests/pool_allocation.rs +++ b/cuda-async/tests/pool_allocation.rs @@ -416,7 +416,7 @@ fn mem_stats_track_alloc_and_free() { let dptr = unsafe { cuda_core::malloc_from_pool_async(N, &pool, &stream) } .expect("pool allocation failed"); assert!(dptr != 0); - unsafe { stream.synchronize() }.expect("stream sync failed"); + stream.synchronize().expect("stream sync failed"); let stats = pool.mem_stats().expect("read stats failed"); assert!( @@ -431,7 +431,7 @@ fn mem_stats_track_alloc_and_free() { ); unsafe { cuda_core::free_async(dptr, &stream) }.expect("free failed"); - unsafe { stream.synchronize() }.expect("stream sync failed"); + stream.synchronize().expect("stream sync failed"); let stats = pool.mem_stats().expect("read stats failed"); assert_eq!( @@ -467,7 +467,7 @@ fn reset_used_high_collapses_watermark_to_current() { let dptr = unsafe { cuda_core::malloc_from_pool_async(N, &pool, &stream) } .expect("pool allocation failed"); unsafe { cuda_core::free_async(dptr, &stream) }.expect("free failed"); - unsafe { stream.synchronize() }.expect("stream sync failed"); + stream.synchronize().expect("stream sync failed"); let before = pool.mem_stats().expect("read stats failed"); assert_eq!(before.used_current, 0); diff --git a/cuda-bindings/Cargo.toml b/cuda-bindings/Cargo.toml index cd7d5f952a..eee40b94ed 100644 --- a/cuda-bindings/Cargo.toml +++ b/cuda-bindings/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" authors = ["Melih Elibol"] edition = "2021" rust-version = "1.89" diff --git a/cuda-core-derive/Cargo.toml b/cuda-core-derive/Cargo.toml index 3ca5a573c0..797531a8b5 100644 --- a/cuda-core-derive/Cargo.toml +++ b/cuda-core-derive/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "cuda-core-derive" description = "Derive macro for cuda-core's DeviceCopy trait." -version = "0.4.0" +version = "0.5.0" edition = "2021" rust-version = "1.89" license = "Apache-2.0" diff --git a/cuda-core-derive/src/device_copy.rs b/cuda-core-derive/src/device_copy.rs index e33ff15a1f..c0477b369f 100644 --- a/cuda-core-derive/src/device_copy.rs +++ b/cuda-core-derive/src/device_copy.rs @@ -20,28 +20,22 @@ pub fn impl_device_copy(input: &DeriveInput, import: TokenStream) -> TokenStream // type-bound on DeviceCopy and call it with every type that's in the struct/union. // This will fail to compile if any of the nested types doesn't implement DeviceCopy. // - // Enums are deliberately rejected: `DeviceCopy`'s safety contract requires - // every bit pattern, including the all-zero pattern written by - // `DeviceBuffer::zeroed`, to be a valid value of the type. That holds for - // product types (structs/unions) when every field is `DeviceCopy`, but NOT - // for enums, whose discriminant leaves most byte patterns invalid. A zeroed - // or device-written buffer could then materialize an out-of-range - // discriminant (undefined behavior). This is the same reason `bool`, `char`, - // and `NonZeroU32` are not `DeviceCopy`. Per-field checking is necessary but - // not sufficient for enums, so we refuse rather than emit an unsound impl. + // Enums need more than field checks: the all-zero representation must also + // select a valid discriminant and payload. This derive does not analyze + // enum layouts, so it conservatively rejects them. DeviceCopy does not + // require every bit pattern to be valid (bool implements it); an enum with + // a valid zero representation can have a manually justified unsafe impl. let check_types_code = match input.data { Data::Struct(ref data_struct) => type_check_struct(data_struct), Data::Union(ref data_union) => type_check_union(data_union), Data::Enum(_) => { return syn::Error::new_spanned( input_type, - "`#[derive(DeviceCopy)]` cannot be applied to enums: `DeviceCopy` requires \ - every bit pattern (including the all-zero pattern written by \ - `DeviceBuffer::zeroed`) to be a valid value, but an enum's discriminant \ - leaves most byte patterns invalid, so a zeroed or device-written buffer \ - could materialize an out-of-range discriminant (undefined behavior). If you \ - can guarantee every device-produced byte pattern is a valid variant, write \ - `unsafe impl DeviceCopy for ... {}` by hand.", + "`#[derive(DeviceCopy)]` cannot be applied to enums: this derive cannot \ + prove that the all-zero representation selects a valid discriminant and \ + payload. `DeviceCopy` requires byte-copyable values and a valid all-zero \ + representation, but not validity of every bit pattern. If the enum meets \ + the full safety contract, write `unsafe impl DeviceCopy for ... {}` by hand.", ) .to_compile_error(); } diff --git a/cuda-core/Cargo.toml b/cuda-core/Cargo.toml index 8eeb103c93..c19d8bc9ee 100644 --- a/cuda-core/Cargo.toml +++ b/cuda-core/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "cuda-core" description = "Idiomatic CUDA API." -version = "0.4.0" +version = "0.5.0" edition = "2021" rust-version = "1.89" license = "Apache-2.0" @@ -13,10 +13,17 @@ readme = "README.md" # Sibling paths, not the root workspace entries: those name the published # release that cuda-oxide consumes, and inheriting them here would make this # copy depend on a crates.io cuda-bindings instead of the one beside it. -cuda-bindings = { version = "=0.4.0", path = "../cuda-bindings" } -cuda-core-derive = { version = "=0.4.0", path = "../cuda-core-derive" } +cuda-bindings = { version = "=0.5.0", path = "../cuda-bindings" } +cuda-core-derive = { version = "=0.5.0", path = "../cuda-core-derive" } +cuda-repr = { version = "=0.5.0", path = "../cuda-repr" } anyhow = "1" -half = { version = "2" } +# `AnyBitPattern` bounds safe reinterpreting casts on `DeviceBuffer`. +# Re-exported at the crate root for downstream element types. +# `min_const_generics` gives `Pod` for arrays of every length. +bytemuck = { version = "1", features = ["min_const_generics"] } +# The `bytemuck` feature makes `f16` and `bf16` `Pod`, so a buffer or tensor +# can be reinterpreted as either of them. +half = { version = "2", features = ["bytemuck"] } oxide-artifacts = { version = "0.2.1", features = ["object-read"] } [dev-dependencies] diff --git a/cuda-core/README.md b/cuda-core/README.md index 2814763b15..cc0547a310 100644 --- a/cuda-core/README.md +++ b/cuda-core/README.md @@ -1,7 +1,43 @@ -# CUDA +# CUDA host runtime `cuda-core` wraps `cuda-bindings` with the CUDA types and functions used by -the rest of the workspace. +the rest of the workspace. Tile and SIMT use the same handles and allocation +owners; there is no `simt` module in this crate. + +`Device` identifies a physical GPU. `Context` owns or borrows a CUDA execution +context on that GPU. `Stream`, `Event`, `Module`, and `Function` retain their +actual context, and compatibility checks compare context handles rather than +device ordinals. The `CudaContext`, `CudaStream`, `CudaEvent`, `CudaModule`, and +`CudaFunction` spellings are aliases of those shared types. + +`DeviceBuffer` provides typed storage. `RawDeviceBuffer` shares allocation +ownership, byte offsets, context identity, and access tracking with typed +buffers and Tile tensors. Owned synchronous allocations, stream-ordered +allocations, borrowed pointers, and foreign owners use this same representation. +Foreign `DeviceAllocation` implementations supply their actual context. + +Safe submissions retain allocation access leases independently of their future. +Conflicting accesses on different streams return `CUDA_ERROR_NOT_READY` until +completion; accesses in one stream remain ordered by CUDA. Buffer and stream +destructors retire resources without waiting for GPU work on the dropping +thread. If completion cannot be established, retirement retains the resource +instead of freeing or reusing it. Unsafe raw-pointer users must establish their +own lifetime, initialization, and ordering guarantees. + +`DeviceCopy` permits byte transfers of valid values. `DType` describes Tile +element types. Neither trait implies the other; types supported by both models +implement both. Safe casts to another representation additionally require +`AnyBitPattern`, while `DeviceRepr` supports ownership-preserving wrappers such +as device atomics. A `DeviceBuffer` remains supported, but arbitrary bytes +cannot safely be cast to `bool`. + +For migration, replace `cuda_core::simt::...` imports with `cuda_core::...`. +Code using the old context-owning `Device` should use `Context`; stream and pool +context accessors are named `context()`. With cuTile, use `Tensor::from_buffer` +and `Tensor::into_buffer` for ownership transfer without copying device memory, +or `Tensor::as_buffer` / `TensorView::as_buffer` for contiguous borrowed kernel +arguments. Ownership transfer requires unique storage; borrowed slices preserve +their byte offset and reject strided layouts. # Testing diff --git a/cuda-core/src/allocation.rs b/cuda-core/src/allocation.rs new file mode 100644 index 0000000000..11358744fb --- /dev/null +++ b/cuda-core/src/allocation.rs @@ -0,0 +1,436 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared allocation ownership, byte views, and in-flight access tracking. + +use crate::{Context, DriverError, Stream}; +use cuda_bindings::CUdeviceptr; +use std::sync::{Arc, Mutex}; + +/// A foreign allocation whose owner keeps both its memory and CUDA context alive. +/// +/// # Safety +/// The pointer, byte extent, and context must remain stable and valid for the +/// owner's lifetime. The pointer must address at least `len_bytes` bytes in that +/// context. Per-operation exclusion from foreign accesses remains the obligation +/// of callers constructing typed views over this memory. +pub unsafe trait DeviceAllocation: Send + Sync + 'static { + fn device_ptr(&self) -> CUdeviceptr; + fn len_bytes(&self) -> usize; + fn context(&self) -> Arc; + fn device_id(&self) -> usize { + self.context().ordinal() + } +} + +struct Access { + id: u64, + stream: usize, + write: bool, + recording: bool, + completion: Option>, +} + +#[derive(Default)] +struct Accesses { + next_id: u64, + active: Vec, +} + +impl Accesses { + fn register( + &mut self, + stream: usize, + write: bool, + recording: bool, + ) -> Result { + // A synchronous wait on another thread may have completed a retired + // operation before the worker got to its lease. Do not report a stale + // conflict in that case. Errors leave the exclusion intact. + self.active + .retain(|a| !matches!(a.completion.as_ref().map(|e| e.query()), Some(Ok(true)))); + if !recording + && self + .active + .iter() + .any(|a| !a.recording && a.stream != stream && (write || a.write)) + { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_READY, + )); + } + let id = self.next_id.checked_add(1).ok_or(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + ))?; + self.next_id = id; + self.active.push(Access { + id, + stream, + write, + recording, + completion: None, + }); + Ok(id) + } +} + +enum Owner { + Synchronous, + StreamOrdered(Arc), + Borrowed, + Foreign(#[allow(dead_code)] Arc), +} + +struct Allocation { + ptr: CUdeviceptr, + bytes: usize, + context: Arc, + owner: Owner, + accesses: Mutex, +} + +impl Drop for Allocation { + fn drop(&mut self) { + if self.ptr == 0 || matches!(self.owner, Owner::Borrowed | Owner::Foreign(_)) { + return; + } + let ptr = self.ptr; + let context = self.context.clone(); + let allocation_stream = match &self.owner { + Owner::StreamOrdered(s) => Some(s.clone()), + _ => None, + }; + // All safe users hold an independent access lease until completion. + // Driver teardown is retired on a worker, never on the dropping thread. + crate::completion::defer(move || { + if let Err(error) = context.bind_to_thread() { + context.record_err(Err::<(), _>(error)); + std::mem::forget((context, allocation_stream)); + return; + } + let result = if let Some(stream) = &allocation_stream { + // SAFETY: the allocation is no longer accessed and the retained + // stream is the one on which its allocation became valid. + unsafe { crate::memory::free_async(ptr, stream.cu_stream()) } + } else { + // SAFETY: ownership of this synchronous allocation is unique. + unsafe { crate::memory::free_sync(ptr) } + }; + if let Err(error) = result { + context.record_err(Err::<(), _>(error)); + std::mem::forget((context, allocation_stream)); + } + }); + } +} + +/// A byte view sharing an allocation owner and access tracker with every derived +/// typed buffer, tensor, and sliced view. Cloning does not copy device memory. +#[derive(Clone)] +pub struct RawDeviceBuffer { + allocation: Arc, + // Only user-visible views share this token. Completion leases retain the + // allocation directly, so creating/dropping them cannot affect uniqueness. + views: Arc<()>, + offset: usize, + len: usize, +} + +impl std::fmt::Debug for RawDeviceBuffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RawDeviceBuffer") + .field("ptr", &self.cu_deviceptr()) + .field("len", &self.len) + .field("context", &self.allocation.context) + .finish_non_exhaustive() + } +} + +impl RawDeviceBuffer { + fn new(ptr: CUdeviceptr, bytes: usize, context: Arc, owner: Owner) -> Self { + assert!( + ptr.checked_add(bytes as u64).is_some(), + "device allocation extent overflows" + ); + Self { + allocation: Arc::new(Allocation { + ptr, + bytes, + context, + owner, + accesses: Mutex::default(), + }), + views: Arc::new(()), + offset: 0, + len: bytes, + } + } + + /// Adopts a synchronous CUDA allocation. + /// + /// # Safety + /// `ptr` must uniquely own `bytes` bytes allocated by `cuMemAlloc` in + /// `context`, with no outstanding untracked accesses. + pub unsafe fn from_raw_parts(ptr: CUdeviceptr, bytes: usize, context: Arc) -> Self { + Self::new(ptr, bytes, context, Owner::Synchronous) + } + + /// Adopts a stream-ordered CUDA allocation. + /// + /// # Safety + /// `ptr` must uniquely own `bytes` bytes from `cuMemAllocAsync` or its pool + /// variant on `stream`. The caller must establish allocation ordering for + /// every use and retain the allocation through outstanding work. + pub unsafe fn from_raw_parts_async( + ptr: CUdeviceptr, + bytes: usize, + stream: Arc, + ) -> Self { + Self::new( + ptr, + bytes, + stream.context().clone(), + Owner::StreamOrdered(stream), + ) + } + + /// Borrows externally managed memory without taking ownership. + /// + /// # Safety + /// Memory and context must outlive every derived view and queued use. No + /// external access may conflict with accesses through this allocation. + pub unsafe fn borrowed_from_raw_parts( + ptr: CUdeviceptr, + bytes: usize, + context: Arc, + ) -> Self { + Self::new(ptr, bytes, context, Owner::Borrowed) + } + + /// Retains an external owner and limits the logical view to `bytes` bytes. + /// No CUDA calls are made, permitting metadata validation before submission. + pub fn foreign(owner: Arc, bytes: usize) -> Self { + assert!( + bytes <= owner.len_bytes(), + "view exceeds foreign allocation" + ); + Self::new( + owner.device_ptr(), + bytes, + owner.context(), + Owner::Foreign(owner), + ) + } + + pub fn cu_deviceptr(&self) -> CUdeviceptr { + self.allocation.ptr + self.offset as u64 + } + pub fn len_bytes(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } + pub fn context(&self) -> &Arc { + &self.allocation.context + } + pub(crate) fn allocation_stream(&self) -> Option<&Arc> { + match &self.allocation.owner { + Owner::StreamOrdered(stream) => Some(stream), + _ => None, + } + } + pub fn device_id(&self) -> usize { + self.context().ordinal() + } + /// Identity of the owner/access tracker, unchanged across byte views. + pub fn allocation_id(&self) -> usize { + Arc::as_ptr(&self.allocation) as usize + } + + /// Creates a bounds-checked view without changing allocation identity. + pub fn slice_bytes(&self, offset: usize, bytes: usize) -> Result { + if offset.checked_add(bytes).is_none_or(|end| end > self.len) { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + Ok(Self { + allocation: self.allocation.clone(), + views: self.views.clone(), + offset: self.offset + offset, + len: bytes, + }) + } + + /// Whether only this user-visible view owns the allocation. Completion leases + /// do not count as user aliases; they still enforce in-flight exclusion. + pub fn is_unique(&self) -> bool { + Arc::strong_count(&self.views) == 1 + } + + /// Acquires an allocation-wide access, rejecting cross-stream conflicts. + /// The returned lease must be retained until the submitted work completes. + pub fn acquire( + &self, + stream: &Stream, + write: bool, + recording: bool, + ) -> Result { + if self.context().as_ref() != stream.context().as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + self.acquire_on_stream(stream.cu_stream() as usize, write, recording) + } + + /// Registers bookkeeping for a known stream handle. This does not submit GPU + /// work or grant permission to dereference pointers. Runtime adapters must + /// separately validate the stream's context before submitting work. + #[doc(hidden)] + pub fn acquire_on_stream( + &self, + stream: usize, + write: bool, + recording: bool, + ) -> Result { + let mut accesses = self + .allocation + .accesses + .lock() + .unwrap_or_else(|e| e.into_inner()); + let id = accesses.register(stream, write, recording)?; + Ok(AllocationLease { + allocation: self.allocation.clone(), + id, + }) + } + + pub(crate) fn into_owned_parts( + self, + ) -> (CUdeviceptr, usize, Arc, Option>) { + assert_eq!(self.offset, 0, "cannot detach a sliced allocation"); + assert_eq!( + self.len, self.allocation.bytes, + "cannot detach a partial allocation" + ); + let mut allocation = Arc::try_unwrap(self.allocation) + .unwrap_or_else(|_| panic!("cannot detach a shared or in-flight allocation")); + let stream = match &allocation.owner { + Owner::Synchronous => None, + Owner::StreamOrdered(stream) => Some(stream.clone()), + _ => panic!("cannot detach an allocation owned externally"), + }; + let ptr = std::mem::replace(&mut allocation.ptr, 0); + (ptr, allocation.bytes, allocation.context.clone(), stream) + } +} + +/// Keeps an allocation alive and excludes conflicting accesses until dropped. +/// A runtime must hold this beyond enqueue, until completion has been proven. +pub struct AllocationLease { + allocation: Arc, + id: u64, +} + +impl Drop for AllocationLease { + fn drop(&mut self) { + let mut accesses = self + .allocation + .accesses + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(pos) = accesses.active.iter().position(|a| a.id == self.id) { + accesses.active.swap_remove(pos); + } + } +} + +impl AllocationLease { + /// Transfers this lease to completion-bound retirement after submission. + pub fn retain_until_complete(self, stream: &Stream) -> Result<(), DriverError> { + let event = match stream.record_event(None) { + Ok(event) => Arc::new(event), + Err(error) => { + std::mem::forget(self); + return Err(error); + } + }; + { + let mut accesses = self + .allocation + .accesses + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(access) = accesses.active.iter_mut().find(|a| a.id == self.id) { + access.completion = Some(event.clone()); + } + } + crate::completion::retain_until_event(self, event); + Ok(()) + } +} + +impl std::fmt::Debug for AllocationLease { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AllocationLease") + .field("id", &self.id) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + + fn metadata() -> RawDeviceBuffer { + // SAFETY: bookkeeping only. These borrowed handles are never passed to + // CUDA, dereferenced, or freed by their destructors. + unsafe { + RawDeviceBuffer::borrowed_from_raw_parts( + 0x1000, + 128, + Context::borrow_raw(0x2000 as *mut _, 0, 0), + ) + } + } + + #[test] + fn uniqueness_counts_byte_views_but_not_completion_leases() { + let raw = metadata(); + let view = raw.slice_bytes(16, 32).unwrap(); + let clone = raw.clone(); + let lease = raw.acquire_on_stream(1, true, false).unwrap(); + assert!(!raw.is_unique()); + drop(view); + assert!(!raw.is_unique()); + drop(clone); + assert!(raw.is_unique()); + assert!(raw.acquire_on_stream(2, false, false).is_err()); + drop(lease); + assert!(raw.acquire_on_stream(2, false, false).is_ok()); + } + + #[test] + fn concurrent_lease_retirement_cannot_change_user_uniqueness() { + let raw = metadata(); + let stop = AtomicBool::new(false); + let mut false_negatives = 0; + std::thread::scope(|scope| { + scope.spawn(|| { + while !stop.load(Ordering::Relaxed) { + drop(raw.acquire_on_stream(1, false, false).unwrap()); + } + }); + for _ in 0..100_000 { + false_negatives += usize::from(!raw.is_unique()); + } + stop.store(true, Ordering::Relaxed); + }); + assert_eq!(false_negatives, 0); + } +} diff --git a/cuda-core/src/api.rs b/cuda-core/src/api.rs index 0b6ea8a9bf..4802bf5dca 100644 --- a/cuda-core/src/api.rs +++ b/cuda-core/src/api.rs @@ -171,6 +171,7 @@ pub unsafe fn malloc_async( num_bytes: usize, stream: &Arc, ) -> Result { + stream.context().bind_to_thread()?; crate::cudarc_shim::memory::malloc_async(stream.cu_stream(), num_bytes) } @@ -185,6 +186,7 @@ pub unsafe fn malloc_from_pool_async( pool: &Arc, stream: &Arc, ) -> Result { + stream.context().bind_to_thread()?; crate::cudarc_shim::pool::malloc_from_pool_async(pool.cu_pool(), stream.cu_stream(), num_bytes) } @@ -199,6 +201,7 @@ pub unsafe fn memset_d8_async( num_bytes: usize, stream: &Arc, ) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; crate::cudarc_shim::memory::memset_d8_async(dptr, value, num_bytes, stream.cu_stream()) } @@ -209,6 +212,7 @@ pub unsafe fn memset_d8_async( /// [`malloc_from_pool_async`], `stream` must be ordered after every use of /// the allocation, and `dptr` must not be used after this call. pub unsafe fn free_async(dptr: sys::CUdeviceptr, stream: &Arc) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; crate::cudarc_shim::memory::free_async(dptr, stream.cu_stream()) } @@ -222,6 +226,7 @@ pub unsafe fn memcpy_htod_async( num_elements: usize, stream: &Arc, ) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; let num_bytes = num_elements * mem::size_of::(); unsafe { crate::cudarc_shim::memory::memcpy_htod_async(dst, src, num_bytes, stream.cu_stream()) @@ -238,6 +243,7 @@ pub unsafe fn memcpy_dtoh_async( num_elements: usize, stream: &Arc, ) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; let num_bytes = num_elements * mem::size_of::(); unsafe { crate::cudarc_shim::memory::memcpy_dtoh_async(dst, src, num_bytes, stream.cu_stream()) @@ -254,6 +260,7 @@ pub unsafe fn memcpy_dtod_async( num_elements: usize, stream: &Arc, ) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; let num_bytes = num_elements * mem::size_of::(); unsafe { crate::cudarc_shim::memory::memcpy_dtod_async(dst, src, num_bytes, stream.cu_stream()) diff --git a/cuda-core/src/completion.rs b/cuda-core/src/completion.rs new file mode 100644 index 0000000000..c6acd9b982 --- /dev/null +++ b/cuda-core/src/completion.rs @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Completion-bound retirement without waiting in destructors. + +use crate::{DriverError, Event, Stream}; +use std::sync::{mpsc, Arc, OnceLock}; +use std::time::Duration; + +type Cleanup = Box; +struct Pending { + event: Option>, + cleanup: Cleanup, +} + +fn enqueue(pending: Pending) { + static QUEUE: OnceLock>> = OnceLock::new(); + let queue = QUEUE.get_or_init(|| { + let (sender, receiver) = mpsc::channel::(); + std::thread::Builder::new() + .name("cuda-resource-retirement".into()) + .spawn(move || { + let mut pending = Vec::::new(); + loop { + match receiver.recv_timeout(Duration::from_millis(1)) { + Ok(item) => pending.push(item), + Err(mpsc::RecvTimeoutError::Disconnected) => break, + Err(mpsc::RecvTimeoutError::Timeout) => {} + } + pending.extend(receiver.try_iter()); + let mut i = 0; + while i < pending.len() { + let ready = pending[i] + .event + .as_ref() + .map_or(Ok(true), |event| event.query()); + match ready { + Ok(false) => i += 1, + Ok(true) => { + let item = pending.swap_remove(i); + // Preserve the event's context until cleanup has finished. + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe( + item.cleanup, + )); + } + Err(_) => { + // Completion is unknown: leaking the owner and its access + // lease is safer than freeing memory or allowing conflicting work. + std::mem::forget(pending.swap_remove(i)); + } + } + } + } + for item in pending { + std::mem::forget(item); + } + }) + .ok() + .map(|_| sender) + }); + if let Some(queue) = queue { + if let Err(error) = queue.send(pending) { + std::mem::forget(error.0); + } + } else { + std::mem::forget(pending); + } +} + +pub(crate) fn defer(cleanup: impl FnOnce() + Send + 'static) { + enqueue(Pending { + event: None, + cleanup: Box::new(cleanup), + }); +} + +/// Keeps an owner or access lease alive until work already queued on `stream` +/// completes. Dropping the caller's future cannot shorten this lifetime. +/// +/// Failure to record or observe completion intentionally leaks the owner; +/// releasing it could invalidate work that the driver has already accepted. +pub fn retain_until_complete( + owner: impl Send + 'static, + stream: &Stream, +) -> Result<(), DriverError> { + let event = match stream.record_event(None) { + Ok(event) => event, + Err(error) => { + std::mem::forget(owner); + return Err(error); + } + }; + enqueue(Pending { + event: Some(Arc::new(event)), + cleanup: Box::new(move || drop(owner)), + }); + Ok(()) +} + +pub(crate) fn retain_until_event(owner: impl Send + 'static, event: Arc) { + enqueue(Pending { + event: Some(event), + cleanup: Box::new(move || drop(owner)), + }); +} diff --git a/cuda-core/src/simt/context.rs b/cuda-core/src/context.rs similarity index 68% rename from cuda-core/src/simt/context.rs rename to cuda-core/src/context.rs index 2a76c5a389..5d8270254d 100644 --- a/cuda-core/src/simt/context.rs +++ b/cuda-core/src/context.rs @@ -3,93 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! CUDA context management (primary context, RAII). -//! -//! [`CudaContext`] retains the **primary context** for a given device ordinal -//! via `cuDevicePrimaryCtxRetain` and releases it on [`Drop`]. The primary -//! context is shared across the process; multiple `CudaContext` instances for -//! the same device share the same underlying `CUcontext`. -//! -//! # Thread binding -//! -//! CUDA driver calls are context-scoped and thread-local. [`CudaContext`] -//! transparently calls `cuCtxSetCurrent` before any driver operation, so -//! callers do not need to manage the context stack manually. +//! Context queries and configuration shared by both programming models. -use crate::error::{DriverError, IntoResult}; -use crate::simt::launch::DeviceLaunchLimits; -use crate::simt::stream::CudaStream; -use std::ffi::c_int; -use std::mem::MaybeUninit; -use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; -use std::sync::Arc; - -/// Owns a reference to a CUDA device's primary context. -/// -/// Created via [`CudaContext::new`] and typically held in an `Arc` so streams, -/// events, and modules can share the same context. Dropping the last reference -/// releases the primary context (`cuDevicePrimaryCtxRelease`). -/// -/// Tracks live stream count and accumulated error state atomically for -/// cross-thread diagnostics. -#[derive(Debug)] -pub struct CudaContext { - /// Raw CUDA device handle (`CUdevice`). - pub(crate) cu_device: cuda_bindings::CUdevice, - /// Raw CUDA context handle (`CUcontext`). Set to null on drop. - pub(crate) cu_ctx: cuda_bindings::CUcontext, - /// Zero-based device ordinal passed to [`CudaContext::new`]. - pub(crate) ordinal: usize, - /// Number of live [`CudaStream`] instances sharing this context. - pub(crate) num_streams: AtomicUsize, - /// When `true`, the first [`new_stream`](CudaContext::new_stream) call - /// synchronizes the context to establish a clean ordering baseline. - pub(crate) event_tracking: AtomicBool, - /// Sticky error state recorded by [`record_err`](CudaContext::record_err). - /// Stores the raw `CUresult` value, or `0` if no error. - pub(crate) error_state: AtomicU32, -} - -/// # Safety -/// -/// `CUdevice` and `CUcontext` are process-wide handles. All mutable state -/// (`num_streams`, `event_tracking`, `error_state`) uses atomics. The CUDA -/// driver itself is thread-safe for distinct contexts, and the -/// [`bind_to_thread`](CudaContext::bind_to_thread) mechanism ensures the -/// correct context is current before each call. -unsafe impl Send for CudaContext {} -/// See [`Send`] impl. -unsafe impl Sync for CudaContext {} - -/// Releases the primary context on drop. -/// -/// Binds the context to the current thread first (required by -/// `cuDevicePrimaryCtxRelease`). Errors during teardown are recorded via -/// [`record_err`](CudaContext::record_err) rather than panicking. -impl Drop for CudaContext { - fn drop(&mut self) { - self.record_err(self.bind_to_thread()); - let ctx = std::mem::replace(&mut self.cu_ctx, std::ptr::null_mut()); - if !ctx.is_null() { - self.record_err(unsafe { - cuda_bindings::cuDevicePrimaryCtxRelease_v2(self.cu_device).result() - }); - } - } -} - -/// Equality is based on device handle, context handle, and ordinal. -impl PartialEq for CudaContext { - fn eq(&self, other: &Self) -> bool { - self.cu_device == other.cu_device - && self.cu_ctx == other.cu_ctx - && self.ordinal == other.ordinal - } -} -impl Eq for CudaContext {} +pub use crate::runtime::{Context, Context as CudaContext}; +use crate::{DeviceLaunchLimits, DriverError, IntoResult, KeepAlive, Stream}; +use std::{ + ffi::c_int, + mem::MaybeUninit, + sync::{atomic::Ordering, Arc}, +}; /// The device's meaningful stream priorities, from -/// [`CudaContext::stream_priority_range`]. +/// [`Context::stream_priority_range`]. /// /// CUDA orders priorities the opposite way round to intuition: **a lower /// number is a higher priority**, so the range runs from @@ -142,17 +67,17 @@ impl StreamPriorityRange { /// The value the driver will actually apply for a request of `priority`. /// /// Answers the silent clamp in - /// [`CudaContext::new_stream_with_priority`] ahead of the call. + /// [`Context::new_stream_with_priority`] ahead of the call. pub fn clamp(&self, priority: i32) -> i32 { priority.clamp(self.greatest, self.least) } } /// A per-context device limit (`CU_LIMIT_*`), read with -/// [`CudaContext::limit`] and written with [`CudaContext::set_limit`]. +/// [`Context::limit`] and written with [`Context::set_limit`]. /// /// Every limit is state on the device's **primary** context, so it is shared -/// by every [`CudaContext`] for that device, by any runtime-API user of the +/// by every [`Context`] for that device, by any runtime-API user of the /// same device in this process, and across library boundaries. A limit set /// here is not scoped to the handle that set it. /// @@ -254,8 +179,8 @@ impl ContextLimit { /// waits when the calling host thread blocks on GPU work (a stream or /// context synchronize). /// -/// Set via [`CudaContext::set_sync_policy`], read via -/// [`CudaContext::sync_policy`]. The trade-off is host CPU against wake +/// Set via [`Context::set_sync_policy`], read via +/// [`Context::sync_policy`]. The trade-off is host CPU against wake /// latency: [`Spin`](Self::Spin) burns a full core for the lowest latency, /// [`BlockingSync`](Self::BlockingSync) frees the core but wakes later. On a /// host with spare cores this trade rarely matters; on a CPU-constrained one @@ -302,7 +227,7 @@ impl SyncPolicy { } } -impl CudaContext { +impl Context { fn device_attribute( &self, attribute: cuda_bindings::CUdevice_attribute, @@ -316,173 +241,6 @@ impl CudaContext { .map_err(|_| DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE)) } } - - /// Creates a new context for the device at `ordinal`. - /// - /// Calls [`cuInit`](crate::init), obtains the device handle, retains the - /// primary context, and binds it to the calling thread. Returns the context - /// wrapped in an `Arc` for shared ownership across streams, events, and - /// modules. - pub fn new(ordinal: usize) -> Result, DriverError> { - unsafe { crate::init(0)? }; - - let cu_device = unsafe { - let mut cu_device = MaybeUninit::uninit(); - cuda_bindings::cuDeviceGet(cu_device.as_mut_ptr(), ordinal as c_int).result()?; - cu_device.assume_init() - }; - - let cu_ctx = unsafe { - let mut cu_ctx = MaybeUninit::uninit(); - cuda_bindings::cuDevicePrimaryCtxRetain(cu_ctx.as_mut_ptr(), cu_device).result()?; - cu_ctx.assume_init() - }; - - let ctx = Arc::new(CudaContext { - cu_device, - cu_ctx, - ordinal, - num_streams: AtomicUsize::new(0), - event_tracking: AtomicBool::new(true), - error_state: AtomicU32::new(0), - }); - ctx.bind_to_thread()?; - Ok(ctx) - } - - /// Returns the zero-based device ordinal. - pub fn ordinal(&self) -> usize { - self.ordinal - } - - /// Returns the raw `CUdevice` handle. - pub fn cu_device(&self) -> cuda_bindings::CUdevice { - self.cu_device - } - - /// Returns the raw `CUcontext` handle. - pub fn cu_ctx(&self) -> cuda_bindings::CUcontext { - self.cu_ctx - } - - /// Binds this context to the calling thread if not already current. - /// - /// Checks [`check_err`](Self::check_err) first and propagates any sticky - /// error. Skips the `cuCtxSetCurrent` call when the context is already - /// bound, avoiding an unnecessary driver round-trip. - /// - /// Most methods on [`CudaStream`], [`CudaEvent`](crate::CudaEvent), and - /// [`CudaModule`](crate::CudaModule) call this internally. - /// - /// CUcontext is the backing runtime object, and CUmodule / CUfunction / CUstream - /// are opaque handles to objects created under that context. bind_to_thread() - /// makes that context, the one, the current host thread is operating against. - /// If the thread is currently bound to some other context, using those handles - /// can fail. - pub fn bind_to_thread(&self) -> Result<(), DriverError> { - self.check_err()?; - let mut current = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuCtxGetCurrent(current.as_mut_ptr()).result()?; - let current = current.assume_init(); - if current.is_null() || current != self.cu_ctx { - cuda_bindings::cuCtxSetCurrent(self.cu_ctx).result()?; - } - } - Ok(()) - } - - /// Blocks the calling thread until all preceding work in this context - /// completes. - /// - /// Binds the context first, then calls `cuCtxSynchronize`. - pub fn synchronize(&self) -> Result<(), DriverError> { - self.bind_to_thread()?; - unsafe { cuda_bindings::cuCtxSynchronize() }.result() - } - - /// Returns a handle to the per-context default stream (stream `0`). - /// - /// The default stream implicitly synchronizes with all blocking streams in - /// the same context. The returned [`CudaStream`] holds a null `CUstream` - /// pointer, which the driver interprets as the default stream. - pub fn default_stream(self: &Arc) -> Arc { - Arc::new(CudaStream { - cu_stream: std::ptr::null_mut(), - ctx: self.clone(), - }) - } - - /// Creates a new non-blocking stream in this context. - /// - /// The stream is created with `CU_STREAM_NON_BLOCKING`, so it does not - /// implicitly synchronize with the default stream. - /// - /// On the first call (when `num_streams` transitions from 0 to 1), the - /// context is synchronized to establish a clean ordering baseline if - /// `event_tracking` is enabled. - pub fn new_stream(self: &Arc) -> Result, DriverError> { - self.create_stream(None) - } - - /// Creates a new non-blocking stream at `priority` in this context. - /// - /// As [`new_stream`](Self::new_stream), with the stream created through - /// `cuStreamCreateWithPriority`. - /// - /// **Lower numbers are higher priorities**, and a value outside the - /// device's range is **clamped silently** to the nearest end of it rather - /// than refused, so passing [`i32::MIN`] yields the greatest priority the - /// device supports and reports nothing. Read - /// [`CudaStream::priority`] back to see what the driver applied, or - /// consult [`stream_priority_range`](Self::stream_priority_range) first. - /// - /// Priority orders **compute kernels only**. Host-to-device and - /// device-to-host transfers are unaffected, so a high-priority stream does - /// not jump a queue of copies. On a device without priority support every - /// priority collapses to the single value 0 and the stream still works. - pub fn new_stream_with_priority( - self: &Arc, - priority: i32, - ) -> Result, DriverError> { - self.create_stream(Some(priority)) - } - - /// Shared body of [`new_stream`](Self::new_stream) and - /// [`new_stream_with_priority`](Self::new_stream_with_priority). - /// - /// `None` takes `cuStreamCreate`, which is what the driver documents as - /// equivalent to a priority of 0, rather than passing 0 explicitly: on a - /// device whose range does not contain 0 those are different requests. - fn create_stream( - self: &Arc, - priority: Option, - ) -> Result, DriverError> { - self.bind_to_thread()?; - let prev = self.num_streams.fetch_add(1, Ordering::Relaxed); - if prev == 0 && self.event_tracking.load(Ordering::Relaxed) { - self.synchronize()?; - } - let flags = cuda_bindings::CUstream_flags_enum_CU_STREAM_NON_BLOCKING; - let mut cu_stream = MaybeUninit::uninit(); - let cu_stream = unsafe { - match priority { - Some(priority) => cuda_bindings::cuStreamCreateWithPriority( - cu_stream.as_mut_ptr(), - flags, - priority, - ), - None => cuda_bindings::cuStreamCreate(cu_stream.as_mut_ptr(), flags), - } - .result()?; - cu_stream.assume_init() - }; - Ok(Arc::new(CudaStream { - cu_stream, - ctx: self.clone(), - })) - } - /// Returns the device's meaningful stream priority range. /// /// Wraps `cuCtxGetStreamPriorityRange`. See [`StreamPriorityRange`] for @@ -501,7 +259,6 @@ impl CudaContext { }) } } - /// Queries the device's marketing name (e.g. `"NVIDIA GeForce RTX 5090"`). /// /// Wraps `cuDeviceGetName` with a 256-byte buffer (driver guarantees @@ -523,7 +280,6 @@ impl CudaContext { let end = buf.iter().position(|&c| c == 0).unwrap_or(buf.len()); Ok(String::from_utf8_lossy(&buf[..end]).into_owned()) } - /// Queries the compute capability (SM version) of the device. /// /// Returns `(major, minor)` -- e.g. `(9, 0)` for Hopper H100. @@ -547,7 +303,6 @@ impl CudaContext { Ok((major.assume_init(), minor.assume_init())) } } - /// Queries dimension, thread-count, and portable shared-memory launch /// limits for this device. /// @@ -586,7 +341,6 @@ impl CudaContext { )?, }) } - /// Queries the non-portable opt-in shared-memory limit per block. /// /// Typed launch preparation calls this only when static plus dynamic @@ -596,7 +350,6 @@ impl CudaContext { cuda_bindings::CUdevice_attribute_enum_CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, ) } - /// Returns whether this device supports cooperative kernel launches. pub fn supports_cooperative_launch(&self) -> Result { self.device_attribute( @@ -604,7 +357,6 @@ impl CudaContext { ) .map(|value| value != 0) } - /// Returns whether this device supports thread-block cluster launches. pub fn supports_cluster_launch(&self) -> Result { self.device_attribute( @@ -612,14 +364,12 @@ impl CudaContext { ) .map(|value| value != 0) } - /// Returns the number of streaming multiprocessors on this device. pub fn multiprocessor_count(&self) -> Result { self.device_attribute( cuda_bindings::CUdevice_attribute_enum_CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, ) } - /// Returns the maximum number of threads resident on one streaming /// multiprocessor. /// @@ -631,7 +381,6 @@ impl CudaContext { cuda_bindings::CUdevice_attribute_enum_CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR, ) } - /// Reads a device limit off this context. /// /// Wraps `cuCtxGetLimit`. See [`ContextLimit`] for what each limit means, @@ -645,7 +394,6 @@ impl CudaContext { Ok(value.assume_init()) } } - /// Writes a device limit on this context. /// /// Wraps `cuCtxSetLimit`. The limit is state on the device's primary @@ -662,7 +410,6 @@ impl CudaContext { self.bind_to_thread()?; unsafe { cuda_bindings::cuCtxSetLimit(limit.to_raw(), value) }.result() } - /// Queries the per-thread device stack size, in bytes. /// /// Equivalent to [`limit`](Self::limit) with @@ -671,7 +418,6 @@ impl CudaContext { pub fn stack_size(&self) -> Result { self.limit(ContextLimit::StackSize) } - /// Sets the per-thread device stack size, in bytes. /// /// Equivalent to [`set_limit`](Self::set_limit) with @@ -688,19 +434,18 @@ impl CudaContext { pub fn set_stack_size(&self, bytes: usize) -> Result<(), DriverError> { self.set_limit(ContextLimit::StackSize, bytes) } - /// Sets the context's scheduling policy (`CU_CTX_SCHED_*`). /// /// Wraps `cuDevicePrimaryCtxSetFlags_v2`, scoped to this context's /// device rather than `cuCtxSetFlags`'s "whatever is current on this - /// thread": [`CudaContext`] retains and owns a specific device's primary + /// thread": [`Context`] retains and owns a specific device's primary /// context, so the device-scoped call is the one that matches what this /// type actually holds. Read-modify-write: only the 3 scheduling bits /// (`CU_CTX_SCHED_MASK`) are replaced, any other flag bit the process has /// set independently (e.g. `CU_CTX_MAP_HOST`) is preserved. /// /// The policy is process-wide state on the device's primary context: - /// it affects every [`CudaContext`] clone of this device, and any + /// it affects every [`Context`] clone of this device, and any /// runtime-API user of the same device in this process, not just this /// handle. The read-modify-write over `CU_CTX_SCHED_MASK` is also not /// atomic: a concurrent flag writer on the same device can interleave @@ -727,7 +472,6 @@ impl CudaContext { cuda_bindings::cuDevicePrimaryCtxSetFlags_v2(self.cu_device, new_flags).result() } } - /// Returns the context's current scheduling policy (`CU_CTX_SCHED_*`). /// /// `None` if the driver reports a scheduling value this enum does not @@ -747,7 +491,6 @@ impl CudaContext { Ok(SyncPolicy::from_raw(flags.assume_init())) } } - /// Atomically reads and clears the sticky error state. /// /// Returns `Ok(())` if no error was recorded, or the stored @@ -760,7 +503,6 @@ impl CudaContext { Err(DriverError(error_state)) } } - /// Records a driver error into the sticky error state. /// /// Used during [`Drop`] paths where returning a `Result` is not possible. @@ -772,4 +514,35 @@ impl CudaContext { self.error_state.store(err.0, Ordering::Relaxed) } } + /// Returns the context's legacy default stream (stream zero). + pub fn default_stream(self: &Arc) -> Arc { + Arc::new(Stream { + cu_stream: std::ptr::null_mut(), + context: self.clone(), + owned: false, + _keep_alive: KeepAlive::none(), + }) + } + /// Creates a nonblocking stream with the requested CUDA scheduling priority. + pub fn new_stream_with_priority( + self: &Arc, + priority: i32, + ) -> Result, DriverError> { + self.bind_to_thread()?; + let mut stream = std::ptr::null_mut(); + unsafe { + cuda_bindings::cuStreamCreateWithPriority( + &mut stream, + cuda_bindings::CUstream_flags_enum_CU_STREAM_NON_BLOCKING, + priority, + ) + .result()?; + } + Ok(Arc::new(Stream { + cu_stream: stream, + context: self.clone(), + owned: true, + _keep_alive: KeepAlive::none(), + })) + } } diff --git a/cuda-core/src/simt/device_buffer.rs b/cuda-core/src/device_buffer.rs similarity index 55% rename from cuda-core/src/simt/device_buffer.rs rename to cuda-core/src/device_buffer.rs index 9057de9be9..3cea3c9d29 100644 --- a/cuda-core/src/simt/device_buffer.rs +++ b/cuda-core/src/device_buffer.rs @@ -12,12 +12,9 @@ //! [`DeviceBuffer::uninitialized_async`] retain their allocation stream for //! deallocation. //! -//! Ordinary drop synchronizes the context before freeing an asynchronous -//! allocation, because safe operations may have submitted work using the -//! buffer on any stream in that context. The unsafe -//! [`DeviceBuffer::drop_async`] avoids that host-side synchronization by -//! freeing on a chosen stream; its caller takes over the obligation to -//! order every other stream that uses the buffer before that stream. +//! Drop does not wait for GPU work. Safe launch and transfer APIs retain the +//! shared allocation owner and access leases independently until completion. +//! Unsafe raw-pointer callers must retain storage and order accesses themselves. //! //! # Quick start //! @@ -33,12 +30,14 @@ use std::mem::MaybeUninit; use std::num::Wrapping; use std::sync::Arc; +use bytemuck::AnyBitPattern; use cuda_bindings::CUdeviceptr; +use cuda_repr::DeviceRepr; +use crate::context::Context; use crate::error::DriverError; -use crate::simt::context::CudaContext; -use crate::simt::pinned_host_buffer::PinnedHostBuffer; -use crate::simt::stream::CudaStream; +use crate::pinned_host_buffer::PinnedHostBuffer; +use crate::stream::Stream; /// Marker trait for values that can be safely copied between host and device /// memory as raw bytes. @@ -50,15 +49,21 @@ use crate::simt::stream::CudaStream; /// /// # Safety /// -/// Implementors must be safe to duplicate with a byte-for-byte copy. Values -/// copied back from device memory must have a bit pattern that is valid for -/// `Self`, and the all-zero bit pattern must also be valid because -/// [`DeviceBuffer::zeroed`] initializes memory with zero bytes. +/// Implementors must be safe to duplicate with a byte-for-byte copy, and the +/// all-zero bit pattern must be a valid value because [`DeviceBuffer::zeroed`] +/// initializes memory with zero bytes. Types whose zero pattern is invalid, +/// such as `NonZeroU32`, cannot implement `DeviceCopy`. /// -/// `Copy` alone is not enough: types such as `bool`, `char`, and -/// `NonZeroU32` are `Copy`, but not every byte pattern is a valid value of -/// those types. `DeviceCopy` is the stronger promise required when -/// `DeviceBuffer` turns raw device bytes back into initialized Rust values. +/// `DeviceCopy` does not promise that every bit pattern is valid. `bool` +/// implements it although only `0` and `1` are `bool` values. That is sound +/// because a typed `DeviceBuffer` holds valid `T` values by construction: +/// the safe constructors copy `T` values from the host or write zero bytes, a +/// kernel that writes through `&mut T` writes a `T`, and the raw constructors +/// are `unsafe`, so their caller vouches for the contents. The one place +/// arbitrary bytes could become typed values is the safe reinterpreting casts, +/// [`DeviceBuffer::cast_elem`] and [`DeviceBuffer::cast_chunks`]. They require +/// the target to implement [`bytemuck::AnyBitPattern`], which `bool` does not, +/// so a `DeviceBuffer` can only come from memory written as `bool`. pub unsafe trait DeviceCopy: Copy {} macro_rules! impl_device_copy { @@ -87,6 +92,11 @@ impl_device_copy!( f64 ); +// Zero is `false`, and the trait documentation explains why the `0`/`1` +// validity requirement does not stop `bool` from being `DeviceCopy`. It is not +// `bytemuck::AnyBitPattern`, so the reinterpreting casts cannot produce it. +unsafe impl DeviceCopy for bool {} + unsafe impl DeviceCopy for [T; N] {} unsafe impl DeviceCopy for *const T {} unsafe impl DeviceCopy for *mut T {} @@ -120,36 +130,42 @@ unsafe impl DeviceCopy for f16 {} unsafe impl DeviceCopy for half::bf16 {} unsafe impl DeviceCopy for half::f16 {} +// SAFETY: these transparent numeric storage wrappers own no resources, have +// valid all-zero representations, and can be duplicated byte-for-byte. Packed +// FP4 remains one byte holding two logical values; copying does not unpack it. +impl_device_copy!( + crate::tf32, + crate::f8e4m3fn, + crate::f8e5m2, + crate::f8e8m0fnu, + crate::f8e5m3fnu, + crate::f4e2m1fnx2 +); + /// Owning handle to a contiguous device allocation of `T` elements. /// -/// Holds a raw device pointer, element count, and a reference-counted -/// context that keeps the CUDA context alive. Synchronous allocations are -/// freed with `cuMemFree`. Dropping a stream-ordered allocation synchronizes -/// its context before enqueueing `cuMemFreeAsync` on its retained allocation -/// stream. Use the unsafe [`DeviceBuffer::drop_async`] when the caller can -/// provide explicit stream ordering and must avoid that context-wide -/// synchronization. +/// Shares an allocation owner, its CUDA context, and an access tracker with +/// tensor and byte views. Safe submissions retain independent access leases +/// until completion, including when a launch future is dropped or forgotten. +/// Drop retires deallocation on a worker without waiting on the caller's +/// thread. Unsafe raw-pointer users must retain storage through their own work. /// /// Device buffers may only transfer plain device-copyable values. Owning host /// types such as [`String`] are rejected because copying their bytes to and /// from device memory would not preserve Rust ownership invariants. /// /// ```compile_fail -/// # use cuda_core::{CudaStream, DeviceBuffer}; -/// # fn rejects_non_device_copy(stream: &CudaStream) { +/// # use cuda_core::{Stream, DeviceBuffer}; +/// # fn rejects_non_device_copy(stream: &Stream) { /// let _ = DeviceBuffer::::zeroed(stream, 1); /// # } /// ``` pub struct DeviceBuffer { + storage: crate::RawDeviceBuffer, + shared_access_writes: bool, ptr: CUdeviceptr, len: usize, num_bytes: usize, - ctx: Arc, - /// Retains the allocation stream for a stream-ordered (`cuMemAllocAsync`) - /// allocation. Ordinary `Drop` first synchronizes the context so work - /// submitted on any stream has completed, then frees on this stream. - /// `None` identifies a synchronous (`cuMemAlloc`) allocation. - dealloc_stream: Option>, _marker: PhantomData, } @@ -160,28 +176,90 @@ unsafe impl Send for DeviceBuffer {} // return Copy values. No interior mutability. unsafe impl Sync for DeviceBuffer {} -impl Drop for DeviceBuffer { - fn drop(&mut self) { - if self.ptr != 0 { - self.ctx.record_err(self.ctx.bind_to_thread()); - // Safe buffer operations can enqueue work on any stream in this - // context. Synchronize all of them before implicitly freeing a - // stream-ordered allocation. - let result = match &self.dealloc_stream { - Some(stream) => match self.ctx.synchronize() { - Ok(()) => unsafe { - crate::simt::memory::free_async(self.ptr, stream.cu_stream()) - }, - Err(error) => Err(error), - }, - None => unsafe { crate::simt::memory::free_sync(self.ptr) }, - }; - self.ctx.record_err(result); +impl DeviceBuffer { + fn enqueue_access( + &self, + stream: &Stream, + write: bool, + enqueue: impl FnOnce() -> Result<(), DriverError>, + ) -> Result<(), DriverError> { + let lease = self.storage.acquire(stream, write, false)?; + self.storage.context().bind_to_thread()?; + let submitted = enqueue(); + let retained = lease.retain_until_complete(stream); + submitted.and(retained) + } + + /// The shared byte storage and allocation identity underlying this view. + pub fn storage(&self) -> &crate::RawDeviceBuffer { + &self.storage + } + /// Whether a shared kernel reference may mutate elements through interior + /// mutability. Representation wrappers and unclassified raw types are + /// conservative; ordinary DeviceCopy values can share read accesses. + #[doc(hidden)] + pub fn shared_access_writes(&self) -> bool { + self.shared_access_writes + } + + /// Erases the element type without copying or releasing storage. + pub fn into_storage(self) -> crate::RawDeviceBuffer { + self.storage + } + + /// Constructs a typed view over shared byte storage. + /// + /// # Safety + /// Every initialized element must be a valid `T`; the caller must preserve + /// element validity across all aliases (including future writes through + /// aliases with another element type). The storage must be initialized before + /// any read, and raw foreign accesses must obey its access-exclusion contract. + pub unsafe fn from_storage( + storage: crate::RawDeviceBuffer, + len: usize, + ) -> Result { + let num_bytes = allocation_size::(len)?; + let ptr = storage.cu_deviceptr(); + if num_bytes != storage.len_bytes() + || (num_bytes != 0 && !(ptr as usize).is_multiple_of(std::mem::align_of::())) + { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + Ok(Self { + ptr, + len, + num_bytes, + storage, + shared_access_writes: true, + _marker: PhantomData, + }) + } + + /// Borrows a contiguous range of elements while retaining the same storage + /// owner and access tracker. No device memory is copied. + pub fn slice( + &self, + range: std::ops::Range, + ) -> Result, DriverError> { + if range.start > range.end || range.end > self.len { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); } + let storage = self.storage.slice_bytes( + allocation_size::(range.start)?, + allocation_size::(range.end - range.start)?, + )?; + Ok(DeviceBufferView { + storage, + shared_access_writes: self.shared_access_writes, + len: range.end - range.start, + _borrow: PhantomData, + }) } -} -impl DeviceBuffer { /// Returns the raw `CUdeviceptr` for use in kernel argument lists. #[inline] pub fn cu_deviceptr(&self) -> CUdeviceptr { @@ -208,8 +286,8 @@ impl DeviceBuffer { /// Returns a reference to the owning context. #[inline] - pub fn context(&self) -> &Arc { - &self.ctx + pub fn context(&self) -> &Arc { + self.storage.context() } /// Constructs a `DeviceBuffer` from pre-existing raw parts. @@ -227,7 +305,7 @@ impl DeviceBuffer { /// # Panics /// /// Panics if `len * size_of::()` overflows `usize`. - pub unsafe fn from_raw_parts(ptr: CUdeviceptr, len: usize, ctx: Arc) -> Self { + pub unsafe fn from_raw_parts(ptr: CUdeviceptr, len: usize, ctx: Arc) -> Self { // SAFETY: `from_raw_parts` has the same raw-allocation safety contract, // with no stream-ordered deallocation metadata attached. unsafe { Self::from_raw_parts_with_dealloc_stream(ptr, len, ctx, None) } @@ -236,17 +314,22 @@ impl DeviceBuffer { unsafe fn from_raw_parts_with_dealloc_stream( ptr: CUdeviceptr, len: usize, - ctx: Arc, - dealloc_stream: Option>, + ctx: Arc, + dealloc_stream: Option>, ) -> Self { let num_bytes = allocation_size::(len).expect("DeviceBuffer::from_raw_parts byte size overflow"); + let storage = if let Some(stream) = &dealloc_stream { + unsafe { crate::RawDeviceBuffer::from_raw_parts_async(ptr, num_bytes, stream.clone()) } + } else { + unsafe { crate::RawDeviceBuffer::from_raw_parts(ptr, num_bytes, ctx.clone()) } + }; Self { + storage, + shared_access_writes: true, ptr, len, num_bytes, - ctx, - dealloc_stream, _marker: PhantomData, } } @@ -257,63 +340,104 @@ impl DeviceBuffer { /// allocator that matches how it was created. For stream-ordered /// allocations, this does not return the stored deallocation stream; the /// caller must already know which stream to use for `cuMemFreeAsync`. - pub fn into_raw_parts(self) -> (CUdeviceptr, usize, Arc) { + pub fn into_raw_parts(self) -> (CUdeviceptr, usize, Arc) { let (ptr, len, ctx, _dealloc_stream) = self.into_all_raw_parts(); (ptr, len, ctx) } - fn into_all_raw_parts( - self, - ) -> ( - CUdeviceptr, - usize, - Arc, - Option>, - ) { - // Suppress the buffer's `Drop` (which would free `ptr`) while still - // moving out the heap-owned fields. Callers that do not need - // `dealloc_stream` can drop it after this helper returns. - let this = std::mem::ManuallyDrop::new(self); - let ptr = this.ptr; - let len = this.len; - // SAFETY: `this` is `ManuallyDrop` and is never used again, so reading - // out its non-`Copy` fields takes ownership without a double drop. - let ctx = unsafe { std::ptr::read(&this.ctx) }; - let dealloc_stream = unsafe { std::ptr::read(&this.dealloc_stream) }; - (ptr, len, ctx, dealloc_stream) + fn into_all_raw_parts(self) -> (CUdeviceptr, usize, Arc, Option>) { + let Self { storage, len, .. } = self; + let (ptr, _, ctx, stream) = storage.into_owned_parts(); + (ptr, len, ctx, stream) } /// Reinterpret the element type of this buffer as `A`. /// - /// `A` must have the same size and alignment as `T` (e.g. `A` is - /// `#[repr(transparent)]` over `T`). This is the "atomic-slice launch - /// mapping" for issue #151: allocate and initialize a plain - /// `DeviceBuffer`, then hand it to a kernel that takes - /// `&[DeviceAtomicU64]`. The pointer, length, and bytes are unchanged; - /// only the element type the kernel sees changes to one whose pointee - /// permits shared mutation (so rustc does not mark it `readonly`/`noalias`). + /// `A` must have the same size and alignment as `T`. The pointer, length, + /// and bytes are unchanged; only the element type changes. Element counts + /// are preserved because `size_of::() == size_of::()`. /// - /// Element counts are preserved because `size_of::() == size_of::()`. - pub fn cast_elem(self) -> DeviceBuffer { + /// `A: AnyBitPattern` is what makes this safe: every bit pattern is a valid + /// `A`, so no byte written as a `T` can become an invalid `A`. Types with + /// invalid bit patterns, `bool` for one, are rejected at compile time, so a + /// `DeviceBuffer` can only come from memory written as `bool`. + /// + /// Atomic types and other `#[repr(transparent)]` wrappers that are not + /// `Copy` cannot be `AnyBitPattern`; use [`Self::wrap_elem`] and + /// [`Self::peel_elem`] for those. + pub fn cast_elem(self) -> DeviceBuffer { + // SAFETY: every bit pattern is a valid `A`, so the bytes of any `T` + // are valid `A` values; `relabel` checks the layout. + let mut buffer = unsafe { self.relabel() }; + buffer.shared_access_writes = false; + buffer + } + + /// View this buffer through a `#[repr(transparent)]` wrapper of `T`. + /// + /// This is the atomic-slice launch mapping for issue #151: allocate and + /// initialize a plain `DeviceBuffer`, then hand it to a kernel that + /// takes `&[DeviceAtomicU64]`. The pointer, length, and bytes are + /// unchanged; only the element type the kernel sees changes to one whose + /// pointee permits shared mutation, so rustc does not mark the slice + /// `readonly`/`noalias`. + /// + /// `W: DeviceRepr` guarantees the same layout and valid initialized + /// representations in both directions, without element destruction or + /// additional ownership obligations. It permits no shared-reference conversions. + /// The wrapper need not be `Copy` or `DeviceCopy`; the + /// atomic types are neither. [`Self::peel_elem`] goes back to `T`, for + /// example before [`Self::to_host_vec`]. + /// + /// This consumes the buffer and retains its allocation, context, and + /// deallocation stream. It does not synchronize pending kernel accesses; + /// those must obey the launch API's borrowing and ordering requirements. + pub fn wrap_elem>(self) -> DeviceBuffer { + // SAFETY: DeviceRepr guarantees every initialized T representation + // is valid as W; relabel preserves ownership and checks the layout. + unsafe { self.relabel() } + } + + /// Remove a `#[repr(transparent)]` wrapper from the element type. + /// + /// The inverse of [`Self::wrap_elem`]: `T` must implement `DeviceRepr`, + /// so a `DeviceBuffer` peels back to the + /// `DeviceBuffer` it was made from, which can be read on the host. + pub fn peel_elem(self) -> DeviceBuffer + where + T: DeviceRepr, + { + // SAFETY: DeviceRepr guarantees every initialized T representation + // is valid as I; relabel preserves ownership and checks the layout. + let mut buffer = unsafe { self.relabel() }; + buffer.shared_access_writes = false; + buffer + } + + /// Hand this allocation to a `DeviceBuffer` of the same length. + /// + /// # Safety + /// + /// Every value of `T` in the allocation must also be a valid `A`. The + /// three public casts above each establish that through a trait bound. + unsafe fn relabel(self) -> DeviceBuffer { assert_eq!( std::mem::size_of::(), std::mem::size_of::(), - "cast_elem requires the same element size" + "element casts require the same element size" ); assert_eq!( std::mem::align_of::(), std::mem::align_of::(), - "cast_elem requires the same element alignment" + "element casts require the same element alignment" ); - let (ptr, len, ctx, dealloc_stream) = self.into_all_raw_parts(); - // SAFETY: `ptr` came from a valid `DeviceBuffer` allocation of `len` - // elements; `A` has identical size and alignment, so the same allocation - // is a valid `DeviceBuffer` of the same length and the same byte - // extent. Ownership transfers; the original buffer's `Drop` is - // suppressed by `into_all_raw_parts`, and the allocation metadata is - // preserved for the new element type. - unsafe { - DeviceBuffer::::from_raw_parts_with_dealloc_stream(ptr, len, ctx, dealloc_stream) + DeviceBuffer { + storage: self.storage, + shared_access_writes: true, + ptr: self.ptr, + len: self.len, + num_bytes: self.num_bytes, + _marker: PhantomData, } } @@ -357,7 +481,12 @@ impl DeviceBuffer { /// Returning the original on failure rather than panicking keeps the /// fallback path available, since a caller that cannot widen usually has a /// scalar version to fall back to. - pub fn cast_chunks(self) -> Result, Self> { + /// + /// As with [`Self::cast_elem`], `A: AnyBitPattern` is what makes the + /// reinterpretation safe: every bit pattern is a valid `A`. A host-defined + /// vector element gets there with `unsafe impl bytemuck::Pod` next to its + /// `DeviceCopy` impl. + pub fn cast_chunks(self) -> Result, Self> { let bytes = self.len.saturating_mul(std::mem::size_of::()); let Some(new_len) = chunk_cast_len( bytes, @@ -367,14 +496,13 @@ impl DeviceBuffer { ) else { return Err(self); }; - let (ptr, _len, ctx, dealloc_stream) = self.into_all_raw_parts(); - // SAFETY: `ptr` came from a valid allocation of `bytes` bytes, checked - // above to be exactly `new_len` elements of `A` and to be aligned for - // `A`. The allocation and its ownership are unchanged; only the element - // type and the count describing the same bytes change. `Drop` on the - // original is suppressed by `into_all_raw_parts`. - Ok(unsafe { - DeviceBuffer::::from_raw_parts_with_dealloc_stream(ptr, new_len, ctx, dealloc_stream) + Ok(DeviceBuffer { + storage: self.storage, + shared_access_writes: false, + ptr: self.ptr, + len: new_len, + num_bytes: self.num_bytes, + _marker: PhantomData, }) } @@ -383,7 +511,7 @@ impl DeviceBuffer { /// For choosing between a wide and a scalar path without consuming the /// buffer to find out. #[must_use] - pub fn can_cast_chunks(&self) -> bool { + pub fn can_cast_chunks(&self) -> bool { let bytes = self.len.saturating_mul(std::mem::size_of::()); chunk_cast_len( bytes, @@ -396,6 +524,28 @@ impl DeviceBuffer { } impl DeviceBuffer { + fn empty_copy(len: usize, ctx: Arc) -> Self { + // SAFETY: a null pointer with zero bytes is never dereferenced + // and Drop/drop_async ignore it. DeviceCopy excludes interior mutability. + let mut buf = unsafe { Self::from_raw_parts(0, len, ctx) }; + buf.shared_access_writes = false; + buf + } + + /// Constructs a plain-value typed view of shared byte storage. + /// + /// # Safety + /// The initialization, validity, and aliasing requirements of + /// `from_storage` apply. DeviceCopy additionally excludes interior mutability. + pub unsafe fn from_storage_copy( + storage: crate::RawDeviceBuffer, + len: usize, + ) -> Result { + let mut buffer = unsafe { Self::from_storage(storage, len)? }; + buffer.shared_access_writes = false; + Ok(buffer) + } + /// Allocates device memory, copies `data` from the host on `stream`, and /// synchronizes `stream` before returning. /// @@ -414,7 +564,7 @@ impl DeviceBuffer { /// stream synchronization run. If either step fails, the early return /// drops the buffer and its `Drop` impl frees the allocation, so no /// device memory is leaked. - pub fn from_host(stream: &CudaStream, data: &[T]) -> Result { + pub fn from_host(stream: &Stream, data: &[T]) -> Result { let ctx = stream.context().clone(); let len = data.len(); let num_bytes = allocation_size::(len)?; @@ -422,27 +572,22 @@ impl DeviceBuffer { // cuMemAlloc rejects zero-byte requests with CUDA_ERROR_INVALID_VALUE, // so represent an empty buffer as a null pointer (Drop skips it). if num_bytes == 0 { - // SAFETY: a null pointer with zero bytes is never dereferenced - // and Drop ignores it. - return Ok(unsafe { Self::from_raw_parts(0, len, ctx) }); + return Ok(Self::empty_copy(len, ctx)); } - let ptr = unsafe { crate::simt::memory::malloc_sync(num_bytes)? }; + ctx.bind_to_thread()?; + let ptr = unsafe { crate::memory::malloc_sync(num_bytes)? }; // SAFETY: `ptr` was just allocated with `num_bytes` bytes in the // stream's context; ownership transfers to `buf` here so any early // return below frees it through the buffer's own `Drop`. - let buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; + let mut buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; let enqueue_result = unsafe { - crate::simt::memory::memcpy_htod_async( - buf.ptr, - data.as_ptr(), - num_bytes, - stream.cu_stream(), - ) + crate::memory::memcpy_htod_async(buf.ptr, data.as_ptr(), num_bytes, stream.cu_stream()) }; let sync_result = stream.synchronize(); enqueue_result?; sync_result?; + buf.shared_access_writes = false; Ok(buf) } @@ -455,9 +600,9 @@ impl DeviceBuffer { /// still be reading from `data` after the borrow is released. The caller /// must ensure `data` is not dropped, freed, mutated, or aliased until the /// enqueued copy has completed, typically after the next - /// [`CudaStream::synchronize`] call or a stream-ordered event wait. + /// [`Stream::synchronize`] call or a stream-ordered event wait. pub unsafe fn from_host_async_unchecked( - stream: &CudaStream, + stream: &Stream, data: &[T], ) -> Result { let ctx = stream.context().clone(); @@ -467,24 +612,19 @@ impl DeviceBuffer { // cuMemAlloc rejects zero-byte requests with CUDA_ERROR_INVALID_VALUE, // so represent an empty buffer as a null pointer (Drop skips it). if num_bytes == 0 { - // SAFETY: a null pointer with zero bytes is never dereferenced - // and Drop ignores it. - return Ok(unsafe { Self::from_raw_parts(0, len, ctx) }); + return Ok(Self::empty_copy(len, ctx)); } - let ptr = unsafe { crate::simt::memory::malloc_sync(num_bytes)? }; + ctx.bind_to_thread()?; + let ptr = unsafe { crate::memory::malloc_sync(num_bytes)? }; // SAFETY: `ptr` was just allocated with `num_bytes` bytes in the // stream's context; ownership transfers to `buf` here so any early // return below frees it through the buffer's own `Drop`. - let buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; - unsafe { - crate::simt::memory::memcpy_htod_async( - buf.ptr, - data.as_ptr(), - num_bytes, - stream.cu_stream(), - )?; - } + let mut buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; + buf.enqueue_access(stream, true, || unsafe { + crate::memory::memcpy_htod_async(buf.ptr, data.as_ptr(), num_bytes, stream.cu_stream()) + })?; + buf.shared_access_writes = false; Ok(buf) } @@ -497,8 +637,8 @@ impl DeviceBuffer { /// /// `PinnedHostBuffer` currently uses `cuMemAllocHost` without the /// `PORTABLE` flag, so the allocation is only pinned in the context that - /// created it. In debug builds this asserts that `data` and `stream` - /// share the same [`CudaContext`]. + /// created it. This checks that `data` and `stream` share the same + /// CUDA context. /// /// The device-to-host counterparts are [`Self::copy_to_pinned_host`] /// (blocking) and [`Self::copy_to_pinned_host_async`] (non-blocking). To @@ -512,25 +652,30 @@ impl DeviceBuffer { /// after this function returns. The caller is responsible for ensuring /// `data` is not dropped, freed, mutated, or aliased until the enqueued /// copy has completed, typically after the next - /// [`CudaStream::synchronize`] call or a stream-ordered event wait. + /// [`Stream::synchronize`] call or a stream-ordered event wait. /// Dropping `data` before that synchronization point calls /// `cuMemFreeHost` while the in-flight transfer is still reading the /// buffer, which is undefined behavior. pub unsafe fn from_pinned_host( - stream: &CudaStream, + stream: &Stream, data: &PinnedHostBuffer, ) -> Result { - debug_assert!( - Arc::ptr_eq(data.context(), stream.context()), - "pinned host buffer and stream must belong to the same CUDA context" - ); + if data.context().as_ref() != stream.context().as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } // SAFETY: this method's safety contract requires the caller to keep // the pinned source valid until the enqueued copy completes. unsafe { Self::from_host_async_unchecked(stream, data.as_slice()) } } - /// Allocates zero-initialized device memory of `len` elements, enqueued - /// on `stream`. + /// Allocates zero-initialized device memory of `len` elements. + /// + /// The zero-fill is enqueued on `stream`, then this method synchronizes + /// that stream before returning. The initialized buffer is therefore safe + /// to use on another stream, including for types such as `bool` whose + /// validity depends on initialization having completed. /// /// A `len` of zero (or a zero-sized `T`) yields an empty buffer without /// touching the driver allocator. @@ -539,29 +684,30 @@ impl DeviceBuffer { /// /// The returned buffer takes ownership of the device allocation /// immediately after `malloc_sync`, before the fallible - /// `memset_d8_async` enqueue runs. If the enqueue fails, the early - /// return drops the buffer and its `Drop` impl frees the allocation, so - /// no device memory is leaked. - pub fn zeroed(stream: &CudaStream, len: usize) -> Result { + /// `memset_d8_async` enqueue and stream synchronization run. If either + /// fails, the early return drops the buffer and its `Drop` impl frees the + /// allocation, so no device memory is leaked. + pub fn zeroed(stream: &Stream, len: usize) -> Result { let ctx = stream.context().clone(); let num_bytes = allocation_size::(len)?; // cuMemAlloc rejects zero-byte requests with CUDA_ERROR_INVALID_VALUE, // so represent an empty buffer as a null pointer (Drop skips it). if num_bytes == 0 { - // SAFETY: a null pointer with zero bytes is never dereferenced - // and Drop ignores it. - return Ok(unsafe { Self::from_raw_parts(0, len, ctx) }); + return Ok(Self::empty_copy(len, ctx)); } - let ptr = unsafe { crate::simt::memory::malloc_sync(num_bytes)? }; + ctx.bind_to_thread()?; + let ptr = unsafe { crate::memory::malloc_sync(num_bytes)? }; // SAFETY: `ptr` was just allocated with `num_bytes` bytes in the // stream's context; ownership transfers to `buf` here so any early // return below frees it through the buffer's own `Drop`. - let buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; + let mut buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; unsafe { - crate::simt::memory::memset_d8_async(buf.ptr, 0, num_bytes, stream.cu_stream())?; + crate::memory::memset_d8_async(buf.ptr, 0, num_bytes, stream.cu_stream())?; } + stream.synchronize()?; + buf.shared_access_writes = false; Ok(buf) } @@ -569,18 +715,35 @@ impl DeviceBuffer { /// /// Synchronizes on `stream` before returning so the host vector is safe /// to read immediately. - pub fn to_host_vec(&self, stream: &CudaStream) -> Result, DriverError> { - let mut host = Vec::with_capacity(self.len); - unsafe { - crate::simt::memory::memcpy_dtoh_async( + pub fn to_host_vec(&self, stream: &Stream) -> Result, DriverError> { + let mut host = Vec::::with_capacity(self.len); + if self.num_bytes() == 0 { + // SAFETY: empty buffers or zero-sized DeviceCopy values need no bytes. + unsafe { + host.set_len(self.len); + } + return Ok(host); + } + let lease = self.storage.acquire(stream, false, false)?; + self.storage.context().bind_to_thread()?; + let submitted = unsafe { + crate::memory::memcpy_dtoh_async( host.as_mut_ptr(), self.ptr, self.num_bytes(), stream.cu_stream(), - )?; + ) + }; + let completed = stream.synchronize(); + if let Err(error) = completed { + // The driver has not established that it stopped writing host memory. + std::mem::forget((host, lease)); + return Err(error); + } + submitted?; + unsafe { + host.set_len(self.len); } - stream.synchronize()?; - unsafe { host.set_len(self.len) }; Ok(host) } @@ -588,22 +751,26 @@ impl DeviceBuffer { /// /// Synchronizes on `stream` before returning. Panics if /// `dst.len() < self.len()`. - pub fn copy_to_host(&self, stream: &CudaStream, dst: &mut [T]) -> Result<(), DriverError> { - assert!( - dst.len() >= self.len, - "destination slice too small: {} < {}", - dst.len(), - self.len - ); - unsafe { - crate::simt::memory::memcpy_dtoh_async( + pub fn copy_to_host(&self, stream: &Stream, dst: &mut [T]) -> Result<(), DriverError> { + assert!(dst.len() >= self.len(), "destination host slice too small"); + if self.num_bytes() == 0 { + return Ok(()); + } + let lease = self.storage.acquire(stream, false, false)?; + self.storage.context().bind_to_thread()?; + let submitted = unsafe { + crate::memory::memcpy_dtoh_async( dst.as_mut_ptr(), self.ptr, self.num_bytes(), stream.cu_stream(), - )?; + ) + }; + let completed = stream.synchronize(); + if completed.is_err() { + std::mem::forget(lease); } - stream.synchronize() + submitted.and(completed) } /// Copies the buffer contents into an existing pinned host buffer and @@ -617,7 +784,7 @@ impl DeviceBuffer { /// synchronize the stream later. pub fn copy_to_pinned_host( &self, - stream: &CudaStream, + stream: &Stream, dst: &mut PinnedHostBuffer, ) -> Result<(), DriverError> { // SAFETY: we synchronize the stream below before returning, so the @@ -635,7 +802,7 @@ impl DeviceBuffer { /// `PinnedHostBuffer` currently uses `cuMemAllocHost` without the /// `PORTABLE` flag, so the allocation is only pinned in the context that /// created it. In debug builds this asserts that `dst` and `stream` - /// share the same [`CudaContext`]. + /// share the same [`Context`]. /// /// # Safety /// @@ -643,33 +810,35 @@ impl DeviceBuffer { /// returns; CUDA may still be writing into `dst`'s pinned pointer long /// after this function returns. The caller is responsible for ensuring /// `dst` is not dropped, freed, read, or aliased until the enqueued copy - /// has completed, typically after the next [`CudaStream::synchronize`] + /// has completed, typically after the next [`Stream::synchronize`] /// call or a stream-ordered event wait. Dropping `dst` before that /// synchronization point calls `cuMemFreeHost` while the in-flight /// transfer is still writing the buffer, which is undefined behavior. pub unsafe fn copy_to_pinned_host_async( &self, - stream: &CudaStream, + stream: &Stream, dst: &mut PinnedHostBuffer, ) -> Result<(), DriverError> { - debug_assert!( - Arc::ptr_eq(dst.context(), stream.context()), - "pinned host buffer and stream must belong to the same CUDA context" - ); assert!( - dst.len() >= self.len, - "destination pinned host buffer too small: {} < {}", - dst.len(), - self.len + dst.len() >= self.len(), + "destination pinned host buffer too small" ); - unsafe { - crate::simt::memory::memcpy_dtoh_async( + if dst.context().as_ref() != self.storage.context().as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + if self.num_bytes() == 0 { + return Ok(()); + } + self.enqueue_access(stream, false, || unsafe { + crate::memory::memcpy_dtoh_async( dst.as_mut_ptr(), self.ptr, self.num_bytes(), stream.cu_stream(), ) - } + }) } /// Enqueues a host-to-device copy from a pinned host buffer into this @@ -685,8 +854,8 @@ impl DeviceBuffer { /// /// `PinnedHostBuffer` currently uses `cuMemAllocHost` without the /// `PORTABLE` flag, so the allocation is only pinned in the context that - /// created it. In debug builds this asserts that `src` and `stream` - /// share the same [`CudaContext`]. + /// created it. This checks that `src` and `stream` share the same + /// CUDA context. /// /// # Safety /// @@ -695,34 +864,35 @@ impl DeviceBuffer { /// after this function returns. The caller is responsible for ensuring /// `src` is not dropped, freed, mutated, or aliased until the enqueued /// copy has completed, typically after the next - /// [`CudaStream::synchronize`] call or a stream-ordered event wait. + /// [`Stream::synchronize`] call or a stream-ordered event wait. /// Dropping `src` before that synchronization point calls /// `cuMemFreeHost` while the in-flight transfer is still reading the /// buffer, which is undefined behavior. pub unsafe fn copy_from_pinned_host_async( &mut self, - stream: &CudaStream, + stream: &Stream, src: &PinnedHostBuffer, ) -> Result<(), DriverError> { - debug_assert!( - Arc::ptr_eq(src.context(), stream.context()), - "pinned host buffer and stream must belong to the same CUDA context" - ); assert!( - src.len() <= self.len, - "source pinned host buffer too large: {} > {}", - src.len(), - self.len + src.len() <= self.len(), + "source pinned host buffer too large" ); - let num_bytes = src.num_bytes(); - unsafe { - crate::simt::memory::memcpy_htod_async( + if src.context().as_ref() != self.storage.context().as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + if src.num_bytes() == 0 { + return Ok(()); + } + self.enqueue_access(stream, true, || unsafe { + crate::memory::memcpy_htod_async( self.ptr, src.as_ptr(), - num_bytes, + src.num_bytes(), stream.cu_stream(), ) - } + }) } /// Allocates `len` elements of uninitialized device memory, enqueued on @@ -732,35 +902,30 @@ impl DeviceBuffer { /// the returned buffer are undefined until the caller writes them. /// /// The buffer co-owns `stream` (via the `Arc`) so its implicit `Drop` can - /// release the stream-ordered allocation after synchronizing the context. - /// Call the unsafe [`Self::drop_async`] to free explicitly on a chosen - /// stream without a context-wide synchronization. + /// retire the stream-ordered allocation after tracked accesses complete. + /// Unsafe callers must establish allocation and use ordering. /// /// # Safety /// /// Reading from the returned buffer before any kernel or memcpy has /// written it is undefined behavior. pub unsafe fn uninitialized_async( - stream: &Arc, + stream: &Arc, len: usize, ) -> Result { let ctx = stream.context().clone(); let num_bytes = allocation_size::(len)?; if num_bytes == 0 { - // SAFETY: a null pointer with zero bytes is never dereferenced - // and Drop/drop_async ignore it. - return Ok(unsafe { Self::from_raw_parts(0, len, ctx) }); + return Ok(Self::empty_copy(len, ctx)); } - let ptr = unsafe { crate::simt::memory::malloc_async(stream.cu_stream(), num_bytes)? }; - Ok(Self { - ptr, - len, - num_bytes, - ctx, - dealloc_stream: Some(stream.clone()), - _marker: PhantomData, - }) + ctx.bind_to_thread()?; + let ptr = unsafe { crate::memory::malloc_async(stream.cu_stream(), num_bytes)? }; + let mut buffer = unsafe { + Self::from_raw_parts_with_dealloc_stream(ptr, len, ctx, Some(stream.clone())) + }; + buffer.shared_access_writes = false; + Ok(buffer) } /// Copies `other` into `self` device-to-device, enqueued on `stream`. @@ -769,24 +934,39 @@ impl DeviceBuffer { pub fn copy_from_device_async( &mut self, other: &DeviceBuffer, - stream: &CudaStream, + stream: &Stream, ) -> Result<(), DriverError> { - assert_eq!( - self.len, other.len, - "device-to-device copy length mismatch: dst {} != src {}", - self.len, other.len - ); + assert_eq!(self.len, other.len, "device-to-device copy length mismatch"); if self.num_bytes() == 0 { return Ok(()); } - unsafe { - crate::simt::memory::memcpy_dtod_async( - self.ptr, - other.ptr, - self.num_bytes(), - stream.cu_stream(), - ) + if self.storage.context().as_ref() != stream.context().as_ref() + || other.storage.context().as_ref() != self.storage.context().as_ref() + { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + if self.ptr == other.ptr { + return Ok(()); + } + if self.ptr < other.ptr + other.num_bytes as u64 + && other.ptr < self.ptr + self.num_bytes as u64 + { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); } + self.enqueue_access(stream, true, || { + other.enqueue_access(stream, false, || unsafe { + crate::memory::memcpy_dtod_async( + self.ptr, + other.ptr, + self.num_bytes(), + stream.cu_stream(), + ) + }) + }) } /// Copies `src` into `self` host-to-device on `stream` and synchronizes @@ -795,7 +975,7 @@ impl DeviceBuffer { /// The synchronization keeps this safe for borrowed host slices: `src` /// may be dropped, reused, or mutated immediately after this function /// returns. Panics if `src.len() != self.len()`. - pub fn copy_from_host(&mut self, stream: &CudaStream, src: &[T]) -> Result<(), DriverError> { + pub fn copy_from_host(&mut self, stream: &Stream, src: &[T]) -> Result<(), DriverError> { // SAFETY: this safe wrapper synchronizes `stream` before returning, // so the borrowed host slice cannot be used by CUDA after this call. let enqueue_result = unsafe { self.copy_from_host_async_unchecked(stream, src) }; @@ -817,31 +997,25 @@ impl DeviceBuffer { /// still be reading from `src` after the borrow is released. The caller /// must ensure `src` is not dropped, freed, mutated, or aliased until the /// enqueued copy has completed, typically after the next - /// [`CudaStream::synchronize`] call or a stream-ordered event wait. + /// [`Stream::synchronize`] call or a stream-ordered event wait. /// Panics if `src.len() != self.len()`. pub unsafe fn copy_from_host_async_unchecked( &mut self, - stream: &CudaStream, + stream: &Stream, src: &[T], ) -> Result<(), DriverError> { - assert_eq!( - self.len, - src.len(), - "host-to-device copy length mismatch: dst {} != src {}", - self.len, - src.len() - ); + assert_eq!(self.len, src.len(), "host-to-device copy length mismatch"); if self.num_bytes() == 0 { return Ok(()); } - unsafe { - crate::simt::memory::memcpy_htod_async( + self.enqueue_access(stream, true, || unsafe { + crate::memory::memcpy_htod_async( self.ptr, src.as_ptr(), self.num_bytes(), stream.cu_stream(), ) - } + }) } /// Consumes the buffer and frees it asynchronously on `stream`. @@ -861,42 +1035,33 @@ impl DeviceBuffer { /// /// The caller must ensure that every stream other than the allocation /// stream that has pending work touching this buffer is ordered before - /// `stream` (for example via [`CudaStream::join`]). Only the allocation + /// `stream` (for example via [`Stream::join`]). Only the allocation /// stream is joined automatically; an unordered third stream still /// racing the free is a driver-level use-after-free. This is the same /// deferred-use contract as [`Self::copy_from_host_async_unchecked`]: - /// ordinary [`Drop`] is the safe alternative and synchronizes the whole - /// context first. - pub unsafe fn drop_async(mut self, stream: &CudaStream) -> Result<(), DriverError> { - if self.ctx.as_ref() != stream.context().as_ref() { + /// ordinary [`Drop`] relies on completion leases acquired by safe APIs. + /// Unsafe raw-pointer launches must retain their own completion leases. + pub unsafe fn drop_async(self, stream: &Stream) -> Result<(), DriverError> { + if self.context().as_ref() != stream.context().as_ref() { return Err(DriverError( cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, )); } - if self.ptr == 0 { - return Ok(()); - } - - self.ctx.bind_to_thread()?; - if let Some(allocation_stream) = &self.dealloc_stream { - if allocation_stream.as_ref() != stream { - stream.join(allocation_stream)?; - } + stream.context().bind_to_thread()?; + if let Some(allocation_stream) = self.storage.allocation_stream() { + stream.join(allocation_stream)?; } - - let ptr = self.ptr; - self.ptr = 0; - unsafe { crate::simt::memory::free_async(ptr, stream.cu_stream()) } + crate::completion::retain_until_complete(self.storage, stream) } /// Zeroes every byte in the buffer asynchronously on `stream`. - pub fn zero_async(&mut self, stream: &CudaStream) -> Result<(), DriverError> { + pub fn zero_async(&mut self, stream: &Stream) -> Result<(), DriverError> { if self.num_bytes() == 0 { return Ok(()); } - unsafe { - crate::simt::memory::memset_d8_async(self.ptr, 0, self.num_bytes(), stream.cu_stream()) - } + self.enqueue_access(stream, true, || unsafe { + crate::memory::memset_d8_async(self.ptr, 0, self.num_bytes(), stream.cu_stream()) + }) } } @@ -925,6 +1090,69 @@ fn chunk_cast_len(bytes: usize, addr: usize, elem_size: usize, align: usize) -> Some(bytes / elem_size) } +/// A borrowed contiguous typed view sharing its parent's allocation and access tracker. +pub struct DeviceBufferView<'a, T> { + storage: crate::RawDeviceBuffer, + shared_access_writes: bool, + len: usize, + _borrow: PhantomData<&'a DeviceBuffer>, +} +impl DeviceBufferView<'_, T> { + pub fn cu_deviceptr(&self) -> CUdeviceptr { + self.storage.cu_deviceptr() + } + pub fn len(&self) -> usize { + self.len + } + pub fn is_empty(&self) -> bool { + self.len == 0 + } + pub fn num_bytes(&self) -> usize { + self.storage.len_bytes() + } + pub fn context(&self) -> &Arc { + self.storage.context() + } + pub fn storage(&self) -> &crate::RawDeviceBuffer { + &self.storage + } + /// Whether a shared kernel reference may mutate elements through interior + /// mutability. Representation wrappers and unclassified raw types are + /// conservative; ordinary DeviceCopy values can share read accesses. + #[doc(hidden)] + pub fn shared_access_writes(&self) -> bool { + self.shared_access_writes + } +} + +impl<'a, T: DeviceCopy> DeviceBufferView<'a, T> { + /// Borrows initialized elements from shared byte storage without a copy. + /// + /// # Safety + /// The range must contain valid `T` values. All aliases must preserve that + /// validity and use the same access tracker for every device access. + pub unsafe fn borrow_storage( + storage: &'a crate::RawDeviceBuffer, + offset: usize, + len: usize, + ) -> Result { + let storage = storage.slice_bytes(offset, allocation_size::(len)?)?; + if !storage.is_empty() + && !(storage.cu_deviceptr() as usize).is_multiple_of(std::mem::align_of::()) + { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + Ok(Self { + storage, + len, + shared_access_writes: false, + _borrow: PhantomData, + }) + } +} + #[cfg(test)] mod chunk_cast_tests { use super::chunk_cast_len; diff --git a/cuda-core/src/dtype.rs b/cuda-core/src/dtype.rs index 9929a00d28..07cf16a5df 100644 --- a/cuda-core/src/dtype.rs +++ b/cuda-core/src/dtype.rs @@ -126,6 +126,28 @@ impl f4e2m1fnx2 { #[allow(non_camel_case_types)] pub struct i4(pub u8); +// Every bit pattern of the underlying `u8` or `u32` is a valid value of each +// wrapper: the fields are public or reachable through a safe `from_bits`, and +// nothing is reserved. That makes them plain old data in bytemuck's sense, +// which is what lets byte storage be reinterpreted as one of them through +// `DeviceBuffer::cast_elem`, `DeviceBuffer::cast_chunks`, and +// `Tensor::reinterpret`, all bounded on `AnyBitPattern` (#1506). `bool` is +// a `DType` too but deliberately has no such impl: only `0` and `1` are +// valid, so a `bool` buffer or tensor can only come from memory written as +// `bool`. +macro_rules! impl_pod_newtype { + ($($ty:ty),+ $(,)?) => { + $( + // SAFETY: `#[repr(transparent)]` over an integer, so there is no + // padding, and every bit pattern is a valid value (see above). + unsafe impl bytemuck::Zeroable for $ty {} + unsafe impl bytemuck::Pod for $ty {} + )+ + }; +} + +impl_pod_newtype!(tf32, f8e4m3fn, f8e5m2, f8e8m0fnu, f8e5m3fnu, f4e2m1fn, f4e2m1fnx2, i4); + /// Runtime identifier for a `DType`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum DTypeId { diff --git a/cuda-core/src/simt/embedded.rs b/cuda-core/src/embedded.rs similarity index 96% rename from cuda-core/src/simt/embedded.rs rename to cuda-core/src/embedded.rs index c146f58753..1bb8565184 100644 --- a/cuda-core/src/simt/embedded.rs +++ b/cuda-core/src/embedded.rs @@ -5,7 +5,7 @@ //! Loading CUDA modules from embedded device artifact bundles. -use crate::{CudaContext, CudaModule, DriverError}; +use crate::{Context, DriverError, Module}; use oxide_artifacts::ArtifactError; pub use oxide_artifacts::{ ArtifactCompileOptions, ArtifactDebugPolicy, ArtifactPayloadKind, OwnedArtifactBundle, @@ -43,7 +43,7 @@ impl EmbeddedModule { self.bundle.payload(kind) } - pub fn load(&self, ctx: &Arc) -> Result, EmbeddedModuleError> { + pub fn load(&self, ctx: &Arc) -> Result, EmbeddedModuleError> { let image = loadable_payload(&self.bundle).expect("EmbeddedModule always has a loadable payload"); ctx.load_module_from_image(image) @@ -78,9 +78,9 @@ pub fn embedded_modules_from_current_exe() -> Result, Embedd } pub fn load_embedded_module( - ctx: &Arc, + ctx: &Arc, name: &str, -) -> Result, EmbeddedModuleError> { +) -> Result, EmbeddedModuleError> { let module = embedded_modules_from_current_exe()? .into_iter() .find(|module| module.name() == name) @@ -90,9 +90,7 @@ pub fn load_embedded_module( module.load(ctx) } -pub fn load_first_embedded_module( - ctx: &Arc, -) -> Result, EmbeddedModuleError> { +pub fn load_first_embedded_module(ctx: &Arc) -> Result, EmbeddedModuleError> { let module = embedded_modules_from_current_exe()? .into_iter() .next() diff --git a/cuda-core/src/event.rs b/cuda-core/src/event.rs new file mode 100644 index 0000000000..93de0bddc2 --- /dev/null +++ b/cuda-core/src/event.rs @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Event timing and configuration helpers. + +pub use crate::runtime::{Event, Event as CudaEvent}; +use crate::{Context, DriverError, IntoResult}; +use std::{mem::MaybeUninit, sync::Arc}; + +impl Context { + /// Creates a new CUDA event in this context. + /// + /// `flags` defaults to `CU_EVENT_DISABLE_TIMING` when `None`. Use + /// `Some(CU_EVENT_DEFAULT)` to enable timing queries via + /// [`Event::elapsed_ms`]. + pub fn new_event_with_flags( + self: &Arc, + flags: Option, + ) -> Result { + let flags = flags.unwrap_or(cuda_bindings::CUevent_flags_enum_CU_EVENT_DISABLE_TIMING); + self.bind_to_thread()?; + let mut cu_event = MaybeUninit::uninit(); + let cu_event = unsafe { + cuda_bindings::cuEventCreate(cu_event.as_mut_ptr(), flags).result()?; + cu_event.assume_init() + }; + Ok(Event { + cu_event, + context: self.clone(), + }) + } +} +impl Event { + /// Returns the raw `CUevent` handle. + pub fn cu_event(&self) -> cuda_bindings::CUevent { + self.cu_event + } + /// Returns the parent [`Context`]. + pub fn context(&self) -> &Arc { + &self.context + } + /// Returns the elapsed time in milliseconds between `self` (start) and + /// `end`. + /// + /// Both events are synchronized before querying. Both events must have + /// been created **without** `CU_EVENT_DISABLE_TIMING`; otherwise the + /// driver returns `CUDA_ERROR_INVALID_HANDLE`. + /// + /// `self` must have been recorded before `end` in wall-clock time. + pub fn elapsed_ms(&self, end: &Self) -> Result { + self.synchronize()?; + end.synchronize()?; + let mut ms: f32 = 0.0; + unsafe { + cuda_bindings::cu_event_elapsed_time(&mut ms as *mut _, self.cu_event, end.cu_event) + .result()?; + } + Ok(ms) + } +} diff --git a/cuda-core/src/simt/launch.rs b/cuda-core/src/launch.rs similarity index 99% rename from cuda-core/src/simt/launch.rs rename to cuda-core/src/launch.rs index ff1f37684b..a3a870ebc8 100644 --- a/cuda-core/src/simt/launch.rs +++ b/cuda-core/src/launch.rs @@ -10,7 +10,7 @@ //! proving that it matches a kernel is unsafe. Use a rank-preserving config and //! [`PreparedLaunch`] for the checked path. -use crate::{CudaFunction, CudaStream, DriverError}; +use crate::{DriverError, Function, Stream}; use std::error::Error; use std::fmt::{self, Display, Formatter}; use std::marker::PhantomData; @@ -867,7 +867,7 @@ impl From for LaunchContractError { /// value performs no contract query; it only compares the stream's context /// handle before the normal launch path binds that context. pub struct PreparedLaunch { - function: CudaFunction, + function: Function, config: LaunchConfig, _contract: PhantomData C>, } @@ -897,7 +897,7 @@ impl PreparedLaunch { /// relationship. #[doc(hidden)] pub unsafe fn __prepare( - function: CudaFunction, + function: Function, config: C::Config, ) -> Result { let raw = config.__raw(); @@ -1043,7 +1043,7 @@ impl PreparedLaunch { } /// Returns the prepared CUDA function. - pub fn function(&self) -> &CudaFunction { + pub fn function(&self) -> &Function { &self.function } @@ -1057,7 +1057,7 @@ impl PreparedLaunch { } /// Rejects a stream from a different context without making a driver call. - pub fn validate_stream(&self, stream: &CudaStream) -> Result<(), LaunchContractError> { + pub fn validate_stream(&self, stream: &Stream) -> Result<(), LaunchContractError> { let function_context = self.function.context(); let stream_context = stream.context(); if function_context.cu_ctx() == stream_context.cu_ctx() { diff --git a/cuda-core/src/simt/mod.rs b/cuda-core/src/launch_raw.rs similarity index 80% rename from cuda-core/src/simt/mod.rs rename to cuda-core/src/launch_raw.rs index 7f07b4c361..53c0b0f1fc 100644 --- a/cuda-core/src/simt/mod.rs +++ b/cuda-core/src/launch_raw.rs @@ -3,83 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! The cuda-oxide host surface. -//! -//! Everything under this module is the cuda-oxide host API, brought over -//! as-is for the shared host-crate migration rather than reconciled: -//! [`CudaContext`] lands next to this crate's own `Device`, and -//! [`LaunchConfig`] next to this crate's own `LaunchConfig`. The two surfaces -//! get reconciled after the repository migration, with both models in one -//! tree, and this module is deletable as a unit once that happens. -//! -//! Naming: everything cuda-oxide exposed at its crate root is re-exported -//! here, so `cuda_core::simt::X` is oxide's `cuda_core::X`. The crate root -//! additionally re-exports every name that does not collide with the -//! existing surface; the one collision is [`LaunchConfig`], which is -//! reachable only through this module. -//! -//! Differences from the cuda-oxide original, recorded for the -//! reconciliation: -//! -//! - `DeviceCopy` for the unstable `f16` primitive sits behind the -//! default-off `f16` cargo feature, since it is nightly-only; the -//! `half::f16` and `half::bf16` impls are always present. -//! - `#[derive(DeviceCopy)]` comes from the sibling `cuda-core-derive` -//! crate, extracted from cuda-oxide's `cuda-macros` so it is publishable. -//! - `vmm` and `peer` are copied like the rest; #202 remains the idiomatic -//! port of the VMM wrappers onto this crate's own runtime types. -//! - `init` and `launch_kernel` are not copied: the crate root already -//! exposes identical ones, re-exported below. -//! - `oxide-artifacts` comes from crates.io (cuda-oxide publishes it) and is -//! re-exported as [`artifacts`]. +//! Context-aware raw CUDA launch helpers. -pub mod context; -pub mod device_buffer; -pub mod embedded; -pub mod event; -pub mod launch; -pub mod memory; -pub mod module; -pub mod peer; -pub mod pinned_host_buffer; -pub mod stream; -pub mod vmm; +use crate::{launch_kernel, DriverError, Function, IntoResult, Stream}; -pub use context::{ContextLimit, CudaContext, StreamPriorityRange, SyncPolicy}; -/// `#[derive(DeviceCopy)]`, re-exported next to the trait so -/// `use cuda_core::DeviceCopy;` brings both into scope (serde pattern). -pub use cuda_core_derive::DeviceCopy; -pub use device_buffer::{DeviceBuffer, DeviceCopy}; -pub use embedded::{EmbeddedModule, EmbeddedModuleError}; -pub use event::CudaEvent; -pub use launch::{ - BlockRequirement, DeviceLaunchLimits, DynamicSharedMemoryRequirement, KernelLaunchConfig, - KernelLaunchContract, LaunchAxis, LaunchConfig, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D, - LaunchContractError, LaunchContractSpec, LaunchDimension, PreparedLaunch, -}; -pub use module::{ConstantHandle, CudaFunction, CudaModule}; -pub use pinned_host_buffer::PinnedHostBuffer; -pub use stream::CudaStream; - -// The remaining items of oxide's crate root are identical on both surfaces; -// they are re-exported so `cuda_core::simt` mirrors it completely. -pub use crate::error::{DriverError, IntoResult}; -pub use crate::init; -pub use crate::launch_kernel; -pub use cuda_bindings as sys; -/// Artifact-bundle metadata, from the published crate cuda-oxide releases. -pub use oxide_artifacts as artifacts; - -/// Launches a CUDA kernel on a specific [`CudaStream`], binding its context first. +/// Launches a CUDA kernel on a specific [`Stream`], binding its context first. /// /// This is the usual host-side helper for `cuda_launch!` and async launches. /// It ensures the stream's owning context is current before calling the raw /// [`launch_kernel`] entry point, so callers do not need to manually call -/// [`CudaContext::bind_to_thread`] before every launch. +/// [`crate::Context::bind_to_thread`] before every launch. /// /// Unlike [`launch_kernel`], this helper works with typed wrappers rather than /// raw driver handles. It is therefore the preferred API whenever you already -/// have a [`CudaFunction`] and [`CudaStream`]. +/// have a [`Function`] and [`Stream`]. /// /// # Safety /// @@ -100,11 +37,11 @@ pub use oxide_artifacts as artifacts; /// `cuLaunchKernel` call rejects the launch. #[inline] pub unsafe fn launch_kernel_on_stream( - func: &CudaFunction, + func: &Function, grid_dim: (u32, u32, u32), block_dim: (u32, u32, u32), shared_mem_bytes: u32, - stream: &CudaStream, + stream: &Stream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; @@ -218,12 +155,12 @@ pub unsafe fn launch_kernel_ex( /// `cuLaunchKernelEx` call rejects the launch. #[inline] pub unsafe fn launch_kernel_ex_on_stream( - func: &CudaFunction, + func: &Function, grid_dim: (u32, u32, u32), block_dim: (u32, u32, u32), shared_mem_bytes: u32, cluster_dim: (u32, u32, u32), - stream: &CudaStream, + stream: &Stream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; @@ -330,11 +267,11 @@ pub unsafe fn launch_kernel_cooperative( /// `cuLaunchKernelEx` call rejects the launch. #[inline] pub unsafe fn launch_kernel_cooperative_on_stream( - func: &CudaFunction, + func: &Function, grid_dim: (u32, u32, u32), block_dim: (u32, u32, u32), shared_mem_bytes: u32, - stream: &CudaStream, + stream: &Stream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; @@ -447,12 +384,12 @@ pub unsafe fn launch_kernel_ex_cooperative( /// `cuLaunchKernelEx` call rejects the launch. #[inline] pub unsafe fn launch_kernel_ex_cooperative_on_stream( - func: &CudaFunction, + func: &Function, grid_dim: (u32, u32, u32), block_dim: (u32, u32, u32), shared_mem_bytes: u32, cluster_dim: (u32, u32, u32), - stream: &CudaStream, + stream: &Stream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; diff --git a/cuda-core/src/lib.rs b/cuda-core/src/lib.rs index ce2a802e5c..9d48e11132 100644 --- a/cuda-core/src/lib.rs +++ b/cuda-core/src/lib.rs @@ -7,33 +7,50 @@ #![cfg_attr(feature = "f16", feature(f16))] +pub mod allocation; mod api; +pub use allocation::{AllocationLease, DeviceAllocation, RawDeviceBuffer}; +pub mod completion; +pub mod context; pub(crate) mod cudarc_shim; +pub mod device_buffer; mod dtype; +pub mod embedded; mod error; +pub mod event; +pub mod launch; +mod launch_raw; +pub mod memory; +pub mod module; +pub mod peer; +pub mod pinned_host_buffer; mod runtime; -pub mod simt; +pub mod stream; pub mod vmm; pub use api::*; +// `DeviceBuffer::cast_elem` and `cast_chunks` are bounded on bytemuck's +// marker traits. The re-export lets a downstream crate +// implement them for its own element types against the same bytemuck. +pub use bytemuck; pub use cuda_bindings as sys; +pub use cuda_repr::DeviceRepr; pub use dtype::*; pub use error::*; pub use runtime::*; -// The cuda-oxide surface, re-exported at the root where its consumers -// expect it. `simt::LaunchConfig` and `simt::vmm` are deliberately absent: -// they collide with this crate's own `LaunchConfig` and with the reviewed -// `vmm` module above (both forks of the same cuda-oxide ancestor), and stay -// reachable only through `simt::`. -pub use simt::embedded; -pub use simt::{ - launch_kernel_cooperative, launch_kernel_cooperative_on_stream, launch_kernel_ex, - launch_kernel_ex_cooperative, launch_kernel_ex_cooperative_on_stream, - launch_kernel_ex_on_stream, launch_kernel_on_stream, BlockRequirement, ConstantHandle, - ContextLimit, CudaContext, CudaEvent, CudaFunction, CudaModule, CudaStream, DeviceBuffer, - DeviceCopy, DeviceLaunchLimits, DynamicSharedMemoryRequirement, EmbeddedModule, - EmbeddedModuleError, KernelLaunchConfig, KernelLaunchContract, LaunchAxis, LaunchConfig1D, - LaunchConfig2D, LaunchConfig3D, LaunchContractError, LaunchContractSpec, LaunchDimension, - PinnedHostBuffer, PreparedLaunch, StreamPriorityRange, SyncPolicy, +pub use context::{ContextLimit, StreamPriorityRange, SyncPolicy}; +pub use cuda_core_derive::DeviceCopy; +pub use device_buffer::{DeviceBuffer, DeviceBufferView, DeviceCopy}; +pub use embedded::{EmbeddedModule, EmbeddedModuleError}; +pub use launch::*; +pub use launch_raw::*; +pub use module::ConstantHandle; +pub use oxide_artifacts as artifacts; +pub use pinned_host_buffer::PinnedHostBuffer; + +// Compatibility spellings are aliases of the shared handles, not separate runtimes. +pub use runtime::{ + Context as CudaContext, Event as CudaEvent, Function as CudaFunction, Module as CudaModule, + Stream as CudaStream, }; diff --git a/cuda-core/src/simt/memory.rs b/cuda-core/src/memory.rs similarity index 100% rename from cuda-core/src/simt/memory.rs rename to cuda-core/src/memory.rs diff --git a/cuda-core/src/simt/module.rs b/cuda-core/src/module.rs similarity index 64% rename from cuda-core/src/simt/module.rs rename to cuda-core/src/module.rs index a2b4cb646f..952a8a5f88 100644 --- a/cuda-core/src/simt/module.rs +++ b/cuda-core/src/module.rs @@ -3,98 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! CUDA module and function management (RAII, PTX/cubin loading). -//! -//! A [`CudaModule`] wraps a `CUmodule` loaded from PTX source or a cubin file. -//! [`CudaFunction`] extracts a kernel entry point from a loaded module by -//! symbol name. Both types are reference-counted and tie their lifetime to the -//! parent [`CudaContext`] / [`CudaModule`] respectively. -//! -//! # Typical workflow -//! -//! ```ignore -//! let ctx = CudaContext::new(0)?; -//! let module = ctx.load_module_from_ptx_src(ptx)?; -//! let kernel = module.load_function("my_kernel")?; -//! ``` -//! -//! # Raw CUDA interop -//! -//! Most users should load kernels with [`CudaModule::load_function`] and -//! launch them through cuda-oxide's typed launch helpers. Some CUDA-adjacent -//! libraries need the underlying driver handle to inspect or register -//! module-scope device state before launch. For those cases, -//! [`CudaModule::cu_module`] exposes a non-owning raw `CUmodule` handle under -//! an explicit `unsafe` contract. - -use crate::error::{DriverError, IntoResult}; -use crate::simt::context::CudaContext; -use std::borrow::Cow; -use std::ffi::{c_void, CString}; -use std::mem::MaybeUninit; -use std::sync::Arc; - -/// An RAII wrapper around a `CUmodule` handle. -/// -/// Holds an `Arc` to ensure the context outlives the module. -/// Unloaded automatically via `cuModuleUnload` on [`Drop`]. -#[derive(Debug)] -pub struct CudaModule { - /// Raw CUDA module handle. - pub(crate) cu_module: cuda_bindings::CUmodule, - /// Owning context. Kept alive for the lifetime of this module. - pub(crate) ctx: Arc, -} - -/// # Safety -/// -/// `CUmodule` handles are not thread-local. The CUDA driver permits querying -/// functions from a module on any thread, provided the owning context is bound. -unsafe impl Send for CudaModule {} -/// See [`Send`] impl. -unsafe impl Sync for CudaModule {} - -/// Unloads the module on drop. -/// -/// Binds the context to the current thread first (required by -/// `cuModuleUnload`). Errors are recorded on the context rather than -/// panicking. -impl Drop for CudaModule { - fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); - self.ctx - .record_err(unsafe { cuda_bindings::cuModuleUnload(self.cu_module).result() }); - } -} - -impl CudaContext { - /// JIT-compiles PTX source and loads the resulting module into this - /// context. - /// - /// `ptx_src` must be a valid, null-terminator-free PTX string. The driver - /// performs JIT compilation targeting the current device architecture. - /// - /// # Panics - /// - /// Panics if `ptx_src` contains interior null bytes. - pub fn load_module_from_ptx_src( - self: &Arc, - ptx_src: &str, - ) -> Result, DriverError> { - self.bind_to_thread()?; - let c_src = CString::new(ptx_src).unwrap(); - let cu_module = unsafe { - let mut cu_module = MaybeUninit::uninit(); - cuda_bindings::cuModuleLoadData(cu_module.as_mut_ptr(), c_src.as_ptr() as *const _) - .result()?; - cu_module.assume_init() - }; - Ok(Arc::new(CudaModule { - cu_module, - ctx: self.clone(), - })) - } - +//! Module globals and function resource queries. + +pub use crate::runtime::{Function, Module}; +use crate::{Context, DriverError, IntoResult}; +use std::{ + borrow::Cow, + ffi::{c_void, CString}, + mem::MaybeUninit, + sync::Arc, +}; + +impl Context { /// Loads a CUDA module from an in-memory image. /// /// `image` may be PTX source bytes, a cubin, or a fatbin. PTX text is @@ -104,7 +24,7 @@ impl CudaContext { pub fn load_module_from_image( self: &Arc, image: &[u8], - ) -> Result, DriverError> { + ) -> Result, DriverError> { self.bind_to_thread()?; let image = null_terminated_image(image); let cu_module = unsafe { @@ -113,151 +33,24 @@ impl CudaContext { .result()?; cu_module.assume_init() }; - Ok(Arc::new(CudaModule { + Ok(Arc::new(Module { cu_module, - ctx: self.clone(), - })) - } - - /// Loads a module from a cubin or PTX file on disk. - /// - /// `filename` is the filesystem path. The driver selects the loader based - /// on file contents (PTX text or cubin ELF). - /// - /// # Panics - /// - /// Panics if `filename` contains interior null bytes. - pub fn load_module_from_file( - self: &Arc, - filename: &str, - ) -> Result, DriverError> { - self.bind_to_thread()?; - let c_str = CString::new(filename).unwrap(); - let mut cu_module = MaybeUninit::uninit(); - let cu_module = unsafe { - cuda_bindings::cuModuleLoad(cu_module.as_mut_ptr(), c_str.as_ptr()).result()?; - cu_module.assume_init() - }; - Ok(Arc::new(CudaModule { - cu_module, - ctx: self.clone(), + context: self.clone(), + owned: true, })) } } - -fn null_terminated_image(image: &[u8]) -> Cow<'_, [u8]> { - if image.last() == Some(&0) { - Cow::Borrowed(image) - } else { - let mut owned = Vec::with_capacity(image.len() + 1); - owned.extend_from_slice(image); - owned.push(0); - Cow::Owned(owned) - } -} - -/// A handle to a device kernel entry point within a loaded [`CudaModule`]. -/// -/// Holds an `Arc` so the module (and transitively the context) -/// remains loaded for the lifetime of this handle. Cloning is cheap (just an -/// `Arc` bump). -#[derive(Debug, Clone)] -pub struct CudaFunction { - /// Raw CUDA function handle. - pub(crate) cu_function: cuda_bindings::CUfunction, - /// Owning module. Prevents unloading while this function handle exists. - #[allow(unused)] - pub(crate) module: Arc, -} - -/// # Safety -/// -/// `CUfunction` handles are derived from a `CUmodule` and valid in any thread -/// that has the owning context bound. -unsafe impl Send for CudaFunction {} -/// See [`Send`] impl. -unsafe impl Sync for CudaFunction {} - -impl CudaModule { - /// Returns the parent [`CudaContext`]. +impl Module { + /// Returns the parent [`Context`]. /// /// This is mainly useful when interoperating with raw CUDA driver APIs: - /// call [`CudaContext::bind_to_thread`] on this context before passing raw + /// call [`Context::bind_to_thread`] on this context before passing raw /// module/function handles to APIs that require the owning context to be /// current on the calling host thread. - pub fn context(&self) -> &Arc { - &self.ctx - } - - /// Returns the raw `CUmodule` handle owned by this wrapper. - /// - /// This is an escape hatch for CUDA driver interop libraries that need to - /// inspect or register a loaded module directly. For example, NVSHMEM uses - /// module-level state and may need the raw `CUmodule` before launching a - /// kernel that calls into NVSHMEM device code. - /// - /// The returned handle is copied by value, but it is non-owning. - /// cuda-oxide still owns the module and will unload it when the last - /// [`Arc`] owning this module is dropped. - /// - /// # Safety - /// - /// - The returned handle is valid only while this [`CudaModule`] remains - /// alive. If the handle is stored outside the immediate call, the caller - /// must keep an [`Arc`] to this module alive for at least as long. - /// - The caller must not unload the module through the raw handle, transfer - /// ownership of it, or pass it to any API that may invalidate module, - /// function, or global handles owned by cuda-oxide. - /// - Before passing the handle to CUDA driver APIs or interop libraries - /// that make driver calls, the caller must ensure this module's owning - /// context is current on the calling host thread, for example with - /// [`CudaModule::context`] followed by - /// [`CudaContext::bind_to_thread`]. - /// - Any foreign library that retains this handle must obey the same - /// lifetime and context-current requirements. cuda-oxide cannot enforce - /// those requirements once the raw handle leaves Rust's type system. - pub unsafe fn cu_module(&self) -> cuda_bindings::CUmodule { - self.cu_module - } - - /// Looks up a kernel entry point by `fn_name` in this module. - /// - /// The returned [`CudaFunction`] holds an `Arc` back to this module, - /// preventing unloading while the handle is live. - /// - /// This method first binds the module's owning context to the calling - /// thread, then performs `cuModuleGetFunction`. That makes it safe to look - /// up functions from any host thread, provided the module and its context - /// are still alive. - /// - /// # Errors - /// - /// Returns an error if binding the module's context fails or if - /// `cuModuleGetFunction` cannot resolve `fn_name` in this module. - /// - /// # Panics - /// - /// Panics if `fn_name` contains interior null bytes. - pub fn load_function(self: &Arc, fn_name: &str) -> Result { - self.ctx.bind_to_thread()?; - let c_name = CString::new(fn_name).unwrap(); - let cu_function = unsafe { - let mut cu_function = MaybeUninit::uninit(); - cuda_bindings::cuModuleGetFunction( - cu_function.as_mut_ptr(), - self.cu_module, - c_name.as_ptr(), - ) - .result()?; - cu_function.assume_init() - }; - Ok(CudaFunction { - cu_function, - module: self.clone(), - }) + pub fn context(&self) -> &Arc { + &self.context } } - /// A resolved handle to a `#[constant]` device global. Macro-generated /// `set_` methods resolve these lazily on first use and cache the /// handle on the `LoadedModule` struct. Callers pass `size_of::()` on @@ -270,7 +63,7 @@ pub struct ConstantHandle { impl ConstantHandle { /// Construct from a raw device pointer. Used by macro-generated - /// `LoadedModule` initializers after [`CudaModule::get_global`] has + /// `LoadedModule` initializers after [`Module::get_global`] has /// resolved the symbol and the size has been asserted against /// `size_of::()`. /// @@ -293,14 +86,12 @@ impl ConstantHandle { /// - The bytes must have a layout compatible with the device-side type. pub unsafe fn write_async( &self, - stream: &crate::CudaStream, + stream: &crate::Stream, src: *const u8, num_bytes: usize, ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; - unsafe { - crate::simt::memory::memcpy_htod_async(self.dptr, src, num_bytes, stream.cu_stream()) - } + unsafe { crate::memory::memcpy_htod_async(self.dptr, src, num_bytes, stream.cu_stream()) } } /// Stream-ordered `cuMemcpyHtoDAsync` from owned host bytes into the @@ -311,7 +102,7 @@ impl ConstantHandle { /// caller passes a temporary such as `&3.0`. pub fn write_async_staged( &self, - stream: &crate::CudaStream, + stream: &crate::Stream, bytes: Box<[MaybeUninit]>, ) -> Result<(), DriverError> { stream.context().bind_to_thread()?; @@ -321,7 +112,7 @@ impl ConstantHandle { } unsafe { - crate::simt::memory::memcpy_htod_async( + crate::memory::memcpy_htod_async( self.dptr, bytes.as_ptr() as *const u8, num_bytes, @@ -370,7 +161,7 @@ impl ConstantHandle { /// Same contract as [`write_async`](Self::write_async). pub unsafe fn write_blocking( &self, - module: &Arc, + module: &Arc, src: *const u8, num_bytes: usize, ) -> Result<(), DriverError> { @@ -378,7 +169,7 @@ impl ConstantHandle { } } -impl CudaModule { +impl Module { /// Resolves a device global by name and returns its device pointer and /// size in bytes. /// @@ -402,7 +193,7 @@ impl CudaModule { self: &Arc, name: &str, ) -> Result<(cuda_bindings::CUdeviceptr, usize), DriverError> { - self.ctx.bind_to_thread()?; + self.context.bind_to_thread()?; let c_name = CString::new(name).unwrap(); let mut dptr = MaybeUninit::::uninit(); let mut size = MaybeUninit::::uninit(); @@ -439,12 +230,12 @@ impl CudaModule { src: *const u8, num_bytes: usize, ) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { crate::simt::memory::memcpy_htod_sync(dptr, src, num_bytes) } + self.context.bind_to_thread()?; + unsafe { crate::memory::memcpy_htod_sync(dptr, src, num_bytes) } } } -impl CudaFunction { +impl Function { fn attribute( &self, attribute: cuda_bindings::CUfunction_attribute, @@ -460,7 +251,7 @@ impl CudaFunction { } /// Returns the context that owns this function. - pub fn context(&self) -> &Arc { + pub fn context(&self) -> &Arc { self.module.context() } @@ -497,7 +288,7 @@ impl CudaFunction { /// bytes. /// /// This is the per-thread quantity that - /// [`CudaContext::set_stack_size`](crate::simt::context::CudaContext::set_stack_size) + /// [`Context::set_stack_size`](crate::context::Context::set_stack_size) /// budgets for: the driver reserves the stack limit for every resident /// thread on the device, so a kernel with a large frame can reserve /// device memory far in excess of anything it allocates explicitly. @@ -681,20 +472,15 @@ impl CudaFunction { .map_err(|_| DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE)) } } +} - /// Returns the raw `CUfunction` handle. - /// - /// # Safety - /// - /// The returned handle is copied by value, but it is non-owning. It is - /// invalidated if the parent [`CudaModule`] is dropped. - /// - /// Because [`CudaFunction`] holds an [`Arc`] to its parent module, the - /// module cannot be unloaded while `self` is alive. If the raw handle is - /// stored outside the immediate call, the caller must keep this - /// [`CudaFunction`] or another [`Arc`] owning the parent module alive for - /// at least as long as the raw handle is used. - pub unsafe fn cu_function(&self) -> cuda_bindings::CUfunction { - self.cu_function +fn null_terminated_image(image: &[u8]) -> Cow<'_, [u8]> { + if image.last() == Some(&0) { + Cow::Borrowed(image) + } else { + let mut owned = Vec::with_capacity(image.len() + 1); + owned.extend_from_slice(image); + owned.push(0); + Cow::Owned(owned) } } diff --git a/cuda-core/src/simt/peer.rs b/cuda-core/src/peer.rs similarity index 89% rename from cuda-core/src/simt/peer.rs rename to cuda-core/src/peer.rs index 4b67f79337..076207697e 100644 --- a/cuda-core/src/simt/peer.rs +++ b/cuda-core/src/peer.rs @@ -12,8 +12,8 @@ //! P2P access depends on hardware topology. Use `can_access_peer` to query //! support before calling `enable_peer_access`. +use crate::context::Context; use crate::error::{DriverError, IntoResult}; -use crate::simt::context::CudaContext; use std::mem::MaybeUninit; use std::sync::Arc; @@ -25,7 +25,7 @@ use std::sync::Arc; /// /// This is a query only -- it does not enable access. Call `enable_peer_access` /// to actually enable it. -pub fn can_access_peer(from: &CudaContext, to: &CudaContext) -> Result { +pub fn can_access_peer(from: &Context, to: &Context) -> Result { let mut can_access = MaybeUninit::uninit(); unsafe { cuda_bindings::cuDeviceCanAccessPeer( @@ -49,10 +49,7 @@ pub fn can_access_peer(from: &CudaContext, to: &CudaContext) -> Result, - to: &Arc, -) -> Result<(), DriverError> { +pub fn enable_peer_access(from: &Arc, to: &Arc) -> Result<(), DriverError> { from.bind_to_thread()?; let result = unsafe { cuda_bindings::cuCtxEnablePeerAccess(to.cu_ctx(), 0) }; match result { @@ -67,10 +64,7 @@ pub fn enable_peer_access( /// After this call, direct memory access from `from` to `to` is no longer /// possible. Any in-flight operations accessing peer memory must have completed /// before calling this. -pub fn disable_peer_access( - from: &Arc, - to: &Arc, -) -> Result<(), DriverError> { +pub fn disable_peer_access(from: &Arc, to: &Arc) -> Result<(), DriverError> { from.bind_to_thread()?; let result = unsafe { cuda_bindings::cuCtxDisablePeerAccess(to.cu_ctx()) }; match result { diff --git a/cuda-core/src/simt/pinned_host_buffer.rs b/cuda-core/src/pinned_host_buffer.rs similarity index 59% rename from cuda-core/src/simt/pinned_host_buffer.rs rename to cuda-core/src/pinned_host_buffer.rs index f39d823d82..5ed2382fb1 100644 --- a/cuda-core/src/simt/pinned_host_buffer.rs +++ b/cuda-core/src/pinned_host_buffer.rs @@ -18,29 +18,33 @@ use std::ptr::NonNull; use std::slice; use std::sync::Arc; +use crate::context::Context; +use crate::device_buffer::DeviceCopy; use crate::error::DriverError; -use crate::simt::context::CudaContext; -use crate::simt::device_buffer::DeviceCopy; /// Owned page-locked host buffer. /// /// The buffer is initialized and exposes normal Rust slices. The backing -/// allocation is freed with `cuMemFreeHost` on drop. +/// allocation is freed with `cuMemFreeHost` on drop. Extra space is reserved +/// when needed to align the elements, even for types whose alignment exceeds +/// CUDA's host-allocation alignment. /// /// Pinned transfer buffers require `T: DeviceCopy`, so host-owned values such /// as [`String`] are rejected. /// /// ```compile_fail -/// # use cuda_core::{CudaContext, PinnedHostBuffer}; -/// # fn rejects_non_device_copy(ctx: &std::sync::Arc) { +/// # use cuda_core::{Context, PinnedHostBuffer}; +/// # fn rejects_non_device_copy(ctx: &std::sync::Arc) { /// let _ = PinnedHostBuffer::::zeroed(ctx, 1); /// # } /// ``` pub struct PinnedHostBuffer { ptr: NonNull, + // CUDA must free the original allocation, not the aligned interior pointer. + allocation: Option>, len: usize, num_bytes: usize, - ctx: Arc, + ctx: Arc, _marker: PhantomData, } @@ -52,7 +56,7 @@ unsafe impl Sync for PinnedHostBuffer {} impl PinnedHostBuffer { /// Allocates a pinned host buffer and fills it with zero bytes. - pub fn zeroed(ctx: &Arc, len: usize) -> Result { + pub fn zeroed(ctx: &Arc, len: usize) -> Result { let mut buffer = Self::allocate(ctx, len)?; if buffer.num_bytes != 0 { unsafe { @@ -64,7 +68,7 @@ impl PinnedHostBuffer { } /// Allocates a pinned host buffer and copies `data` into it. - pub fn from_slice(ctx: &Arc, data: &[T]) -> Result { + pub fn from_slice(ctx: &Arc, data: &[T]) -> Result { let buffer = Self::allocate(ctx, data.len())?; if !data.is_empty() { unsafe { @@ -94,7 +98,7 @@ impl PinnedHostBuffer { /// Returns the CUDA context used to allocate this buffer. #[inline] - pub fn context(&self) -> &Arc { + pub fn context(&self) -> &Arc { &self.ctx } @@ -122,20 +126,29 @@ impl PinnedHostBuffer { unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) } } - fn allocate(ctx: &Arc, len: usize) -> Result { - let num_bytes = allocation_size::(len)?; - let ptr = if num_bytes == 0 { - NonNull::dangling() + fn allocate(ctx: &Arc, len: usize) -> Result { + let (num_bytes, allocation_bytes) = allocation_sizes::(len)?; + let (ptr, allocation) = if num_bytes == 0 { + (NonNull::dangling(), None) } else { ctx.bind_to_thread()?; - let ptr = unsafe { crate::simt::memory::malloc_host(num_bytes)? }; - NonNull::new(ptr.cast::()).ok_or(DriverError( + let raw = unsafe { crate::memory::malloc_host(allocation_bytes)? }; + let allocation = NonNull::new(raw.cast::()).ok_or(DriverError( cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, - ))? + ))?; + let offset = allocation.as_ptr().align_offset(std::mem::align_of::()); + // SAFETY: for a byte pointer, align_offset is at most align_of::() + // - 1. allocation_sizes reserves that much slack, leaving num_bytes + // accessible after this aligned, non-null interior pointer. The + // entire allocation fits in isize::MAX for pointer arithmetic. + let ptr = + unsafe { NonNull::new_unchecked(allocation.as_ptr().add(offset).cast::()) }; + (ptr, Some(allocation)) }; Ok(Self { ptr, + allocation, len, num_bytes, ctx: ctx.clone(), @@ -146,10 +159,10 @@ impl PinnedHostBuffer { impl Drop for PinnedHostBuffer { fn drop(&mut self) { - if self.num_bytes != 0 { + if let Some(allocation) = self.allocation { self.ctx.record_err(self.ctx.bind_to_thread()); self.ctx - .record_err(unsafe { crate::simt::memory::free_host(self.ptr.as_ptr().cast()) }); + .record_err(unsafe { crate::memory::free_host(allocation.as_ptr().cast()) }); } } } @@ -191,8 +204,39 @@ impl DerefMut for PinnedHostBuffer { } } -fn allocation_size(len: usize) -> Result { - len.checked_mul(std::mem::size_of::()).ok_or(DriverError( - cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, - )) +fn allocation_sizes(len: usize) -> Result<(usize, usize), DriverError> { + let invalid_size = DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + // Rust slices and pointer arithmetic require sizes bounded by isize::MAX, + // even when len * size_of::() would fit in usize. + let layout = std::alloc::Layout::array::(len).map_err(|_| invalid_size)?; + let num_bytes = layout.size(); + if num_bytes == 0 { + return Ok((0, 0)); + } + let allocation_bytes = num_bytes + .checked_add(layout.align() - 1) + .filter(|&size| size <= isize::MAX as usize) + .ok_or(invalid_size)?; + Ok((num_bytes, allocation_bytes)) +} + +#[cfg(test)] +mod tests { + use super::allocation_sizes; + + #[test] + fn rejects_sizes_that_cannot_form_a_rust_slice() { + assert!(allocation_sizes::(isize::MAX as usize + 1).is_err()); + assert!(allocation_sizes::(usize::MAX).is_err()); + assert_eq!( + allocation_sizes::(isize::MAX as usize).unwrap(), + (isize::MAX as usize, isize::MAX as usize) + ); + } + + #[test] + fn empty_and_zero_sized_buffers_need_no_allocation() { + assert_eq!(allocation_sizes::(0).unwrap(), (0, 0)); + assert_eq!(allocation_sizes::<()>(usize::MAX).unwrap(), (0, 0)); + } } diff --git a/cuda-core/src/runtime.rs b/cuda-core/src/runtime.rs index b1893e92b5..07200048cf 100644 --- a/cuda-core/src/runtime.rs +++ b/cuda-core/src/runtime.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! CUDA runtime types: Device, Stream, Module, Function, LaunchConfig. +//! CUDA runtime types: Context, Stream, Module, Function, LaunchConfig. //! //! These are the public API of `cuda-core`. They wrap raw CUDA driver //! handles with RAII lifetimes and provide `borrow_raw` constructors @@ -16,17 +16,6 @@ use crate::cudarc_shim::{ctx, device, module, pool, primary_ctx, stream}; use crate::error::*; use crate::init; -/// Kernel launch configuration specifying grid, block, and shared memory sizes. -#[derive(Clone, Copy, Debug)] -pub struct LaunchConfig { - /// Grid dimensions `(x, y, z)` in thread blocks. - pub grid_dim: (u32, u32, u32), - /// Block dimensions `(x, y, z)` in threads. - pub block_dim: (u32, u32, u32), - /// Bytes of dynamic shared memory per block. - pub shared_mem_bytes: u32, -} - /// Anything that owns an external CUDA resource. A borrowed handle can hold an /// `Arc` as a *liveness token*: while the handle (and anything /// derived from it) is alive, the token's refcount is nonzero, so the external @@ -72,38 +61,88 @@ impl PartialEq for KeepAlive { } impl Eq for KeepAlive {} -/// A GPU device handle wrapping a CUDA primary context. +/// A physical CUDA device identity, independent of any CUDA context. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Device { + cu_device: cuda_bindings::CUdevice, + ordinal: usize, +} + +impl Device { + /// Looks up a device without retaining its primary context. + pub fn new(ordinal: usize) -> Result { + let raw = ordinal_to_c_int(ordinal)?; + unsafe { init(0)? }; + Ok(Self { + cu_device: device::get(raw)?, + ordinal, + }) + } + /// Returns the physical device ordinal. + pub fn ordinal(&self) -> usize { + self.ordinal + } + /// Returns the raw physical device handle. + pub fn cu_device(&self) -> cuda_bindings::CUdevice { + self.cu_device + } + /// Returns the device name. + pub fn name(&self) -> Result { + device::get_name(self.cu_device) + } + /// Retains this device's primary context for resource creation and submission. + pub fn primary_context(&self) -> Result, DriverError> { + Context::new(self.ordinal) + } + /// Returns the number of CUDA devices. + pub fn device_count() -> Result { + Context::device_count() + } +} + +impl Context { + /// The physical device hosting this context. Equal devices do not imply equal contexts. + pub fn device(&self) -> Device { + Device { + cu_device: self.cu_device, + ordinal: self.ordinal, + } + } +} + +/// A CUDA execution context, distinct from its physical [`Device`]. /// -/// Can be **owned** (created via [`Device::new`], releases the primary context -/// on drop), **borrowed** (created via [`Device::borrow_raw`], does NOT release -/// on drop), or **foreign** (created via [`Device::borrow_with_owner`], holds a +/// Can be **owned** (created via [`Context::new`], releases the primary context +/// on drop), **borrowed** (created via [`Context::borrow_raw`], does NOT release +/// on drop), or **foreign** (created via [`Context::borrow_with_owner`], holds a /// liveness token so the external owner outlives it). #[derive(Debug)] -pub struct Device { +pub struct Context { pub(crate) cu_device: cuda_bindings::CUdevice, pub(crate) cu_ctx: cuda_bindings::CUcontext, pub(crate) ordinal: usize, - owned: bool, - _keep_alive: KeepAlive, + pub(crate) owned: bool, + pub(crate) _keep_alive: KeepAlive, + pub(crate) error_state: std::sync::atomic::AtomicU32, } -unsafe impl Send for Device {} -unsafe impl Sync for Device {} +unsafe impl Send for Context {} +unsafe impl Sync for Context {} -impl Drop for Device { +impl Drop for Context { fn drop(&mut self) { if !self.owned { return; } let _guard = teardown_lock(); - // Streams hold an Arc, so by the time the last device handle + // Streams hold an Arc, so by the time the last device handle // for this ordinal drops, every pooled handle is idle; the context // release below reclaims them. Discard so a later re-retain cannot // pop a handle from a torn-down context. stream_pool() .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .remove(&self.ordinal); + .remove(&(self.cu_ctx as usize)); let _ = self.bind_to_thread(); let ctx = std::mem::replace(&mut self.cu_ctx, std::ptr::null_mut()); if !ctx.is_null() { @@ -112,14 +151,14 @@ impl Drop for Device { } } -impl PartialEq for Device { +impl PartialEq for Context { fn eq(&self, other: &Self) -> bool { self.cu_device == other.cu_device && self.cu_ctx == other.cu_ctx && self.ordinal == other.ordinal } } -impl Eq for Device {} +impl Eq for Context {} /// The driver indexes devices with a C `int`. An ordinal that does not fit is /// not a device at all; `as c_int` would wrap it onto some other device's @@ -129,7 +168,7 @@ fn ordinal_to_c_int(ordinal: usize) -> Result { .map_err(|_| DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_DEVICE)) } -impl Device { +impl Context { /// Creates a new owned device on the specified ordinal. /// /// Errors with `CUDA_ERROR_INVALID_DEVICE` (before touching the driver) @@ -139,15 +178,16 @@ impl Device { unsafe { init(0)? }; let cu_device = device::get(cu_ordinal)?; let cu_ctx = unsafe { primary_ctx::retain(cu_device) }?; - let device = Arc::new(Device { + let context = Arc::new(Context { cu_device, cu_ctx, ordinal, owned: true, _keep_alive: KeepAlive::none(), + error_state: std::sync::atomic::AtomicU32::new(0), }); - device.bind_to_thread()?; - Ok(device) + context.bind_to_thread()?; + Ok(context) } /// Wraps externally-owned CUDA handles without taking ownership. @@ -162,7 +202,7 @@ impl Device { /// /// The caller must ensure: /// - `cu_ctx` points to a valid retained `CUcontext` for `cu_device` - /// - The handles outlive the returned `Device` **and everything derived + /// - The handles outlive the returned `Context` **and everything derived /// from it**. Streams from [`new_stream`](Self::new_stream) hold the /// device alive and, because the device is borrowed, are never parked in /// the process-wide stream pool: each is synchronized and destroyed @@ -170,12 +210,13 @@ impl Device { /// therefore still be valid when the last such stream drops. /// - No concurrent destruction of the handles pub unsafe fn borrow_raw(cu_ctx: *mut c_void, cu_device: c_int, ordinal: usize) -> Arc { - Arc::new(Device { + Arc::new(Context { cu_device: cu_device as cuda_bindings::CUdevice, cu_ctx: cu_ctx as cuda_bindings::CUcontext, ordinal, owned: false, _keep_alive: KeepAlive::none(), + error_state: std::sync::atomic::AtomicU32::new(0), }) } @@ -184,8 +225,8 @@ impl Device { /// /// Same as [`borrow_raw`](Self::borrow_raw), but the liveness obligation is /// discharged by construction: `owner` is whatever owns the context (a - /// cudarc device, a torch context handle, ...), and holding it here - /// guarantees the handles stay valid as long as this `Device` — or anything + /// cudarc context, a torch context handle, ...), and holding it here + /// guarantees the handles stay valid as long as this `Context` — or anything /// derived from it — lives. Only the point-in-time validity of the handles /// remains a caller assertion. /// @@ -200,12 +241,13 @@ impl Device { ordinal: usize, owner: Arc, ) -> Arc { - Arc::new(Device { + Arc::new(Context { cu_device: cu_device as cuda_bindings::CUdevice, cu_ctx: cu_ctx as cuda_bindings::CUcontext, ordinal, owned: false, _keep_alive: KeepAlive::owner(owner), + error_state: std::sync::atomic::AtomicU32::new(0), }) } @@ -216,7 +258,7 @@ impl Device { } /// Returns the raw `CUdevice` handle for a given ordinal without - /// creating a full `Device` (no context retained). Rejects an ordinal + /// creating a full `Context` (no context retained). Rejects an ordinal /// that does not fit `c_int` like [`new`](Self::new) does. pub fn raw_device(ordinal: usize) -> Result { let cu_ordinal = ordinal_to_c_int(ordinal)?; @@ -229,7 +271,7 @@ impl Device { self.ordinal } - /// Get the name of this device. + /// Get the name of this context. pub fn name(&self) -> Result { device::get_name(self.cu_device) } @@ -246,6 +288,7 @@ impl Device { /// Binds this context to the calling thread if not already current. pub fn bind_to_thread(&self) -> Result<(), DriverError> { + self.check_err()?; if match ctx::get_current()? { Some(curr_ctx) => curr_ctx != self.cu_ctx, None => true, @@ -257,18 +300,17 @@ impl Device { /// Blocks until all work on this device's context is complete. /// - /// # Safety - /// The caller must ensure this device's context is current on the - /// calling thread (via [`bind_to_thread`](Device::bind_to_thread)). - pub unsafe fn synchronize(&self) -> Result<(), DriverError> { + /// Binds this context to the calling thread before waiting. + pub fn synchronize(&self) -> Result<(), DriverError> { + self.bind_to_thread()?; ctx::synchronize() } - /// Creates a new non-blocking CUDA stream on this device. + /// Creates a new non-blocking CUDA stream on this context. /// /// On an owned device the handle may be one parked by an earlier stream's /// drop (see [`Stream`]). On a borrowed device it is always freshly - /// created: the pool is keyed by ordinal alone and holds handles from the + /// created: the pool is keyed by context identity and holds handles from the /// primary contexts this crate retains, which need not be the context the /// external owner handed to [`borrow_raw`](Self::borrow_raw). pub fn new_stream(self: &Arc) -> Result, DriverError> { @@ -277,7 +319,7 @@ impl Device { stream_pool() .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .get_mut(&self.ordinal) + .get_mut(&(self.cu_ctx as usize)) .and_then(Vec::pop) } else { None @@ -288,7 +330,7 @@ impl Device { }; Ok(Arc::new(Stream { cu_stream, - device: self.clone(), + context: self.clone(), owned: true, _keep_alive: KeepAlive::none(), })) @@ -306,7 +348,7 @@ impl Device { }?; Ok(Arc::new(Module { cu_module, - device: self.clone(), + context: self.clone(), owned: true, })) } @@ -320,7 +362,7 @@ impl Device { let cu_module = { module::load(filename) }?; Ok(Arc::new(Module { cu_module, - device: self.clone(), + context: self.clone(), owned: true, })) } @@ -364,12 +406,12 @@ impl Device { let cu_module = unsafe { module::load_data(image.as_ptr().cast()) }?; Ok(Arc::new(Module { cu_module, - device: self.clone(), + context: self.clone(), owned: true, })) } - /// Creates a new memory pool on this device. + /// Creates a new memory pool on this context. /// /// The returned pool is owned — it will be destroyed when the last `Arc` /// is dropped. Pair with [`MemPool::set_release_threshold`] to control @@ -384,21 +426,21 @@ impl Device { let cu_pool = unsafe { pool::create(&props) }?; Ok(Arc::new(MemPool { cu_pool, - device: self.clone(), + context: self.clone(), owned: true, })) } - /// Returns the driver-owned default memory pool for this device. + /// Returns the driver-owned default memory pool for this context. /// /// The returned wrapper is **not owned** — dropping it does not destroy the - /// default pool, which is shared across all users of the device. + /// default pool, which is shared across all users of the context. pub fn default_mem_pool(self: &Arc) -> Result, DriverError> { self.bind_to_thread()?; let cu_pool = unsafe { pool::get_default(self.cu_device) }?; Ok(Arc::new(MemPool { cu_pool, - device: self.clone(), + context: self.clone(), owned: false, })) } @@ -406,17 +448,17 @@ impl Device { /// A CUDA memory pool handle. /// -/// Can be either **owned** (created via [`Device::new_mem_pool`], destroyed on -/// drop) or **borrowed** (created via [`Device::default_mem_pool`], does NOT +/// Can be either **owned** (created via [`Context::new_mem_pool`], destroyed on +/// drop) or **borrowed** (created via [`Context::default_mem_pool`], does NOT /// destroy on drop). /// /// Used by async tensor allocation via `cuMemAllocFromPoolAsync` when a pool -/// is registered via `cuda_async::device_context::set_device_pool`. +/// is registered via `cuda_async::device_device::set_device_pool`. #[derive(Debug)] pub struct MemPool { pub(crate) cu_pool: cuda_bindings::CUmemoryPool, - pub(crate) device: Arc, - owned: bool, + pub(crate) context: Arc, + pub(crate) owned: bool, } unsafe impl Send for MemPool {} @@ -427,7 +469,7 @@ impl Drop for MemPool { if !self.owned { return; } - let _ = self.device.bind_to_thread(); + let _ = self.context.bind_to_thread(); let _ = unsafe { pool::destroy(self.cu_pool) }; } } @@ -438,9 +480,9 @@ impl MemPool { self.cu_pool } - /// Returns a reference to the parent device. - pub fn device(&self) -> &Arc { - &self.device + /// Returns a reference to the parent context. + pub fn context(&self) -> &Arc { + &self.context } /// Sets the release threshold for this pool. @@ -449,7 +491,7 @@ impl MemPool { /// below this threshold. Use `u64::MAX` to prevent the OS from reclaiming /// pool memory (useful for inference workloads with stable memory footprints). pub fn set_release_threshold(&self, threshold: u64) -> Result<(), DriverError> { - self.device.bind_to_thread()?; + self.context.bind_to_thread()?; unsafe { pool::set_release_threshold(self.cu_pool, threshold) } } @@ -460,7 +502,7 @@ impl MemPool { /// a synchronously consistent reading. For stable values, sync the streams /// that allocate from this pool first. pub fn mem_stats(&self) -> Result { - self.device.bind_to_thread()?; + self.context.bind_to_thread()?; unsafe { Ok(PoolMemStats { used_current: pool::get_attribute_u64( @@ -485,7 +527,7 @@ impl MemPool { /// Resets `used_high` to the current `used_current`. pub fn reset_used_high(&self) -> Result<(), DriverError> { - self.device.bind_to_thread()?; + self.context.bind_to_thread()?; unsafe { pool::reset_high_watermark( self.cu_pool, @@ -496,7 +538,7 @@ impl MemPool { /// Resets `reserved_high` to the current `reserved_current`. pub fn reset_reserved_high(&self) -> Result<(), DriverError> { - self.device.bind_to_thread()?; + self.context.bind_to_thread()?; unsafe { pool::reset_high_watermark( self.cu_pool, @@ -523,41 +565,32 @@ pub struct PoolMemStats { /// A CUDA stream handle. /// -/// Can be **owned** (created via [`Device::new_stream`]), **borrowed** (created +/// Can be **owned** (created via [`Context::new_stream`]), **borrowed** (created /// via [`Stream::borrow_raw`], does NOT destroy on drop), or **foreign** /// (created via [`Stream::borrow_with_owner`], holds a liveness token so the /// external owner outlives it). /// -/// Dropping an owned stream blocks in `cuStreamSynchronize` until the work -/// enqueued on it has drained. What happens to the handle next depends on the -/// device it was created on: -/// -/// - On an **owned** device the handle is *not* destroyed: it is parked in a -/// process-wide, per-ordinal pool and handed out again by the next -/// [`Device::new_stream`] on that ordinal (see `stream_pool` for why). The -/// last owned `Device` for the ordinal discards the pool entry, and the -/// primary-context release reclaims the handles. -/// - On a **borrowed** device ([`Device::borrow_raw`] / -/// [`Device::borrow_with_owner`]) the handle is destroyed, under the -/// teardown lock. It must never enter the pool: the pool is keyed by -/// ordinal only, so an owned `Device` created later for the same ordinal -/// would pop a handle belonging to a context the external owner may already -/// have destroyed. +/// Drop returns without waiting for queued GPU work. An event-backed retirement +/// queue retains the context and foreign owner until completion. Idle default +/// priority streams on owned contexts enter a pool keyed by CUDA context +/// identity; streams on borrowed contexts and priority streams are destroyed +/// after completion. Completion errors retain resources rather than reuse a +/// handle that may still be active. #[derive(Debug, PartialEq, Eq)] pub struct Stream { pub(crate) cu_stream: cuda_bindings::CUstream, - pub(crate) device: Arc, - owned: bool, - _keep_alive: KeepAlive, + pub(crate) context: Arc, + pub(crate) owned: bool, + pub(crate) _keep_alive: KeepAlive, } -/// Per-device pool of idle stream handles (all created `NonBlocking` by -/// [`Device::new_stream`]). `cuStreamDestroy` racing concurrent driver +/// Per-context pool of idle stream handles (all created `NonBlocking` by +/// [`Context::new_stream`]). `cuStreamDestroy` racing concurrent driver /// activity from other threads segfaults inside libcuda (reproducible in /// parallel test runs; a leak-streams experiment ran 10/10 clean where /// destroy-paths crashed ~half the runs), so owned streams are never /// destroyed: drops return the handle here, `new_stream` reuses it, and the -/// last `Device` drop for an ordinal discards the pooled handles — the +/// last `Context` drop for an ordinal discards the pooled handles — the /// primary-context release reclaims them. fn stream_pool() -> &'static std::sync::Mutex>> { static POOL: std::sync::OnceLock< @@ -590,8 +623,8 @@ unsafe impl Sync for Stream {} /// early is safe). An event is bound to its device at construction; /// recording it on a stream of another device is rejected. pub struct Event { - cu_event: cuda_bindings::CUevent, - device: Arc, + pub(crate) cu_event: cuda_bindings::CUevent, + pub(crate) context: Arc, } unsafe impl Send for Event {} @@ -602,20 +635,21 @@ impl Event { /// /// Errors with `CUDA_ERROR_INVALID_VALUE` if the stream belongs to a /// different device than the one this event was created on. - pub fn record(&self, stream: &Arc) -> Result<(), DriverError> { - if stream.device().ordinal() != self.device.ordinal() { + pub fn record(&self, stream: &Stream) -> Result<(), DriverError> { + if stream.context().as_ref() != self.context.as_ref() { return Err(DriverError( cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, )); } - // Safety: both handles are valid by construction (RAII wrappers), - // and the same-device check above pins them to one context. + self.context.bind_to_thread()?; + // SAFETY: the handles have equal context identity and that context is current. unsafe { crate::cudarc_shim::event::record(self.cu_event, stream.cu_stream()) } } /// Blocks the calling thread until this event has completed. pub fn synchronize(&self) -> Result<(), DriverError> { - // Safety: the handle is valid by construction. + self.context.bind_to_thread()?; + // SAFETY: the handle is valid and its context is current. unsafe { crate::cudarc_shim::event::synchronize(self.cu_event) } } @@ -624,7 +658,8 @@ impl Event { /// (or the event was never recorded), `Ok(false)` while it is still in /// flight. Any other driver result is returned as the error. pub fn query(&self) -> Result { - // Safety: the handle is valid by construction. + self.context.bind_to_thread()?; + // SAFETY: the handle is valid and its context is current. match unsafe { crate::cudarc_shim::event::query(self.cu_event) } { Ok(()) => Ok(true), Err(DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_READY)) => Ok(false), @@ -640,27 +675,34 @@ impl Event { /// call [`synchronize`](Self::synchronize) on it first, or the driver /// reports not-ready. pub fn elapsed_time(&self, end: &Event) -> Result { - // Safety: both handles are valid by construction. + if self.context.as_ref() != end.context.as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + self.context.bind_to_thread()?; + // SAFETY: both handles belong to the current context. unsafe { crate::cudarc_shim::event::elapsed(self.cu_event, end.cu_event) } } } impl Drop for Event { fn drop(&mut self) { + self.context.record_err(self.context.bind_to_thread()); // Safety: owned handle; the driver defers destruction while in use. let _ = unsafe { crate::cudarc_shim::event::destroy(self.cu_event) }; } } -impl Device { - /// Creates a timing-enabled event on this device. +impl Context { + /// Creates a timing-enabled event on this context. pub fn new_event(self: &Arc) -> Result { self.bind_to_thread()?; let cu_event = crate::cudarc_shim::event::create(cuda_bindings::CUevent_flags_enum_CU_EVENT_DEFAULT)?; Ok(Event { cu_event, - device: self.clone(), + context: self.clone(), }) } @@ -681,31 +723,58 @@ impl Device { } } -impl Drop for Stream { +// Retains both context and foreign liveness tokens while a stream drains. +struct RetiredStream { + handle: usize, + context: Arc, + owned: bool, + reusable: bool, + _keep_alive: KeepAlive, +} + +impl Drop for RetiredStream { fn drop(&mut self) { - if !self.owned || self.cu_stream.is_null() { + if !self.owned || self.handle == 0 { return; } - let _ = self.device.bind_to_thread(); - if self.device.owned { - // Never destroyed (see `stream_pool`): drain, then return the - // handle for reuse by the next `new_stream` on this device. - let _ = unsafe { stream::synchronize(self.cu_stream) }; + if self.reusable { stream_pool() .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .entry(self.device.ordinal) + .unwrap_or_else(|e| e.into_inner()) + .entry(self.context.cu_ctx() as usize) .or_default() - .push(self.cu_stream as usize); + .push(self.handle); } else { - // Borrowed device: the pool is keyed by ordinal only, so parking - // this handle would let an owned `Device` created later for the - // same ordinal pop a stream from a context the external owner may - // by then have destroyed (reproduced as a segfault). Destroy it - // instead, serialized with the other destructive driver calls. let _guard = teardown_lock(); - let _ = unsafe { stream::synchronize(self.cu_stream) }; - let _ = unsafe { stream::destroy(self.cu_stream) }; + if self.context.bind_to_thread().is_err() { + // Preserve the context rather than releasing a resource whose + // teardown the driver could not perform safely. + std::mem::forget(self.context.clone()); + return; + } + self.context + .record_err(unsafe { stream::destroy(self.handle as cuda_bindings::CUstream) }); + } + } +} + +impl Drop for Stream { + fn drop(&mut self) { + let retired = RetiredStream { + handle: self.cu_stream as usize, + context: self.context.clone(), + owned: self.owned, + // Preserve priority semantics: only default-priority streams enter + // the pool used by Context::new_stream. + reusable: self.context.owned && self.priority().ok() == Some(0), + _keep_alive: self._keep_alive.clone(), + }; + match self.query() { + Ok(true) => drop(retired), + Ok(false) => { + let _ = crate::completion::retain_until_complete(retired, self); + } + Err(_) => std::mem::forget(retired), } } } @@ -714,7 +783,7 @@ impl Stream { /// Wraps an externally-owned CUDA stream without taking ownership. /// /// `cu_stream` is the raw `CUstream` opaque pointer. See - /// [`Device::borrow_raw`] for why this is a `*mut c_void` rather than a + /// [`Context::borrow_raw`] for why this is a `*mut c_void` rather than a /// `cuda_bindings::CUstream` typedef. /// /// # Safety @@ -723,10 +792,10 @@ impl Stream { /// - `cu_stream` points to a valid CUDA stream on `device` /// - The stream outlives the returned `Stream` /// - No concurrent destruction of the stream - pub unsafe fn borrow_raw(cu_stream: *mut c_void, device: &Arc) -> Arc { + pub unsafe fn borrow_raw(cu_stream: *mut c_void, context: &Arc) -> Arc { Arc::new(Stream { cu_stream: cu_stream as cuda_bindings::CUstream, - device: device.clone(), + context: context.clone(), owned: false, _keep_alive: KeepAlive::none(), }) @@ -747,12 +816,12 @@ impl Stream { /// that stream (i.e. `owner` really is what keeps it alive). pub unsafe fn borrow_with_owner( cu_stream: *mut c_void, - device: &Arc, + context: &Arc, owner: Arc, ) -> Arc { Arc::new(Stream { cu_stream: cu_stream as cuda_bindings::CUstream, - device: device.clone(), + context: context.clone(), owned: false, _keep_alive: KeepAlive::owner(owner), }) @@ -763,28 +832,26 @@ impl Stream { self.cu_stream } - /// Returns a reference to the parent device. - pub fn device(&self) -> &Arc { - &self.device + /// Returns a reference to the parent context. + pub fn context(&self) -> &Arc { + &self.context } /// Blocks until all work on this stream is complete. /// - /// # Safety - /// The caller must ensure the parent device's context is current on - /// the calling thread. - pub unsafe fn synchronize(&self) -> Result<(), DriverError> { - stream::synchronize(self.cu_stream) + /// Binds the parent context to the calling thread before submitting. + pub fn synchronize(&self) -> Result<(), DriverError> { + self.context.bind_to_thread()?; + unsafe { stream::synchronize(self.cu_stream) } } /// Queries stream completion without blocking: `Ok(true)` when all prior /// work has completed, `Ok(false)` when work is still in flight. /// - /// # Safety - /// The caller must ensure the parent device's context is current on - /// the calling thread. - pub unsafe fn query(&self) -> Result { - stream::query(self.cu_stream) + /// Binds the parent context to the calling thread before submitting. + pub fn query(&self) -> Result { + self.context.bind_to_thread()?; + unsafe { stream::query(self.cu_stream) } } /// Makes all work subsequently enqueued on this stream wait until @@ -795,13 +862,13 @@ impl Stream { /// Errors with `CUDA_ERROR_INVALID_VALUE` if the event belongs to a /// different device than this stream. pub fn wait_event(&self, event: &Event) -> Result<(), DriverError> { - if event.device.ordinal() != self.device.ordinal() { + if event.context.as_ref() != self.context.as_ref() { return Err(DriverError( cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, )); } - // Safety: both handles are valid by construction (RAII wrappers), - // and the same-device check above pins them to one context. + self.context.bind_to_thread()?; + // SAFETY: the handles have equal context identity and that context is current. unsafe { stream::wait_event( self.cu_stream, @@ -831,13 +898,12 @@ impl Stream { /// callback fires; on `Err` the driver never invokes the callback, so the /// box is reclaimed here and the closure's captures are dropped. /// - /// # Safety - /// The caller must ensure the parent device's context is current on - /// the calling thread. - pub unsafe fn launch_host_function( + /// Binds the parent context to the calling thread before submitting. + pub fn launch_host_function( &self, host_func: F, ) -> Result<(), DriverError> { + self.context.bind_to_thread()?; Self::enqueue_boxed(host_func, |func, arg| unsafe { stream::launch_host_function(self.cu_stream, func, arg) }) @@ -849,14 +915,13 @@ impl Stream { /// rules: no CUDA calls inside `host_func`, panics are caught and /// discarded, and a refused launch reclaims the boxed closure. /// - /// # Safety - /// The caller must ensure the parent device's context is current on - /// the calling thread. - pub unsafe fn launch_host_function_with_sync_mode( + /// Binds the parent context to the calling thread before submitting. + pub fn launch_host_function_with_sync_mode( &self, host_func: F, sync_mode: ::core::ffi::c_uint, ) -> Result<(), DriverError> { + self.context.bind_to_thread()?; Self::enqueue_boxed(host_func, |func, arg| unsafe { stream::launch_host_function_v2(self.cu_stream, func, arg, sync_mode) }) @@ -922,14 +987,14 @@ impl Stream { /// A loaded CUDA module (PTX/cubin). /// -/// Can be either **owned** (created via [`Device::load_module_from_ptx_src`] -/// / [`Device::load_module_from_file`], unloads on drop) or **borrowed** +/// Can be either **owned** (created via [`Context::load_module_from_ptx_src`] +/// / [`Context::load_module_from_file`], unloads on drop) or **borrowed** /// (created via [`Module::borrow_raw`], does NOT unload on drop). #[derive(Debug)] pub struct Module { pub(crate) cu_module: cuda_bindings::CUmodule, - pub(crate) device: Arc, - owned: bool, + pub(crate) context: Arc, + pub(crate) owned: bool, } unsafe impl Send for Module {} @@ -941,7 +1006,7 @@ impl Drop for Module { return; } let _guard = teardown_lock(); - let _ = self.device.bind_to_thread(); + let _ = self.context.bind_to_thread(); let _ = unsafe { module::unload(self.cu_module) }; } } @@ -950,7 +1015,7 @@ impl Module { /// Wraps an externally-owned CUDA module without taking ownership. /// /// `cu_module` is the raw `CUmodule` opaque pointer. See - /// [`Device::borrow_raw`] for why this is a `*mut c_void`. + /// [`Context::borrow_raw`] for why this is a `*mut c_void`. /// /// # Safety /// @@ -958,10 +1023,10 @@ impl Module { /// - `cu_module` points to a valid module loaded on `device` /// - The module outlives the returned `Module` /// - No concurrent unload of the module - pub unsafe fn borrow_raw(cu_module: *mut c_void, device: &Arc) -> Arc { + pub unsafe fn borrow_raw(cu_module: *mut c_void, context: &Arc) -> Arc { Arc::new(Module { cu_module: cu_module as cuda_bindings::CUmodule, - device: device.clone(), + context: context.clone(), owned: false, }) } @@ -973,6 +1038,7 @@ impl Module { /// Looks up a device function by name within this module. pub fn load_function(self: &Arc, fn_name: &str) -> Result { + self.context.bind_to_thread()?; let cu_function = unsafe { module::get_function(self.cu_module, fn_name) }?; Ok(Function { cu_function, @@ -1104,33 +1170,33 @@ mod tests { } fn has_gpu() -> bool { - Device::device_count().map(|n| n > 0).unwrap_or(false) + Context::device_count().map(|n| n > 0).unwrap_or(false) } - /// The stream pool is process-wide and keyed by ordinal, so the tests that - /// assert on its contents for ordinal 0 must not overlap: a concurrent - /// test's `new_stream` pops, and its `Device` drop discards the whole + /// The stream pool is process-wide and keyed by context, so the tests that + /// assert on the primary context pool must not overlap: a concurrent + /// test's `new_stream` pops, and its `Context` drop discards the whole /// entry. Nothing else in this test binary touches the pool. fn pool_tests_lock() -> std::sync::MutexGuard<'static, ()> { static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) } - /// Snapshot of the parked handles for `ordinal`. - fn pooled_handles(ordinal: usize) -> Vec { + /// Snapshot of parked handles for this CUDA context. + fn pooled_handles(context: &Context) -> Vec { stream_pool() .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .get(&ordinal) + .get(&(context.cu_ctx() as usize)) .cloned() .unwrap_or_default() } - fn borrow(owner: &Arc) -> Arc { + fn borrow(owner: &Arc) -> Arc { // SAFETY: `owner` retains the primary context for the duration of // every test below, and nothing destroys it concurrently. unsafe { - Device::borrow_raw( + Context::borrow_raw( owner.cu_ctx().cast(), owner.cu_device() as c_int, owner.ordinal(), @@ -1142,11 +1208,11 @@ mod tests { fn ordinals_beyond_c_int_are_rejected_not_truncated() { let invalid = DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_DEVICE); // Checked before `cuInit`, so this holds with or without a driver. - assert_eq!(Device::new(usize::MAX).err(), Some(invalid)); - assert_eq!(Device::raw_device(usize::MAX).err(), Some(invalid)); + assert_eq!(Context::new(usize::MAX).err(), Some(invalid)); + assert_eq!(Context::raw_device(usize::MAX).err(), Some(invalid)); // The first ordinal past `c_int` would have wrapped to device 0. let wraps_to_zero = (c_int::MAX as usize) + 1 + (c_int::MAX as usize) + 1; - assert_eq!(Device::new(wraps_to_zero).err(), Some(invalid)); + assert_eq!(Context::new(wraps_to_zero).err(), Some(invalid)); assert_eq!(ordinal_to_c_int(0), Ok(0)); assert_eq!(ordinal_to_c_int(c_int::MAX as usize), Ok(c_int::MAX)); } @@ -1157,12 +1223,12 @@ mod tests { return; } let _serialized = pool_tests_lock(); - let owner = Device::new(0).unwrap(); + let owner = Context::new(0).unwrap(); let first = owner.new_stream().unwrap(); let handle = first.cu_stream() as usize; drop(first); assert!( - pooled_handles(0).contains(&handle), + pooled_handles(&owner).contains(&handle), "an owned device's stream is parked, not destroyed" ); let second = owner.new_stream().unwrap(); @@ -1175,14 +1241,14 @@ mod tests { return; } let _serialized = pool_tests_lock(); - let owner = Device::new(0).unwrap(); + let owner = Context::new(0).unwrap(); // Seed the pool with a handle from the owned device so that a // borrowed `new_stream` that consulted the pool would be caught. let parked = { let seed = owner.new_stream().unwrap(); seed.cu_stream() as usize }; - assert!(pooled_handles(0).contains(&parked)); + assert!(pooled_handles(&owner).contains(&parked)); let borrowed = borrow(&owner); let stream = borrowed.new_stream().unwrap(); @@ -1193,18 +1259,18 @@ mod tests { belong to contexts this crate retained, not to the borrowed one" ); assert!( - pooled_handles(0).contains(&parked), + pooled_handles(&owner).contains(&parked), "and leaves the pool as it was" ); - unsafe { stream.synchronize() }.unwrap(); + stream.synchronize().unwrap(); drop(stream); assert!( - !pooled_handles(0).contains(&handle), + !pooled_handles(&owner).contains(&handle), "a borrowed device's stream must be destroyed on drop, never parked \ - where a later owned Device for this ordinal could pop it" + where a later owned Context for this ordinal could pop it" ); - assert!(pooled_handles(0).contains(&parked)); + assert!(pooled_handles(&owner).contains(&parked)); } #[test] diff --git a/cuda-core/src/simt/event.rs b/cuda-core/src/simt/event.rs deleted file mode 100644 index 0ef390afa3..0000000000 --- a/cuda-core/src/simt/event.rs +++ /dev/null @@ -1,144 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! CUDA event management (RAII, timing, synchronization). -//! -//! A [`CudaEvent`] wraps a `CUevent` handle and ties its lifetime to its -//! parent [`CudaContext`]. Events are the fundamental synchronization -//! primitive between CUDA streams: record an event on one stream, then wait -//! on it from another to establish ordering. -//! -//! # Timing -//! -//! By default, events are created with `CU_EVENT_DISABLE_TIMING` for lower -//! overhead. Pass `Some(CU_EVENT_DEFAULT)` to -//! [`CudaContext::new_event`] or [`CudaStream::record_event`] if you -//! need [`elapsed_ms`](CudaEvent::elapsed_ms). - -use crate::error::{DriverError, IntoResult}; -use crate::simt::context::CudaContext; -use crate::simt::stream::CudaStream; -use std::mem::MaybeUninit; -use std::sync::Arc; - -/// An RAII wrapper around a `CUevent` handle. -/// -/// Holds an `Arc` to ensure the context outlives the event. -/// Destroyed automatically via `cuEventDestroy` on [`Drop`]. -#[derive(Debug)] -pub struct CudaEvent { - /// Raw CUDA event handle. - pub(crate) cu_event: cuda_bindings::CUevent, - /// Owning context. Kept alive for the lifetime of this event. - pub(crate) ctx: Arc, -} - -/// # Safety -/// -/// `CUevent` handles are not thread-local. The CUDA driver permits recording -/// and waiting on events from any thread, provided the owning context is bound. -unsafe impl Send for CudaEvent {} -/// See [`Send`] impl. -unsafe impl Sync for CudaEvent {} - -/// Destroys the underlying `CUevent` on drop. -/// -/// Binds the context to the current thread first (required by -/// `cuEventDestroy`). Errors are recorded on the context rather than -/// panicking. -impl Drop for CudaEvent { - fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); - self.ctx - .record_err(unsafe { cuda_bindings::cuEventDestroy_v2(self.cu_event).result() }); - } -} - -impl CudaContext { - /// Creates a new CUDA event in this context. - /// - /// `flags` defaults to `CU_EVENT_DISABLE_TIMING` when `None`. Use - /// `Some(CU_EVENT_DEFAULT)` to enable timing queries via - /// [`CudaEvent::elapsed_ms`]. - pub fn new_event( - self: &Arc, - flags: Option, - ) -> Result { - let flags = flags.unwrap_or(cuda_bindings::CUevent_flags_enum_CU_EVENT_DISABLE_TIMING); - self.bind_to_thread()?; - let mut cu_event = MaybeUninit::uninit(); - let cu_event = unsafe { - cuda_bindings::cuEventCreate(cu_event.as_mut_ptr(), flags).result()?; - cu_event.assume_init() - }; - Ok(CudaEvent { - cu_event, - ctx: self.clone(), - }) - } -} - -impl CudaEvent { - /// Returns the raw `CUevent` handle. - pub fn cu_event(&self) -> cuda_bindings::CUevent { - self.cu_event - } - - /// Returns the parent [`CudaContext`]. - pub fn context(&self) -> &Arc { - &self.ctx - } - - /// Records this event on `stream`. - /// - /// The event captures the point in the stream's work queue at the time of - /// the call. A subsequent [`CudaStream::wait`] on this event will block - /// the waiting stream until all work prior to the record point completes. - pub fn record(&self, stream: &CudaStream) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { cuda_bindings::cuEventRecord(self.cu_event, stream.cu_stream()).result() } - } - - /// Blocks the calling thread until this event has been recorded and all - /// preceding stream work has completed. - pub fn synchronize(&self) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { cuda_bindings::cuEventSynchronize(self.cu_event).result() } - } - - /// Returns `true` when all work captured by this event has completed, - /// `false` when it is still in flight. Never blocks. - /// - /// Wraps `cuEventQuery`, mapping `CUDA_SUCCESS` to `Ok(true)` and - /// `CUDA_ERROR_NOT_READY` to `Ok(false)`. Any other code is a real - /// driver error. - pub fn query(&self) -> Result { - self.ctx.bind_to_thread()?; - match unsafe { cuda_bindings::cuEventQuery(self.cu_event) } { - cuda_bindings::cudaError_enum_CUDA_SUCCESS => Ok(true), - cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_READY => Ok(false), - err => Err(DriverError(err)), - } - } - - /// Returns the elapsed time in milliseconds between `self` (start) and - /// `end`. - /// - /// Both events are synchronized before querying. Both events must have - /// been created **without** `CU_EVENT_DISABLE_TIMING`; otherwise the - /// driver returns `CUDA_ERROR_INVALID_HANDLE`. - /// - /// `self` must have been recorded before `end` in wall-clock time. - pub fn elapsed_ms(&self, end: &Self) -> Result { - self.synchronize()?; - end.synchronize()?; - let mut ms: f32 = 0.0; - unsafe { - cuda_bindings::cu_event_elapsed_time(&mut ms as *mut _, self.cu_event, end.cu_event) - .result()?; - } - Ok(ms) - } -} diff --git a/cuda-core/src/simt/stream.rs b/cuda-core/src/simt/stream.rs deleted file mode 100644 index 8a97417cc4..0000000000 --- a/cuda-core/src/simt/stream.rs +++ /dev/null @@ -1,252 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! CUDA stream management (RAII, host callbacks, fork/join). -//! -//! A [`CudaStream`] wraps a `CUstream` handle and ties its lifetime to its -//! parent [`CudaContext`]. Streams are created via -//! [`CudaContext::new_stream`] or [`CudaStream::fork`] and destroyed -//! automatically on [`Drop`]. -//! -//! # Ordering model -//! -//! Operations enqueued on the same stream execute in FIFO order. Operations on -//! different streams may overlap. Use [`fork`](CudaStream::fork) / -//! [`join`](CudaStream::join) or explicit events ([`record_event`](CudaStream::record_event), -//! [`wait`](CudaStream::wait)) to establish cross-stream ordering. -//! -//! # Host callbacks -//! -//! [`launch_host_function`](CudaStream::launch_host_function) enqueues a -//! host-side closure that the driver invokes after all preceding stream work -//! completes. This is the primary bridge between CUDA stream completion and -//! Rust `async` futures. - -use crate::error::{DriverError, IntoResult}; -use crate::simt::context::CudaContext; -use crate::simt::event::CudaEvent; -use std::ffi::c_void; -use std::mem::MaybeUninit; -use std::sync::atomic::Ordering; -use std::sync::Arc; - -/// An RAII wrapper around a `CUstream` handle. -/// -/// Holds an `Arc` to ensure the context outlives the stream. -/// A null `cu_stream` represents the per-context default stream (stream 0). -#[derive(Debug, PartialEq, Eq)] -pub struct CudaStream { - /// Raw CUDA stream handle. Null for the default stream. - pub(crate) cu_stream: cuda_bindings::CUstream, - /// Owning context. Kept alive for the lifetime of this stream. - pub(crate) ctx: Arc, -} - -/// # Safety -/// -/// `CUstream` handles are not thread-local. The CUDA driver permits enqueuing -/// work onto a stream from any thread, provided the owning context is bound -/// (which [`bind_to_thread`](CudaContext::bind_to_thread) ensures). -unsafe impl Send for CudaStream {} -/// See [`Send`] impl. -unsafe impl Sync for CudaStream {} - -/// Destroys the underlying `CUstream` on drop and decrements the context's -/// live stream count. -/// -/// The default stream (null handle) is never destroyed. Errors during -/// teardown are recorded on the context rather than panicking. -impl Drop for CudaStream { - fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); - if !self.cu_stream.is_null() { - self.ctx.num_streams.fetch_sub(1, Ordering::Relaxed); - self.ctx - .record_err(unsafe { cuda_bindings::cuStreamDestroy_v2(self.cu_stream).result() }); - } - } -} - -impl CudaStream { - /// Returns the raw `CUstream` handle (null for the default stream). - pub fn cu_stream(&self) -> cuda_bindings::CUstream { - self.cu_stream - } - - /// Returns the parent [`CudaContext`]. - pub fn context(&self) -> &Arc { - &self.ctx - } - - /// Returns the priority this stream actually runs at. - /// - /// Wraps `cuStreamGetPriority`. This is the value to read after - /// [`CudaContext::new_stream_with_priority`](crate::simt::context::CudaContext::new_stream_with_priority), - /// which clamps an out-of-range request to the device's range without - /// reporting it. Lower numbers are higher priorities. The default stream - /// and any stream from [`CudaContext::new_stream`](crate::simt::context::CudaContext::new_stream) - /// report `0`. - pub fn priority(&self) -> Result { - self.ctx.bind_to_thread()?; - let mut priority = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuStreamGetPriority(self.cu_stream, priority.as_mut_ptr()).result()?; - Ok(priority.assume_init()) - } - } - - /// Blocks the calling thread until all work enqueued on this stream - /// completes. - pub fn synchronize(&self) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { cuda_bindings::cuStreamSynchronize(self.cu_stream) }.result() - } - - /// Returns `true` when every operation enqueued on this stream has - /// completed, `false` while any is still in flight. Never blocks. - /// - /// Wraps `cuStreamQuery`, mapping `CUDA_SUCCESS` to `Ok(true)` and - /// `CUDA_ERROR_NOT_READY` to `Ok(false)`, as - /// [`CudaEvent::query`](crate::simt::event::CudaEvent::query) does for - /// `cuEventQuery`. Any other code is a real driver error. - /// - /// This is the completion half of the async bridge that - /// [`launch_host_function`](Self::launch_host_function) opens: a - /// `Future::poll` implementation needs to ask whether work has finished - /// without parking the executor's thread, which - /// [`synchronize`](Self::synchronize) cannot answer. - /// - /// A sticky error from earlier work on the stream surfaces here rather - /// than as `Ok(false)`, the same way it would from - /// [`synchronize`](Self::synchronize). - /// - /// On the default stream this reports the completion of every blocking - /// stream in the context, since the default stream orders against them - /// all. - pub fn query(&self) -> Result { - self.ctx.bind_to_thread()?; - match unsafe { cuda_bindings::cuStreamQuery(self.cu_stream) } { - cuda_bindings::cudaError_enum_CUDA_SUCCESS => Ok(true), - cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_READY => Ok(false), - err => Err(DriverError(err)), - } - } - - /// Creates a new non-blocking stream that waits on all prior work in - /// `self` before executing its own. - /// - /// Semantically equivalent to creating a new stream and calling - /// [`join`](Self::join) on it with `self`, establishing a fork point - /// in the stream DAG. - pub fn fork(&self) -> Result, DriverError> { - self.ctx.bind_to_thread()?; - self.ctx.num_streams.fetch_add(1, Ordering::Relaxed); - let mut cu_stream = MaybeUninit::uninit(); - let cu_stream = unsafe { - cuda_bindings::cuStreamCreate( - cu_stream.as_mut_ptr(), - cuda_bindings::CUstream_flags_enum_CU_STREAM_NON_BLOCKING, - ) - .result()?; - cu_stream.assume_init() - }; - let stream = Arc::new(CudaStream { - cu_stream, - ctx: self.ctx.clone(), - }); - stream.join(self)?; - Ok(stream) - } - - /// Makes `self` wait on all prior work in `other`. - /// - /// Records an event on `other` and enqueues a wait on `self`. This is the - /// join side of the fork/join pattern: after this call, work enqueued on - /// `self` is guaranteed to observe all side effects of prior work on - /// `other`. - pub fn join(&self, other: &CudaStream) -> Result<(), DriverError> { - self.wait(&other.record_event(None)?) - } - - /// Records an event on this stream and returns it. - /// - /// `flags` defaults to `CU_EVENT_DISABLE_TIMING` when `None`, which is - /// cheaper than a timing-enabled event. Pass - /// `Some(CU_EVENT_DEFAULT)` if you need [`CudaEvent::elapsed_ms`]. - pub fn record_event( - &self, - flags: Option, - ) -> Result { - let event = self.ctx.new_event(flags)?; - event.record(self)?; - Ok(event) - } - - /// Enqueues a wait on `event` into this stream. - /// - /// All work enqueued on `self` after this call will not begin until - /// `event` has been recorded (i.e., all prior work on the stream that - /// recorded `event` has completed). - pub fn wait(&self, event: &CudaEvent) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { - cuda_bindings::cuStreamWaitEvent( - self.cu_stream, - event.cu_event(), - cuda_bindings::CUevent_wait_flags_enum_CU_EVENT_WAIT_DEFAULT, - ) - .result() - } - } - - /// Enqueues a host-side callback that the driver invokes after all prior - /// work on this stream completes. - /// - /// This is the bridge between CUDA stream ordering and Rust async: wrap a - /// `oneshot::Sender::send` or `Waker::wake` in `host_func` to unblock a - /// future when GPU work finishes. - /// - /// `host_func` is boxed, leaked into a raw pointer, and passed as user - /// data to `cuLaunchHostFunc`. The driver calls - /// `callback_wrapper` on a driver-internal - /// thread, which reclaims the box and invokes the closure. - /// - /// Panics inside the closure are caught and discarded to prevent unwinding - /// across the FFI boundary. - pub fn launch_host_function( - &self, - host_func: F, - ) -> Result<(), DriverError> { - let boxed = Box::new(host_func); - unsafe { - cuda_bindings::cuLaunchHostFunc( - self.cu_stream, - Some(Self::callback_wrapper::), - Box::into_raw(boxed) as *mut c_void, - ) - .result() - } - } - - /// `extern "C"` trampoline invoked by the CUDA driver on a driver-internal - /// thread when a host function callback fires. - /// - /// Reconstructs the `Box` from the raw pointer and calls the closure. - /// Panics are caught with `catch_unwind` to prevent unwinding across the - /// C ABI boundary. - /// - /// # Safety - /// - /// - `callback` must be a pointer produced by `Box::into_raw(Box::new(f))` - /// where `f: F`. - /// - Must be called exactly once per enqueued callback (double-free - /// otherwise). - unsafe extern "C" fn callback_wrapper(callback: *mut c_void) { - let _ = std::panic::catch_unwind(|| { - let callback: Box = unsafe { Box::from_raw(callback as *mut F) }; - callback(); - }); - } -} diff --git a/cuda-core/src/simt/vmm.rs b/cuda-core/src/simt/vmm.rs deleted file mode 100644 index 4784fc6099..0000000000 --- a/cuda-core/src/simt/vmm.rs +++ /dev/null @@ -1,519 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! CUDA Virtual Memory Management (VMM) API wrappers. -//! -//! The VMM APIs provide fine-grained control over physical memory allocation, -//! virtual address reservation, and mapping. Unlike `cuMemAlloc`, which bundles -//! all three steps, VMM separates them so that physical memory from one device -//! can be mapped into another device's virtual address space -- the foundation -//! for P2P symmetric heaps. -//! -//! All handle types are RAII: `PhysicalAllocation` releases via `cuMemRelease`, -//! `VirtualReservation` frees via `cuMemAddressFree`, and `Mapping` unmaps -//! via `cuMemUnmap`. Drop order matters -- mappings must be dropped before the -//! physical allocation or virtual reservation they reference. -//! -//! The module also wraps CUDA's multicast objects (`cuMulticast*`, CUDA -//! 12.1+): `MulticastObject` (compiled out on pre-12.1 toolkits, see the -//! build script probe) builds an NVLink SHARP (NVLS) team whose -//! mapped VA ranges respond to device-side `multimem.*` instructions with -//! switch-side reduction and broadcast across every bound GPU. See -//! `tests/vmm_multicast.rs` and the `nvls_all_reduce` example. - -use crate::error::{DriverError, IntoResult}; -use cuda_bindings::CUdeviceptr; -use std::mem::MaybeUninit; - -/// Sets the device ordinal on a `CUmemLocation_st`. -/// -/// CUDA 13.2 wraps `id` in an anonymous union (`__bindgen_anon_1.id`), while -/// older versions expose it directly. The memory layout is identical -- `id` is -/// always at offset 4 (after the `type_` enum). Writing via pointer works across -/// both layouts. -unsafe fn set_mem_location_device( - loc: &mut cuda_bindings::CUmemLocation_st, - device: cuda_bindings::CUdevice, -) { - loc.type_ = cuda_bindings::CUmemLocationType_enum_CU_MEM_LOCATION_TYPE_DEVICE; - unsafe { - let base = loc as *mut _ as *mut u8; - (base.add(4) as *mut i32).write(device); - } -} - -/// A physical memory allocation created by `cuMemCreate`. -/// -/// Owns the underlying `CUmemGenericAllocationHandle`. The allocation lives on -/// a specific device and can be mapped into any device's VA space that has been -/// granted access. -/// -/// Dropping this releases the physical memory. All `Mapping`s referencing this -/// allocation must be dropped first. -pub struct PhysicalAllocation { - handle: cuda_bindings::CUmemGenericAllocationHandle, - size: usize, - device: cuda_bindings::CUdevice, -} - -impl PhysicalAllocation { - /// Allocates `size` bytes of physical memory on `device`. - /// - /// `size` must be a multiple of the allocation granularity for the device - /// (query via [`allocation_granularity`]). - pub fn new(device: cuda_bindings::CUdevice, size: usize) -> Result { - let mut prop: cuda_bindings::CUmemAllocationProp_st = unsafe { std::mem::zeroed() }; - prop.type_ = cuda_bindings::CUmemAllocationType_enum_CU_MEM_ALLOCATION_TYPE_PINNED; - unsafe { set_mem_location_device(&mut prop.location, device) }; - - let mut handle = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuMemCreate(handle.as_mut_ptr(), size, &prop, 0).result()?; - Ok(Self { - handle: handle.assume_init(), - size, - device, - }) - } - } - - /// Returns the raw `CUmemGenericAllocationHandle`. - pub fn handle(&self) -> cuda_bindings::CUmemGenericAllocationHandle { - self.handle - } - - /// Returns the allocation size in bytes. - pub fn size(&self) -> usize { - self.size - } - - /// Returns the device this allocation lives on. - pub fn device(&self) -> cuda_bindings::CUdevice { - self.device - } -} - -impl Drop for PhysicalAllocation { - fn drop(&mut self) { - unsafe { - let _ = cuda_bindings::cuMemRelease(self.handle).result(); - } - } -} - -/// A reserved virtual address range created by `cuMemAddressReserve`. -/// -/// Owns a contiguous VA range `[base, base + size)`. Physical memory can be -/// mapped into this range via [`Mapping::new`]. The range is freed on drop. -/// -/// All `Mapping`s within this range must be dropped before the reservation. -pub struct VirtualReservation { - base: CUdeviceptr, - size: usize, -} - -impl VirtualReservation { - /// Reserves `size` bytes of virtual address space. - /// - /// The driver chooses the base address. `size` must be a multiple of the - /// allocation granularity. `alignment` can be 0 to let the driver choose. - pub fn new(size: usize, alignment: usize) -> Result { - let mut base = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuMemAddressReserve(base.as_mut_ptr(), size, alignment, 0, 0) - .result()?; - Ok(Self { - base: base.assume_init(), - size, - }) - } - } - - /// Returns the base device pointer of the reserved range. - pub fn base(&self) -> CUdeviceptr { - self.base - } - - /// Returns the reserved size in bytes. - pub fn size(&self) -> usize { - self.size - } -} - -impl Drop for VirtualReservation { - fn drop(&mut self) { - unsafe { - let _ = cuda_bindings::cuMemAddressFree(self.base, self.size).result(); - } - } -} - -/// A mapping of physical memory into a virtual address range. -/// -/// Created by [`Mapping::new`], which calls `cuMemMap` to bind a -/// `PhysicalAllocation` (or a portion of it) to a region within a -/// `VirtualReservation`. Dropped via `cuMemUnmap`. -pub struct Mapping { - va: CUdeviceptr, - size: usize, -} - -impl Mapping { - /// Maps `size` bytes of `phys` at `offset` into virtual address `va`. - /// - /// `va` must lie within a `VirtualReservation`. `offset` is the byte - /// offset into the physical allocation (typically 0 for full mappings). - /// `size` must be a multiple of the allocation granularity. - pub fn new( - va: CUdeviceptr, - size: usize, - phys: &PhysicalAllocation, - offset: usize, - ) -> Result { - unsafe { - cuda_bindings::cuMemMap(va, size, offset, phys.handle(), 0).result()?; - } - Ok(Self { va, size }) - } - - /// Maps `size` bytes of a multicast object at `offset` into virtual - /// address `va`. - /// - /// The resulting VA is a *multicast* view: `multimem.*` PTX instructions - /// issued against it operate on every copy bound to the object (see - /// [`MulticastObject`]). Like [`Mapping::new`], the mapping is not - /// accessible until [`set_access`] grants the accessing device permission. - /// - /// All devices must have been added to the multicast object (via - /// [`MulticastObject::add_device`]) before mapping it. - #[cfg(cuda_has_multicast)] - pub fn new_multicast( - va: CUdeviceptr, - size: usize, - multicast: &MulticastObject, - offset: usize, - ) -> Result { - unsafe { - cuda_bindings::cuMemMap(va, size, offset, multicast.handle(), 0).result()?; - } - Ok(Self { va, size }) - } - - /// Returns the virtual address this mapping occupies. - pub fn va(&self) -> CUdeviceptr { - self.va - } - - /// Returns the mapped size in bytes. - pub fn size(&self) -> usize { - self.size - } -} - -impl Drop for Mapping { - fn drop(&mut self) { - unsafe { - let _ = cuda_bindings::cuMemUnmap(self.va, self.size).result(); - } - } -} - -/// Sets read/write access on a virtual address range for one or more devices. -/// -/// After calling `cuMemMap`, the mapping is not yet accessible. This function -/// grants the specified `devices` read/write permission on the range -/// `[va, va + size)`. -/// -/// Typically called once after all mappings within a reservation are established. -pub fn set_access( - va: CUdeviceptr, - size: usize, - devices: &[cuda_bindings::CUdevice], -) -> Result<(), DriverError> { - let descs: Vec = devices - .iter() - .map(|&dev| { - let mut desc: cuda_bindings::CUmemAccessDesc_st = unsafe { std::mem::zeroed() }; - unsafe { set_mem_location_device(&mut desc.location, dev) }; - desc.flags = cuda_bindings::CUmemAccess_flags_enum_CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - desc - }) - .collect(); - - unsafe { cuda_bindings::cuMemSetAccess(va, size, descs.as_ptr(), descs.len()) }.result() -} - -/// Queries the minimum allocation granularity for VMM operations on `device`. -/// -/// All sizes passed to [`PhysicalAllocation::new`], [`VirtualReservation::new`], -/// and [`Mapping::new`] must be multiples of this value. -pub fn allocation_granularity(device: cuda_bindings::CUdevice) -> Result { - let mut prop: cuda_bindings::CUmemAllocationProp_st = unsafe { std::mem::zeroed() }; - prop.type_ = cuda_bindings::CUmemAllocationType_enum_CU_MEM_ALLOCATION_TYPE_PINNED; - unsafe { set_mem_location_device(&mut prop.location, device) }; - - let mut granularity = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuMemGetAllocationGranularity( - granularity.as_mut_ptr(), - &prop, - cuda_bindings::CUmemAllocationGranularity_flags_enum_CU_MEM_ALLOC_GRANULARITY_MINIMUM, - ) - .result()?; - Ok(granularity.assume_init()) - } -} - -/// Rounds `size` up to the nearest multiple of `granularity`. -pub fn align_size(size: usize, granularity: usize) -> usize { - (size + granularity - 1) & !(granularity - 1) -} - -/// Granularity flavor for [`multicast_granularity`]. -#[cfg(cuda_has_multicast)] -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum MulticastGranularity { - /// Minimum required granularity for multicast sizes and offsets. - Minimum, - /// Recommended granularity for best performance. - Recommended, -} - -#[cfg(cuda_has_multicast)] -impl MulticastGranularity { - fn to_flag(self) -> cuda_bindings::CUmulticastGranularity_flags { - match self { - MulticastGranularity::Minimum => { - cuda_bindings::CUmulticastGranularity_flags_enum_CU_MULTICAST_GRANULARITY_MINIMUM - } - MulticastGranularity::Recommended => { - cuda_bindings::CUmulticastGranularity_flags_enum_CU_MULTICAST_GRANULARITY_RECOMMENDED - } - } - } -} - -/// Fills a `CUmulticastObjectProp_st` for a single-process team of -/// `num_devices` GPUs binding up to `size` bytes each. -/// -/// `handleTypes` is left at 0: the object cannot be exported to other -/// processes. Multi-process teams (exporting the handle over a POSIX file -/// descriptor or fabric handle) are out of scope for these wrappers. -#[cfg(cuda_has_multicast)] -fn multicast_prop(num_devices: u32, size: usize) -> cuda_bindings::CUmulticastObjectProp_st { - let mut prop: cuda_bindings::CUmulticastObjectProp_st = unsafe { std::mem::zeroed() }; - prop.numDevices = num_devices; - prop.size = size; - prop -} - -/// Returns whether `device` supports switch multicast and reduction -/// operations (`CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED`). -/// -/// Multicast requires an NVLink-switch-connected system (e.g. HGX/DGX -/// H100 or B200) and CUDA 12.1+. -#[cfg(cuda_has_multicast)] -pub fn multicast_supported(device: cuda_bindings::CUdevice) -> Result { - let mut value = MaybeUninit::uninit(); - let status = unsafe { - cuda_bindings::cuDeviceGetAttribute( - value.as_mut_ptr(), - cuda_bindings::CUdevice_attribute_enum_CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, - device, - ) - }; - // Drivers older than CUDA 12.1 do not know this attribute and answer - // CUDA_ERROR_INVALID_VALUE. That means "no multicast on this system", - // not a failure, so callers can probe-and-skip uniformly. - if status == cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE { - return Ok(false); - } - status.result()?; - Ok(unsafe { value.assume_init() } != 0) -} - -/// Built against a pre-12.1 CUDA toolkit whose headers lack the multicast -/// API: no device can be used for multicast through these wrappers. -#[cfg(not(cuda_has_multicast))] -pub fn multicast_supported(_device: cuda_bindings::CUdevice) -> Result { - Ok(false) -} - -/// Queries the multicast size/offset granularity for a team of -/// `num_devices` GPUs binding `size` bytes each. -/// -/// All sizes and offsets passed to [`MulticastObject::new`] and -/// [`MulticastObject::bind_mem`] must be multiples of the -/// [`MulticastGranularity::Minimum`] value; use -/// [`MulticastGranularity::Recommended`] for best performance. -#[cfg(cuda_has_multicast)] -pub fn multicast_granularity( - num_devices: u32, - size: usize, - granularity: MulticastGranularity, -) -> Result { - let prop = multicast_prop(num_devices, size); - let mut value = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuMulticastGetGranularity(value.as_mut_ptr(), &prop, granularity.to_flag()) - .result()?; - Ok(value.assume_init()) - } -} - -/// A multicast object created by `cuMulticastCreate`: one virtual "team" -/// handle backed by up to one physical allocation per participating GPU. -/// -/// Once every team device is added ([`add_device`](Self::add_device)) and has -/// bound physical memory ([`bind_mem`](Self::bind_mem)), the object can be -/// mapped into each device's VA space with [`Mapping::new_multicast`]. -/// Device-side `multimem.ld_reduce` / `multimem.st` / `multimem.red` -/// instructions issued against that mapping operate on *all* bound copies at -/// once -- the NVSwitch performs the reduction/broadcast in the fabric -/// (NVLink SHARP, the mechanism behind NCCL's NVLS algorithm). -/// -/// Lifecycle rules, in order: -/// 1. [`MulticastObject::new`] with the final team size. -/// 2. [`add_device`](Self::add_device) exactly `num_devices` times. -/// 3. [`bind_mem`](Self::bind_mem) per device (after ALL devices are added). -/// 4. [`Mapping::new_multicast`] + [`set_access`] per device. -/// -/// Teardown reverses it: mappings first, then bindings ([`MulticastBinding`]), -/// then this object, then the physical allocations. -/// -/// Dropping releases the handle via `cuMemRelease` (the documented release -/// path for multicast objects). -#[cfg(cuda_has_multicast)] -pub struct MulticastObject { - handle: cuda_bindings::CUmemGenericAllocationHandle, - size: usize, - num_devices: u32, -} - -#[cfg(cuda_has_multicast)] -impl MulticastObject { - /// Creates a multicast object for a team of `num_devices` GPUs binding - /// up to `size` bytes each. - /// - /// `size` must be a multiple of the minimum multicast granularity - /// (query via [`multicast_granularity`]). The object is single-process - /// only (no exportable handle types). - pub fn new(num_devices: u32, size: usize) -> Result { - let prop = multicast_prop(num_devices, size); - let mut handle = MaybeUninit::uninit(); - unsafe { - cuda_bindings::cuMulticastCreate(handle.as_mut_ptr(), &prop).result()?; - Ok(Self { - handle: handle.assume_init(), - size, - num_devices, - }) - } - } - - /// Adds `device` to the multicast team. - /// - /// Must be called exactly [`num_devices`](Self::num_devices) times, once - /// per device, before any memory is bound. - pub fn add_device(&self, device: cuda_bindings::CUdevice) -> Result<(), DriverError> { - unsafe { cuda_bindings::cuMulticastAddDevice(self.handle, device).result() } - } - - /// Binds `size` bytes of `phys` (starting at `mem_offset`) to this - /// multicast object at `mc_offset`. - /// - /// The bound device is taken from the physical allocation. All offsets - /// and `size` must be multiples of the minimum multicast granularity, - /// and every team device must already have been added. A CUDA context - /// for the bound device must be current. - /// - /// The returned [`MulticastBinding`] unbinds on drop; it must be dropped - /// before this object and before `phys`. - pub fn bind_mem( - &self, - mc_offset: usize, - phys: &PhysicalAllocation, - mem_offset: usize, - size: usize, - ) -> Result { - unsafe { - cuda_bindings::cuMulticastBindMem( - self.handle, - mc_offset, - phys.handle(), - mem_offset, - size, - 0, - ) - .result()?; - } - Ok(MulticastBinding { - mc_handle: self.handle, - device: phys.device(), - mc_offset, - size, - }) - } - - /// Returns the raw multicast `CUmemGenericAllocationHandle`. - pub fn handle(&self) -> cuda_bindings::CUmemGenericAllocationHandle { - self.handle - } - - /// Returns the per-device bind capacity in bytes. - pub fn size(&self) -> usize { - self.size - } - - /// Returns the team size this object was created for. - pub fn num_devices(&self) -> u32 { - self.num_devices - } -} - -#[cfg(cuda_has_multicast)] -impl Drop for MulticastObject { - fn drop(&mut self) { - unsafe { - let _ = cuda_bindings::cuMemRelease(self.handle).result(); - } - } -} - -/// One device's physical memory bound into a [`MulticastObject`]. -/// -/// Created by [`MulticastObject::bind_mem`]; unbinds via `cuMulticastUnbind` -/// on drop. Must be dropped before the multicast object and before the -/// physical allocation it binds. -#[cfg(cuda_has_multicast)] -pub struct MulticastBinding { - mc_handle: cuda_bindings::CUmemGenericAllocationHandle, - device: cuda_bindings::CUdevice, - mc_offset: usize, - size: usize, -} - -#[cfg(cuda_has_multicast)] -impl MulticastBinding { - /// Returns the device whose memory is bound. - pub fn device(&self) -> cuda_bindings::CUdevice { - self.device - } -} - -#[cfg(cuda_has_multicast)] -impl Drop for MulticastBinding { - fn drop(&mut self) { - unsafe { - let _ = cuda_bindings::cuMulticastUnbind( - self.mc_handle, - self.device, - self.mc_offset, - self.size, - ) - .result(); - } - } -} diff --git a/cuda-core/src/stream.rs b/cuda-core/src/stream.rs new file mode 100644 index 0000000000..ee807b775f --- /dev/null +++ b/cuda-core/src/stream.rs @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Stream ordering and priority helpers. + +pub use crate::runtime::{Stream, Stream as CudaStream}; +use crate::{DriverError, Event, IntoResult}; +use std::{mem::MaybeUninit, sync::Arc}; + +impl Stream { + /// Returns the priority this stream actually runs at. + /// + /// Wraps `cuStreamGetPriority`. This is the value to read after + /// [`Context::new_stream_with_priority`](crate::context::Context::new_stream_with_priority), + /// which clamps an out-of-range request to the device's range without + /// reporting it. Lower numbers are higher priorities. The default stream + /// and any stream from [`Context::new_stream`](crate::context::Context::new_stream) + /// report `0`. + pub fn priority(&self) -> Result { + self.context.bind_to_thread()?; + let mut priority = MaybeUninit::uninit(); + unsafe { + cuda_bindings::cuStreamGetPriority(self.cu_stream, priority.as_mut_ptr()).result()?; + Ok(priority.assume_init()) + } + } + /// Creates a new non-blocking stream that waits on all prior work in + /// `self` before executing its own. + /// + /// Semantically equivalent to creating a new stream and calling + /// [`join`](Self::join) on it with `self`, establishing a fork point + /// in the stream DAG. + pub fn fork(&self) -> Result, DriverError> { + let stream = self.context.new_stream()?; + stream.join(self)?; + Ok(stream) + } + /// Makes `self` wait on all prior work in `other`. + /// + /// Records an event on `other` and enqueues a wait on `self`. This is the + /// join side of the fork/join pattern: after this call, work enqueued on + /// `self` is guaranteed to observe all side effects of prior work on + /// `other`. + pub fn join(&self, other: &Stream) -> Result<(), DriverError> { + self.wait(&other.record_event(None)?) + } + /// Records an event on this stream and returns it. + /// + /// `flags` defaults to `CU_EVENT_DISABLE_TIMING` when `None`, which is + /// cheaper than a timing-enabled event. Pass + /// `Some(CU_EVENT_DEFAULT)` if you need [`Event::elapsed_ms`]. + pub fn record_event( + &self, + flags: Option, + ) -> Result { + let event = self.context.new_event_with_flags(flags)?; + event.record(self)?; + Ok(event) + } + /// Enqueues a wait on `event` into this stream. + /// + /// All work enqueued on `self` after this call will not begin until + /// `event` has been recorded (i.e., all prior work on the stream that + /// recorded `event` has completed). + pub fn wait(&self, event: &Event) -> Result<(), DriverError> { + self.wait_event(event) + } +} diff --git a/cuda-core/tests/borrow_raw.rs b/cuda-core/tests/borrow_raw.rs index b2713d4b3e..6834d9bc3e 100644 --- a/cuda-core/tests/borrow_raw.rs +++ b/cuda-core/tests/borrow_raw.rs @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Verifies the ownership contract of `Device::borrow_raw` and +//! Verifies the ownership contract of `Context::borrow_raw` and //! `Stream::borrow_raw`: dropping a borrowed wrapper must not release the //! primary context or destroy the stream, so the source handles keep //! working afterward. use core::ffi::{c_int, c_void}; -use cuda_core::{Device, Function, Module, Stream}; +use cuda_core::{Context, Function, Module, Stream}; const NOOP_PTX: &str = "\ .version 7.0 @@ -23,12 +23,12 @@ const NOOP_PTX: &str = "\ "; fn has_gpu() -> bool { - Device::device_count().map(|n| n > 0).unwrap_or(false) + Context::device_count().map(|n| n > 0).unwrap_or(false) } -fn borrow_device(src: &Device) -> std::sync::Arc { +fn borrow_device(src: &Context) -> std::sync::Arc { unsafe { - Device::borrow_raw( + Context::borrow_raw( src.cu_ctx() as *mut c_void, src.cu_device() as c_int, src.ordinal(), @@ -36,7 +36,7 @@ fn borrow_device(src: &Device) -> std::sync::Arc { } } -fn borrow_stream(src: &Stream, dev: &std::sync::Arc) -> std::sync::Arc { +fn borrow_stream(src: &Stream, dev: &std::sync::Arc) -> std::sync::Arc { unsafe { Stream::borrow_raw(src.cu_stream() as *mut c_void, dev) } } @@ -45,16 +45,16 @@ fn borrowed_drop_leaves_source_device_usable() { if !has_gpu() { return; } - let source = Device::new(0).unwrap(); + let source = Context::new(0).unwrap(); { let borrowed = borrow_device(&source); borrowed.bind_to_thread().unwrap(); - unsafe { borrowed.synchronize() }.unwrap(); + borrowed.synchronize().unwrap(); } source.bind_to_thread().unwrap(); - unsafe { source.synchronize() }.unwrap(); + source.synchronize().unwrap(); let _stream = source.new_stream().unwrap(); } @@ -63,16 +63,16 @@ fn borrowed_drop_leaves_source_stream_usable() { if !has_gpu() { return; } - let source_dev = Device::new(0).unwrap(); + let source_dev = Context::new(0).unwrap(); let source_stream = source_dev.new_stream().unwrap(); { let borrowed_dev = borrow_device(&source_dev); let borrowed_stream = borrow_stream(&source_stream, &borrowed_dev); - unsafe { borrowed_stream.synchronize() }.unwrap(); + borrowed_stream.synchronize().unwrap(); } - unsafe { source_stream.synchronize() }.unwrap(); + source_stream.synchronize().unwrap(); } #[test] @@ -80,7 +80,7 @@ fn borrowed_module_drop_does_not_unload() { if !has_gpu() { return; } - let source_dev = Device::new(0).unwrap(); + let source_dev = Context::new(0).unwrap(); let source_module = source_dev.load_module_from_ptx_src(NOOP_PTX).unwrap(); { @@ -100,7 +100,7 @@ fn borrowed_function_uses_parent_module() { if !has_gpu() { return; } - let source_dev = Device::new(0).unwrap(); + let source_dev = Context::new(0).unwrap(); let source_module = source_dev.load_module_from_ptx_src(NOOP_PTX).unwrap(); let source_fn = source_module.load_function("noop").unwrap(); @@ -116,15 +116,15 @@ fn many_borrowed_drops_do_not_invalidate_source() { if !has_gpu() { return; } - let source_dev = Device::new(0).unwrap(); + let source_dev = Context::new(0).unwrap(); let source_stream = source_dev.new_stream().unwrap(); for _ in 0..8 { let borrowed_dev = borrow_device(&source_dev); let borrowed_stream = borrow_stream(&source_stream, &borrowed_dev); - unsafe { borrowed_stream.synchronize() }.unwrap(); + borrowed_stream.synchronize().unwrap(); } - unsafe { source_stream.synchronize() }.unwrap(); + source_stream.synchronize().unwrap(); let _another = source_dev.new_stream().unwrap(); } diff --git a/cuda-core/tests/simt_context_limits.rs b/cuda-core/tests/context_limits.rs similarity index 100% rename from cuda-core/tests/simt_context_limits.rs rename to cuda-core/tests/context_limits.rs diff --git a/cuda-core/tests/simt_context_sync_policy.rs b/cuda-core/tests/context_sync_policy.rs similarity index 100% rename from cuda-core/tests/simt_context_sync_policy.rs rename to cuda-core/tests/context_sync_policy.rs diff --git a/cuda-core/tests/derive_device_copy/fail_enum.stderr b/cuda-core/tests/derive_device_copy/fail_enum.stderr index de29338391..021b4ddbec 100644 --- a/cuda-core/tests/derive_device_copy/fail_enum.stderr +++ b/cuda-core/tests/derive_device_copy/fail_enum.stderr @@ -1,4 +1,4 @@ -error: `#[derive(DeviceCopy)]` cannot be applied to enums: `DeviceCopy` requires every bit pattern (including the all-zero pattern written by `DeviceBuffer::zeroed`) to be a valid value, but an enum's discriminant leaves most byte patterns invalid, so a zeroed or device-written buffer could materialize an out-of-range discriminant (undefined behavior). If you can guarantee every device-produced byte pattern is a valid variant, write `unsafe impl DeviceCopy for ... {}` by hand. +error: `#[derive(DeviceCopy)]` cannot be applied to enums: this derive cannot prove that the all-zero representation selects a valid discriminant and payload. `DeviceCopy` requires byte-copyable values and a valid all-zero representation, but not validity of every bit pattern. If the enum meets the full safety contract, write `unsafe impl DeviceCopy for ... {}` by hand. --> tests/derive_device_copy/fail_enum.rs:12:6 | 12 | enum Tag { diff --git a/cuda-core/tests/simt_device_buffer.rs b/cuda-core/tests/device_buffer.rs similarity index 76% rename from cuda-core/tests/simt_device_buffer.rs rename to cuda-core/tests/device_buffer.rs index 9d67b0471d..1e2ec8b925 100644 --- a/cuda-core/tests/simt_device_buffer.rs +++ b/cuda-core/tests/device_buffer.rs @@ -3,8 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DriverError, PinnedHostBuffer}; -use std::sync::{mpsc, Mutex, MutexGuard}; +use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DeviceRepr, DriverError, PinnedHostBuffer}; +use std::sync::{mpsc, Arc, Mutex, MutexGuard}; use std::thread::JoinHandle; use std::time::Duration; @@ -114,6 +114,86 @@ fn device_buffer_zeroed_initializes_with_zeros() { assert_eq!(host_vec, &[0.0, 0.0, 0.0, 0.0]); } +#[test] +fn zeroed_waits_for_initialization_before_cross_stream_readback() { + let _gated_test_guard = lock_gated_test(); + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let writer = ctx.new_stream().expect("failed to create writer stream"); + let reader = ctx.new_stream().expect("failed to create reader stream"); + let release_gate = gate_stream(&writer); + let (started_tx, started_rx) = mpsc::channel(); + let (completion_tx, completion_rx) = mpsc::channel(); + let worker_writer = writer.clone(); + let worker = std::thread::spawn(move || { + worker_writer.context().bind_to_thread().unwrap(); + started_tx.send(()).unwrap(); + // Inspect byte storage, never potentially invalid Rust bool values, + // so a regression in initialization ordering cannot make this test UB. + let buffer = DeviceBuffer::::zeroed(&worker_writer, 1 << 20); + completion_tx.send(buffer).unwrap(); + }); + let buffer = finish_gated_worker( + "zeroed initialization", + release_gate, + started_rx, + completion_rx, + worker, + ) + .expect("zeroed allocation failed"); + assert!(writer.query().expect("writer query failed")); + assert_eq!(buffer.to_host_vec(&reader).unwrap(), vec![0; 1 << 20]); + + // Exercise bool's safe constructor too, reading its representation as + // bytes so the test remains safe even if the ordering regresses. + let flags = DeviceBuffer::::zeroed(&writer, 1024).unwrap(); + assert!(writer.query().unwrap()); + assert_eq!( + flags.cast_elem::().to_host_vec(&reader).unwrap(), + vec![0; 1024] + ); +} + +#[test] +fn owned_wrapper_roundtrip_preserves_allocation_and_stream() { + #[repr(transparent)] + struct Counter(std::cell::UnsafeCell); + // SAFETY: UnsafeCell has u32's layout and valid initialized representations; + // neither representation needs Drop or introduces ownership obligations. + unsafe impl DeviceRepr for Counter {} + + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.new_stream().unwrap(); + let values = [1_u32, 7, 42]; + // SAFETY: the following blocking copy initializes the allocation before + // any read. All accesses use the allocation stream and finish before drop. + let mut words = + unsafe { DeviceBuffer::::uninitialized_async(&stream, values.len()) }.unwrap(); + words.copy_from_host(&stream, &values).unwrap(); + let ptr = words.cu_deviceptr(); + let stream_refs = Arc::strong_count(&stream); + let context_refs = Arc::strong_count(&ctx); + + let counters = words.wrap_elem::(); + assert_eq!(counters.cu_deviceptr(), ptr); + assert_eq!(counters.len(), values.len()); + assert!(Arc::ptr_eq(counters.context(), &ctx)); + assert_eq!(Arc::strong_count(&stream), stream_refs); + assert_eq!(Arc::strong_count(&ctx), context_refs); + + let words = counters.peel_elem::(); + assert_eq!(words.cu_deviceptr(), ptr); + assert_eq!(Arc::strong_count(&stream), stream_refs); + assert_eq!(Arc::strong_count(&ctx), context_refs); + assert_eq!(words.to_host_vec(&stream).unwrap(), values); + drop(words); + stream.synchronize().unwrap(); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while Arc::strong_count(&stream) != stream_refs - 1 && std::time::Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(5)); + } + assert_eq!(Arc::strong_count(&stream), stream_refs - 1); +} + #[test] fn device_buffer_supports_empty_allocations() { let ctx = CudaContext::new(0).expect("failed to create CUDA context"); @@ -191,7 +271,7 @@ fn device_buffer_async_compat_methods_roundtrip() { } #[test] -fn async_allocation_ordinary_drop_waits_for_cross_stream_work() { +fn async_allocation_ordinary_drop_retires_cross_stream_work() { let _gated_test_guard = lock_gated_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let allocation_stream = ctx @@ -219,13 +299,12 @@ fn async_allocation_ordinary_drop_waits_for_cross_stream_work() { .send(()) .expect("failed to send drop completion"); }); - finish_gated_worker( - "ordinary async buffer drop", - release_gate, - started_rx, - completion_rx, - drop_thread, - ); + started_rx.recv_timeout(COMPLETION_TIMEOUT).unwrap(); + let completed = completion_rx.recv_timeout(COMPLETION_TIMEOUT); + // Release the gate even on failure so the test cannot strand the driver. + release_gate.send(()).unwrap(); + drop_thread.join().unwrap(); + completed.expect("ordinary buffer drop must not wait for GPU work"); ctx.synchronize().expect("context cleanup failed"); } @@ -332,11 +411,19 @@ fn drop_async_bind_error_preserves_ordinary_cleanup() { ctx.bind_to_thread() .expect("ordinary drop should leave the context usable"); - assert_eq!( - unsafe { cuda_bindings::cuMemGetAddressRange_v2(&mut base, &mut size, ptr) }, - cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_FOUND, - "ordinary drop must reclaim the still-armed allocation" - ); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + let status = unsafe { cuda_bindings::cuMemGetAddressRange_v2(&mut base, &mut size, ptr) }; + if status == cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_FOUND { + break; + } + assert_eq!(status, cuda_bindings::cudaError_enum_CUDA_SUCCESS); + assert!( + std::time::Instant::now() < deadline, + "retirement failed to reclaim the allocation" + ); + std::thread::sleep(Duration::from_millis(5)); + } } #[test] @@ -382,11 +469,11 @@ fn copy_from_host_with_pinned_source_allows_source_reuse_after_return() { // Dangerous-path regression: a stream-ordered (`cuMemAllocAsync`) buffer that // is dropped *implicitly* while a copy is still in flight, with no explicit -// synchronization by the caller. Ordinary `Drop` synchronizes the context -// before enqueueing the free, so the in-flight copy cannot race deallocation. +// synchronization by the caller. Completion leases retain the allocation +// until the in-flight copy finishes, so it cannot race deallocation. // Run under `compute-sanitizer --tool memcheck` to catch a regression. #[test] -fn uninitialized_async_implicit_drop_waits_for_pending_work() { +fn uninitialized_async_implicit_drop_retains_pending_work() { let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); diff --git a/cuda-core/tests/device_buffer_casts.rs b/cuda-core/tests/device_buffer_casts.rs new file mode 100644 index 0000000000..c2f80df130 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The element casts on `DeviceBuffer` must not forge invalid values (#1506). +//! +//! `cast_elem` and `cast_chunks` reinterpret bytes, so their target must be +//! `bytemuck::AnyBitPattern`; `wrap_elem` and `peel_elem` relabel through a +//! `cuda_core::DeviceRepr`. `bool` is neither, so no safe cast can +//! turn byte storage into a `DeviceBuffer`. These cases only need to +//! compile or fail to compile; none touches a GPU. Regenerate the +//! expectations with `TRYBUILD=overwrite cargo test -p cuda-core --test +//! device_buffer_casts` after an intentional diagnostic change. + +#[test] +fn device_buffer_cast_bounds() { + let t = trybuild::TestCases::new(); + t.pass("tests/device_buffer_casts/pass_bounded_casts.rs"); + t.compile_fail("tests/device_buffer_casts/fail_cast_elem_bool.rs"); + t.compile_fail("tests/device_buffer_casts/fail_cast_chunks_bool.rs"); + t.compile_fail("tests/device_buffer_casts/fail_wrap_elem_bool.rs"); +} diff --git a/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.rs b/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.rs new file mode 100644 index 0000000000..a48d819a58 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.rs @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_core::DeviceBuffer; + +// The length-adjusting cast has the same bound as `cast_elem`: `u32` words +// must not be regrouped into `bool` bytes. +fn forge(words: DeviceBuffer) -> Result, DeviceBuffer> { + words.cast_chunks() +} + +fn main() { + let _ = forge; +} diff --git a/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.stderr b/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.stderr new file mode 100644 index 0000000000..bade663b28 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_cast_chunks_bool.stderr @@ -0,0 +1,22 @@ +error[E0277]: the trait bound `bool: AnyBitPattern` is not satisfied + --> tests/device_buffer_casts/fail_cast_chunks_bool.rs:9:11 + | +9 | words.cast_chunks() + | ^^^^^^^^^^^ the trait `Pod` is not implemented for `bool` + | + = help: the following other types implement trait `Pod`: + () + ManuallyDrop + Option + PhantomData + PhantomPinned + Wrapping + [T; N] + __m128 + and $N others + = note: required for `bool` to implement `AnyBitPattern` +note: required by a bound in `DeviceBuffer::::cast_chunks` + --> src/device_buffer.rs + | + | pub fn cast_chunks(self) -> Result, Self> { + | ^^^^^^^^^^^^^ required by this bound in `DeviceBuffer::::cast_chunks` diff --git a/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.rs b/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.rs new file mode 100644 index 0000000000..8ca0408fe7 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.rs @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_core::DeviceBuffer; + +// Only 0 and 1 are valid `bool` values, so a byte buffer must not become a +// `DeviceBuffer` through the safe reinterpreting cast. +fn forge(bytes: DeviceBuffer) -> DeviceBuffer { + bytes.cast_elem() +} + +fn main() { + let _ = forge; +} diff --git a/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.stderr b/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.stderr new file mode 100644 index 0000000000..ca875fc475 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_cast_elem_bool.stderr @@ -0,0 +1,22 @@ +error[E0277]: the trait bound `bool: AnyBitPattern` is not satisfied + --> tests/device_buffer_casts/fail_cast_elem_bool.rs:9:11 + | +9 | bytes.cast_elem() + | ^^^^^^^^^ the trait `Pod` is not implemented for `bool` + | + = help: the following other types implement trait `Pod`: + () + ManuallyDrop + Option + PhantomData + PhantomPinned + Wrapping + [T; N] + __m128 + and $N others + = note: required for `bool` to implement `AnyBitPattern` +note: required by a bound in `DeviceBuffer::::cast_elem` + --> src/device_buffer.rs + | + | pub fn cast_elem(self) -> DeviceBuffer { + | ^^^^^^^^^^^^^ required by this bound in `DeviceBuffer::::cast_elem` diff --git a/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.rs b/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.rs new file mode 100644 index 0000000000..a03eb9f435 --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.rs @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_core::DeviceBuffer; + +// `bool` is not a `DeviceRepr`, and no downstream crate can make +// it one, so the wrapper path cannot forge a `bool` buffer either. +fn forge(bytes: DeviceBuffer) -> DeviceBuffer { + bytes.wrap_elem() +} + +fn main() { + let _ = forge; +} diff --git a/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.stderr b/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.stderr new file mode 100644 index 0000000000..58cf4ebdef --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/fail_wrap_elem_bool.stderr @@ -0,0 +1,16 @@ +error[E0277]: the trait bound `bool: DeviceRepr` is not satisfied + --> tests/device_buffer_casts/fail_wrap_elem_bool.rs:9:11 + | +9 | bytes.wrap_elem() + | ^^^^^^^^^ the trait `DeviceRepr` is not implemented for `bool` + | +help: the trait `DeviceRepr` is implemented for `Wrapping` + --> $WORKSPACE/cuda-repr/src/lib.rs + | + | unsafe impl DeviceRepr for core::num::Wrapping {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `DeviceBuffer::::wrap_elem` + --> src/device_buffer.rs + | + | pub fn wrap_elem>(self) -> DeviceBuffer { + | ^^^^^^^^^^^^^ required by this bound in `DeviceBuffer::::wrap_elem` diff --git a/cuda-core/tests/device_buffer_casts/pass_bounded_casts.rs b/cuda-core/tests/device_buffer_casts/pass_bounded_casts.rs new file mode 100644 index 0000000000..b00b87502a --- /dev/null +++ b/cuda-core/tests/device_buffer_casts/pass_bounded_casts.rs @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::num::Wrapping; + +use cuda_core::{DeviceBuffer, DeviceRepr}; + +// A wrapper that is not `Copy`, like the atomic types in cuda-device: it can +// be wrapped and peeled, but not reinterpreted into with `cast_elem`. +#[repr(transparent)] +struct Counter(std::cell::UnsafeCell); + +// SAFETY: transparent over UnsafeCell, which has u32's layout and valid +// representations, and neither representation owns resources or needs Drop. +unsafe impl DeviceRepr for Counter {} + +// Every bit pattern is a valid `f32`, so bytes written as `u32` are fine. +fn reinterpret(words: DeviceBuffer) -> DeviceBuffer { + words.cast_elem() +} + +// `Wrapping` is `Pod`, so it can be cast into as well as wrapped. +fn reinterpret_wrapping(words: DeviceBuffer) -> DeviceBuffer> { + words.cast_elem() +} + +fn group(floats: DeviceBuffer) -> Result, DeviceBuffer> { + floats.cast_chunks() +} + +fn wrap(words: DeviceBuffer) -> DeviceBuffer { + words.wrap_elem() +} + +fn peel(counters: DeviceBuffer) -> DeviceBuffer { + counters.peel_elem() +} + +fn wrap_std(words: DeviceBuffer) -> DeviceBuffer> { + words.wrap_elem() +} + +fn main() { + // Type-checking is the test; the functions are never called. + let _ = (reinterpret, reinterpret_wrapping, group, wrap, peel, wrap_std); +} diff --git a/cuda-core/tests/simt_device_buffer_leaks.rs b/cuda-core/tests/device_buffer_leaks.rs similarity index 89% rename from cuda-core/tests/simt_device_buffer_leaks.rs rename to cuda-core/tests/device_buffer_leaks.rs index 2e51c654ea..01d604eb74 100644 --- a/cuda-core/tests/simt_device_buffer_leaks.rs +++ b/cuda-core/tests/device_buffer_leaks.rs @@ -22,6 +22,13 @@ use std::sync::Arc; use cuda_core::{CudaContext, DeviceBuffer}; +fn wait_for_retirement(ctx: &Arc, baseline: usize) { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + while Arc::strong_count(ctx) != baseline && std::time::Instant::now() < deadline { + std::thread::sleep(std::time::Duration::from_millis(5)); + } +} + /// A live buffer must hold exactly one extra context strong count, and the /// count must return to baseline once the buffer drops. A construction /// helper that `mem::forget`s a context clone (as an arm/disarm pointer @@ -44,6 +51,7 @@ fn ctx_strong_count_returns_to_baseline_after_buffer_lifecycle() { drop(buf); } + wait_for_retirement(&ctx, baseline); assert_eq!( Arc::strong_count(&ctx), baseline, @@ -55,6 +63,7 @@ fn ctx_strong_count_returns_to_baseline_after_buffer_lifecycle() { drop(buf); } + wait_for_retirement(&ctx, baseline); assert_eq!( Arc::strong_count(&ctx), baseline, @@ -81,11 +90,14 @@ fn vram_returns_to_baseline_after_buffer_cycles() { free } + let baseline = Arc::strong_count(&ctx); + // Warm up driver allocator caches so the measured window is stable. for _ in 0..4 { let b = DeviceBuffer::::zeroed(&stream, 1 << 20).expect("warmup alloc failed"); drop(b); } + wait_for_retirement(&ctx, baseline); ctx.synchronize().expect("sync failed"); let before = free_mem(); @@ -93,6 +105,7 @@ fn vram_returns_to_baseline_after_buffer_cycles() { let b = DeviceBuffer::::zeroed(&stream, 1 << 20).expect("cycle alloc failed"); drop(b); } + wait_for_retirement(&ctx, baseline); ctx.synchronize().expect("sync failed"); let after = free_mem(); @@ -129,6 +142,7 @@ fn zero_length_construction_succeeds_for_both_constructors() { drop(zeroed); drop(from_host); + wait_for_retirement(&ctx, baseline); assert_eq!( Arc::strong_count(&ctx), baseline, diff --git a/cuda-core/tests/simt_device_copy_derive.rs b/cuda-core/tests/device_copy_derive.rs similarity index 100% rename from cuda-core/tests/simt_device_copy_derive.rs rename to cuda-core/tests/device_copy_derive.rs diff --git a/cuda-core/tests/device_copy_impls.rs b/cuda-core/tests/device_copy_impls.rs new file mode 100644 index 0000000000..68a17a70b6 --- /dev/null +++ b/cuda-core/tests/device_copy_impls.rs @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg_attr(feature = "f16", feature(f16))] + +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::num::Wrapping; + +use cuda_core::DeviceCopy; + +fn assert_device_copy() {} + +#[test] +fn device_copy_supports_half_types() { + assert_device_copy::(); + assert_device_copy::(); +} + +#[test] +fn tile_storage_types_also_implement_device_copy() { + assert_device_copy::(); + assert_device_copy::(); + assert_device_copy::(); + assert_device_copy::(); + assert_device_copy::(); + assert_device_copy::(); +} + +#[cfg(feature = "f16")] +#[test] +fn device_copy_supports_native_f16() { + assert_device_copy::(); +} + +#[test] +fn device_copy_covers_core_parity_types() { + // The representation-preserving wrappers are sound parity additions. + assert_device_copy::>(); + assert_device_copy::>(); + assert_device_copy::>(); +} + +#[test] +fn device_copy_supports_bool() { + // `bool` is `DeviceCopy` although only 0 and 1 are valid: a typed buffer + // holds valid values by construction, and the reinterpreting casts that + // could forge a `bool` from arbitrary bytes require `AnyBitPattern`, which + // `bool` does not implement (see tests/device_buffer_casts, #1506). + assert_device_copy::(); +} diff --git a/cuda-core/tests/simt_launch_contract_types.rs b/cuda-core/tests/launch_contract_types.rs similarity index 100% rename from cuda-core/tests/simt_launch_contract_types.rs rename to cuda-core/tests/launch_contract_types.rs diff --git a/cuda-core/tests/simt_vmm_multicast.rs b/cuda-core/tests/nvlink_multicast.rs similarity index 96% rename from cuda-core/tests/simt_vmm_multicast.rs rename to cuda-core/tests/nvlink_multicast.rs index a5716b9860..eb9220a4d3 100644 --- a/cuda-core/tests/simt_vmm_multicast.rs +++ b/cuda-core/tests/nvlink_multicast.rs @@ -20,8 +20,8 @@ //! (CUDA 12.1); see cuda-core's build script probe. #![cfg(cuda_has_multicast)] -use cuda_core::simt::context::CudaContext; -use cuda_core::simt::vmm; +use cuda_core::context::CudaContext; +use cuda_core::vmm; use cuda_core::IntoResult; use std::mem::MaybeUninit; @@ -115,7 +115,7 @@ fn multicast_two_gpu_team_roundtrip() { let byte_len = pattern.len() * std::mem::size_of::(); let stream0 = ctx0.default_stream(); unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( mappings[0].0.base(), pattern.as_ptr(), byte_len, @@ -127,7 +127,7 @@ fn multicast_two_gpu_team_roundtrip() { let mut readback = vec![0u32; 256]; unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( readback.as_mut_ptr(), mappings[0].0.base(), byte_len, diff --git a/cuda-core/tests/simt_pinned_host_buffer.rs b/cuda-core/tests/pinned_host_buffer.rs similarity index 64% rename from cuda-core/tests/simt_pinned_host_buffer.rs rename to cuda-core/tests/pinned_host_buffer.rs index 5ef153303b..fa069bf7a8 100644 --- a/cuda-core/tests/simt_pinned_host_buffer.rs +++ b/cuda-core/tests/pinned_host_buffer.rs @@ -125,3 +125,66 @@ fn device_buffer_can_be_refilled_from_pinned_host_async() { assert_eq!(output.as_slice(), &payload); } } + +// CUDA's page alignment alone does not satisfy every valid Rust element layout. +#[repr(C, align(1048576))] +#[derive(Clone, Copy, cuda_core::DeviceCopy)] +struct Overaligned([u8; 1048576]); + +#[test] +fn pinned_host_buffer_aligns_elements() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + // Keep a small allocation alive to exercise non-MiB-aligned CUDA allocations. + let _prefix = PinnedHostBuffer::::zeroed(&ctx, 4096).unwrap(); + let mut buffers = Vec::new(); + for _ in 0..8 { + let mut host = PinnedHostBuffer::::zeroed(&ctx, 1).unwrap(); + // Check before forming any references, so the regression fails without UB. + assert!(host.as_ptr().is_aligned()); + assert_eq!(host.num_bytes(), std::mem::size_of::()); + for element in host.as_slice() { + assert!(element.0.iter().all(|&byte| byte == 0)); + } + host.as_mut_slice()[0].0.fill(42); + assert!(host.as_slice()[0].0.iter().all(|&byte| byte == 42)); + buffers.push(host); + } + drop(buffers); + ctx.check_err().expect("failed to free pinned allocations"); +} + +#[test] +fn pinned_host_buffer_copies_overaligned_elements() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + // Initialize on the heap to avoid putting MiB-sized values on the stack. + let mut source = PinnedHostBuffer::::zeroed(&ctx, 2).unwrap(); + source.as_mut_slice()[0].0.fill(17); + source.as_mut_slice()[1].0.fill(42); + let copied = PinnedHostBuffer::from_slice(&ctx, source.as_slice()).unwrap(); + assert!(copied.as_ptr().is_aligned()); + assert_eq!(copied.len(), 2); + assert_eq!(copied.num_bytes(), 2 * std::mem::size_of::()); + assert!(copied.as_slice()[0].0.iter().all(|&byte| byte == 17)); + assert!(copied.as_slice()[1].0.iter().all(|&byte| byte == 42)); + drop((copied, source)); + ctx.check_err().expect("failed to free pinned allocations"); +} + +#[repr(align(1048576))] +#[derive(Clone, Copy, cuda_core::DeviceCopy)] +struct AlignedZst; + +#[test] +fn empty_and_zero_sized_pinned_buffers_preserve_alignment() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let empty = PinnedHostBuffer::::zeroed(&ctx, 0).unwrap(); + assert!(empty.as_ptr().is_aligned()); + assert!(empty.as_slice().is_empty()); + let mut zst = PinnedHostBuffer::::zeroed(&ctx, 8).unwrap(); + assert!(zst.as_ptr().is_aligned()); + assert_eq!(zst.num_bytes(), 0); + assert_eq!(zst.as_mut_slice().len(), 8); + let copied = PinnedHostBuffer::from_slice(&ctx, zst.as_slice()).unwrap(); + assert!(copied.as_ptr().is_aligned()); + assert_eq!(copied.as_slice().len(), 8); +} diff --git a/cuda-core/tests/shared_allocation.rs b/cuda-core/tests/shared_allocation.rs new file mode 100644 index 0000000000..9d3e35a2e1 --- /dev/null +++ b/cuda-core/tests/shared_allocation.rs @@ -0,0 +1,151 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cuda_core::{Context, DeviceBuffer, RawDeviceBuffer}; +use std::sync::{mpsc, Arc}; +use std::time::{Duration, Instant}; + +fn metadata(context: usize) -> RawDeviceBuffer { + // SAFETY: these fake handles are used only for bookkeeping; no CUDA API + // sees them and borrowed allocation/context destruction performs no calls. + unsafe { + RawDeviceBuffer::borrowed_from_raw_parts( + 0x1000, + 128, + Context::borrow_raw(context as *mut _, 0, 0), + ) + } +} + +#[test] +fn byte_views_share_extent_identity_and_access_exclusion() { + let raw = metadata(0x2000); + let view = raw.slice_bytes(16, 32).unwrap(); + assert_eq!(view.cu_deviceptr(), raw.cu_deviceptr() + 16); + assert_eq!(view.allocation_id(), raw.allocation_id()); + assert!(raw.slice_bytes(120, 9).is_err()); + assert!(raw.slice_bytes(usize::MAX, 1).is_err()); + assert!(!raw.is_unique()); + let read = raw.acquire_on_stream(1, false, false).unwrap(); + let second_read = view.acquire_on_stream(2, false, false).unwrap(); + assert!(view.acquire_on_stream(2, true, false).is_err()); + drop((read, second_read)); + let write = view.acquire_on_stream(2, true, false).unwrap(); + assert!(raw.acquire_on_stream(1, false, false).is_err()); + assert!(raw.acquire_on_stream(2, true, false).is_ok()); + drop(view); + assert!(raw.is_unique(), "completion leases are not user aliases"); + drop(write); +} + +#[test] +fn same_ordinal_does_not_mean_same_context() { + let raw = metadata(0x2000); + let other = metadata(0x3000); + assert_eq!(raw.device_id(), other.device_id()); + assert_ne!(raw.context(), other.context()); + // SAFETY: this fake borrowed stream is never submitted or destroyed: forget + // it after testing the context check, before Drop could query the driver. + let stream = unsafe { cuda_core::Stream::borrow_raw(0x4000 as *mut _, other.context()) }; + let result = raw.acquire(&stream, false, false); + std::mem::forget(stream); + assert!(result.is_err()); +} + +struct Gate(mpsc::Sender<()>); +impl Gate { + fn arm(stream: &cuda_core::Stream) -> Self { + let (tx, rx) = mpsc::channel(); + stream + .launch_host_function(move || { + let _ = rx.recv_timeout(Duration::from_secs(10)); + }) + .unwrap(); + Self(tx) + } +} +impl Drop for Gate { + fn drop(&mut self) { + let _ = self.0.send(()); + } +} + +#[test] +fn pending_copy_excludes_other_stream_and_outlives_dropped_buffer() { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + let other = context.new_stream().unwrap(); + let mut buffer = DeviceBuffer::from_host(&stream, &[7u32; 32]).unwrap(); + let storage = buffer.storage().clone(); + let gate = Gate::arm(&stream); + buffer.zero_async(&stream).unwrap(); + assert!(storage.acquire(&other, false, false).is_err()); + let before = Instant::now(); + drop(buffer); + assert!( + before.elapsed() < Duration::from_secs(2), + "buffer drop waited for GPU work" + ); + assert!(storage.acquire(&other, false, false).is_err()); + drop(gate); + stream.synchronize().unwrap(); + // SAFETY: the zero fill completed, u32 accepts every representation, and + // this test is the only remaining user of the allocation. + let buffer = unsafe { DeviceBuffer::::from_storage_copy(storage, 32).unwrap() }; + assert_eq!(buffer.to_host_vec(&other).unwrap(), vec![0; 32]); +} + +#[test] +fn stream_drop_is_nonblocking_and_holds_context_until_completion() { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + let gate = Gate::arm(&stream); + let weak = Arc::downgrade(&context); + drop(context); + let before = Instant::now(); + drop(stream); + assert!( + before.elapsed() < Duration::from_secs(2), + "stream drop waited for GPU work" + ); + assert!(weak.upgrade().is_some()); + drop(gate); + let deadline = Instant::now() + Duration::from_secs(5); + while weak.upgrade().is_some() && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(5)); + } + assert!( + weak.upgrade().is_none(), + "retirement leaked an idle stream/context" + ); +} + +#[test] +fn fresh_thread_allocation_binds_its_context() { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + std::thread::spawn(move || { + let buffer = DeviceBuffer::from_host(&stream, &[11u32, 29]).unwrap(); + assert_eq!(buffer.to_host_vec(&stream).unwrap(), [11, 29]); + }) + .join() + .unwrap(); +} + +#[test] +fn empty_and_peeled_buffers_preserve_shared_read_access() { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + let empty = DeviceBuffer::::from_host(&stream, &[]).unwrap(); + assert!(!empty.shared_access_writes()); + let zeroed_empty = DeviceBuffer::::zeroed(&stream, 0).unwrap(); + assert!(!zeroed_empty.shared_access_writes()); + + let words = DeviceBuffer::from_host(&stream, &[1u32, 2, 3]).unwrap(); + let wrapped = words.wrap_elem::>(); + assert!(wrapped.shared_access_writes()); + let peeled = wrapped.peel_elem::(); + assert!(!peeled.shared_access_writes()); +} diff --git a/cuda-core/tests/simt_device_copy_impls.rs b/cuda-core/tests/simt_device_copy_impls.rs deleted file mode 100644 index 2429f7c7ec..0000000000 --- a/cuda-core/tests/simt_device_copy_impls.rs +++ /dev/null @@ -1,35 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -#![cfg_attr(feature = "f16", feature(f16))] - -use std::marker::PhantomData; -use std::mem::MaybeUninit; -use std::num::Wrapping; - -use cuda_core::DeviceCopy; - -fn assert_device_copy() {} - -#[test] -fn device_copy_supports_half_types() { - assert_device_copy::(); - assert_device_copy::(); -} - -#[cfg(feature = "f16")] -#[test] -fn device_copy_supports_native_f16() { - assert_device_copy::(); -} - -#[test] -fn device_copy_covers_core_parity_types() { - // `bool` and `char` are intentionally NOT `DeviceCopy`: they have validity - // holes (only 0/1 for `bool`, only valid Unicode scalars for `char`), so a - // device-written byte outside that set would be UB on readback. Only the - // representation-preserving wrappers below are sound parity additions. - assert_device_copy::>(); - assert_device_copy::>(); - assert_device_copy::>(); -} diff --git a/cuda-core/tests/simt_stream_priority.rs b/cuda-core/tests/stream_priority.rs similarity index 100% rename from cuda-core/tests/simt_stream_priority.rs rename to cuda-core/tests/stream_priority.rs diff --git a/cuda-core/tests/simt_stream_query.rs b/cuda-core/tests/stream_query.rs similarity index 100% rename from cuda-core/tests/simt_stream_query.rs rename to cuda-core/tests/stream_query.rs diff --git a/cuda-core/tests/vmm.rs b/cuda-core/tests/vmm.rs index ac36a4d42c..b27c6333c8 100644 --- a/cuda-core/tests/vmm.rs +++ b/cuda-core/tests/vmm.rs @@ -13,7 +13,7 @@ //! //! Skips when no GPU is present. -use cuda_core::{vmm, Device, IntoResult}; +use cuda_core::{vmm, Context, IntoResult}; use std::mem::MaybeUninit; #[test] @@ -58,7 +58,7 @@ fn vmm_single_gpu_roundtrip() { return; } - let device = Device::new(0).expect("GPU 0 device"); + let device = Context::new(0).expect("GPU 0 device"); let cu_device = device.cu_device(); let granularity = vmm::allocation_granularity(cu_device).expect("allocation granularity query"); diff --git a/cuda-core/tests/vmm_multicast.rs b/cuda-core/tests/vmm_multicast.rs index 0c8cb398d5..29cbbc198b 100644 --- a/cuda-core/tests/vmm_multicast.rs +++ b/cuda-core/tests/vmm_multicast.rs @@ -17,7 +17,7 @@ //! //! Requires an NVLink-switch system (e.g. HGX H100/B200); skips elsewhere. -use cuda_core::{vmm, Device, IntoResult}; +use cuda_core::{vmm, Context, IntoResult}; use std::mem::MaybeUninit; fn gpu_count() -> Result { @@ -37,8 +37,8 @@ fn multicast_two_gpu_team_roundtrip() { return; } - let ctx0 = Device::new(0).expect("GPU 0 device"); - let ctx1 = Device::new(1).expect("GPU 1 device"); + let ctx0 = Context::new(0).expect("GPU 0 device"); + let ctx1 = Context::new(1).expect("GPU 1 device"); let devices = [ctx0.cu_device(), ctx1.cu_device()]; for (i, &dev) in devices.iter().enumerate() { diff --git a/cuda-core/tests/simt_vmm_p2p.rs b/cuda-core/tests/vmm_p2p.rs similarity index 95% rename from cuda-core/tests/simt_vmm_p2p.rs rename to cuda-core/tests/vmm_p2p.rs index 8adca11b01..c8fa45c202 100644 --- a/cuda-core/tests/simt_vmm_p2p.rs +++ b/cuda-core/tests/vmm_p2p.rs @@ -23,9 +23,9 @@ //! //! Requires a system with at least 2 CUDA-capable GPUs with P2P support. -use cuda_core::simt::context::CudaContext; -use cuda_core::simt::peer; -use cuda_core::simt::vmm; +use cuda_core::context::CudaContext; +use cuda_core::peer; +use cuda_core::vmm; use cuda_core::IntoResult; use std::mem::MaybeUninit; @@ -62,7 +62,7 @@ fn vmm_alloc_map_set_access_roundtrip() { let stream = ctx.default_stream(); let pattern: Vec = (0..64).map(|i| (i * 3 + 7) as u8).collect(); unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( va.base(), pattern.as_ptr(), pattern.len(), @@ -74,7 +74,7 @@ fn vmm_alloc_map_set_access_roundtrip() { let mut readback = vec![0u8; 64]; unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( readback.as_mut_ptr(), va.base(), readback.len(), @@ -142,7 +142,7 @@ fn p2p_vmm_cross_gpu_access() { let stream0 = ctx0.default_stream(); ctx0.bind_to_thread().expect("bind ctx0"); unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( va0.base(), pattern.as_ptr(), byte_len, @@ -157,7 +157,7 @@ fn p2p_vmm_cross_gpu_access() { let stream1 = ctx1.default_stream(); ctx1.bind_to_thread().expect("bind ctx1"); unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( readback.as_mut_ptr(), va1.base(), byte_len, diff --git a/cuda-oxide/.cargo/config.toml b/cuda-oxide/.cargo/config.toml index 92b7403412..e660191546 100644 --- a/cuda-oxide/.cargo/config.toml +++ b/cuda-oxide/.cargo/config.toml @@ -5,6 +5,9 @@ oxide = "run --package cargo-oxide --" # Environment variables for all builds [env] +# Unlock `-Z` flags and `#![feature(rustc_private)]` on the pinned stable +# release (`rust-toolchain.toml`). +RUSTC_BOOTSTRAP = "1" # libmathdx installation path (for libmathdx-sys crate) # Required only when using --features libmathdx with mathdx_ffi_test. # IMPORTANT: Replace the placeholder below with your libmathdx install dir, diff --git a/cuda-oxide/.devcontainer/Dockerfile b/cuda-oxide/.devcontainer/Dockerfile index a92200f9ea..fa16226542 100644 --- a/cuda-oxide/.devcontainer/Dockerfile +++ b/cuda-oxide/.devcontainer/Dockerfile @@ -28,6 +28,7 @@ RUN apt-get update \ pkg-config \ python3 \ python3-pip \ + python3-venv \ xz-utils \ && curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key \ | gpg --dearmor -o /usr/share/keyrings/apt.llvm.org.gpg \ diff --git a/cuda-oxide/.devcontainer/README.md b/cuda-oxide/.devcontainer/README.md index 01bcde0c46..12e9b7f612 100644 --- a/cuda-oxide/.devcontainer/README.md +++ b/cuda-oxide/.devcontainer/README.md @@ -6,7 +6,7 @@ This dev container provides the toolchain expected by cuda-oxide: - CUDA Toolkit 13.0 - LLVM 21 with NVPTX support - Clang 21 resource headers for `bindgen` -- Rust `nightly-2026-10-11` with `rust-src`, `rustc-dev`, `rust-analyzer`, +- Rust `1.99.1` with `rust-src`, `rustc-dev`, `rust-analyzer`, `rustfmt` and `clippy`. `llvm-tools` is not preinstalled in the image; rustup adds it on your first cargo command under `cuda-oxide/`, because `cuda-oxide/rust-toolchain.toml` names it. The image ships LLVM 21 and diff --git a/cuda-oxide/.devcontainer/devcontainer.json b/cuda-oxide/.devcontainer/devcontainer.json index 06fa7b4ba7..e77f29ca6e 100644 --- a/cuda-oxide/.devcontainer/devcontainer.json +++ b/cuda-oxide/.devcontainer/devcontainer.json @@ -17,7 +17,7 @@ "upgradePackages": false }, "ghcr.io/devcontainers/features/rust:1": { - "version": "nightly-2026-10-11", + "version": "1.99.1", "profile": "minimal", "components": "rust-src,rustc-dev,rust-analyzer,rustfmt,clippy" } diff --git a/cuda-oxide/.gitattributes b/cuda-oxide/.gitattributes index b38cd1a8bd..f175be0ec5 100644 --- a/cuda-oxide/.gitattributes +++ b/cuda-oxide/.gitattributes @@ -1 +1,3 @@ intrinsics/probes/*.ll whitespace=-blank-at-eof +# Retain upstream fixture whitespace and unified patch context verbatim. +third_party/pliron-mlir-interop/** -whitespace diff --git a/cuda-oxide/.gitignore b/cuda-oxide/.gitignore index 9f14ab24a8..b6225f31e2 100644 --- a/cuda-oxide/.gitignore +++ b/cuda-oxide/.gitignore @@ -5,6 +5,7 @@ .smoketest-logs/ .regress-gemm-sol/ /.layer-target/ +/lean/.lake/ *.tar.gz # GPU compilation artifacts. @@ -15,6 +16,7 @@ .ltoir_arch *.ll *.bc +*.mlir *.cubin *.identity *.gpu diff --git a/cuda-oxide/Cargo.lock b/cuda-oxide/Cargo.lock index cdef05af07..efac0620fa 100644 --- a/cuda-oxide/Cargo.lock +++ b/cuda-oxide/Cargo.lock @@ -196,6 +196,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] [[package]] name = "bytes" @@ -359,6 +373,17 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", + "thiserror 2.0.21", + "trybuild", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -373,7 +398,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -388,7 +413,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -400,18 +425,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -422,7 +449,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", "trybuild", ] @@ -465,6 +494,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cuda-host", @@ -480,27 +510,72 @@ dependencies = [ name = "cuda-oxide-codegen" version = "0.1.0" dependencies = [ + "cuda-oxide-mlir-export", "cuda-target-spec", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", "iket-lower", + "libloading", "libnvvm-sys", "llvm-export", "mir-lower", "mir-transforms", "nvvm-transforms", + "object", "pliron", "ptx-parse", "rustc-demangle", + "serde_json", + "sha2", "tempfile", "thiserror 2.0.21", ] +[[package]] +name = "cuda-oxide-mlir-export" +version = "0.1.0" +dependencies = [ + "cute-layout", + "dialect-cute", + "dialect-mir", + "dialect-nvvm", + "llvm-export", + "pliron", + "pliron-derive", + "pliron-mlir-interop", + "reserved-oxide-symbols", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.21", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-layout", + "trybuild", +] + [[package]] name = "dashmap" version = "6.2.1" @@ -515,6 +590,20 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dialect-cute" +version = "0.1.0" +dependencies = [ + "combine", + "cute-layout", + "dialect-mir", + "linkme", + "once_cell", + "pliron", + "pliron-derive", + "thiserror 2.0.21", +] + [[package]] name = "dialect-iket" version = "0.1.0" @@ -764,6 +853,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", @@ -989,6 +1079,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -1194,6 +1286,17 @@ dependencies = [ "thiserror 2.0.21", ] +[[package]] +name = "pliron-mlir-interop" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d58884474a71b0e9216d69be5c99dd1b945541f38a3acd76046d7c192a991a" +dependencies = [ + "pliron", + "pliron-llvm", + "thiserror 2.0.21", +] + [[package]] name = "prettyplease" version = "0.2.37" diff --git a/cuda-oxide/Cargo.toml b/cuda-oxide/Cargo.toml index 6aeea9929e..e3c5c28647 100644 --- a/cuda-oxide/Cargo.toml +++ b/cuda-oxide/Cargo.toml @@ -1,6 +1,11 @@ [workspace] resolver = "3" members = [ + # CuTe frontend and semantic IR shared by every CuTe backend. + "cute-rs-preview/cute-layout", + "cute-rs-preview/cute-rs", + "cute-rs-preview/ctm-rs", + "cute-rs-preview/dialect-cute", # Core crates "crates/cuda-target-spec", "crates/cuda-intrinsics", @@ -13,6 +18,7 @@ members = [ "crates/cuda-host", "crates/cuda-macros", "crates/llvm-export", + "crates/cuda-oxide-mlir-export", "crates/dialect-mir", "crates/dialect-iket", "crates/dialect-nvvm", @@ -41,10 +47,12 @@ authors = ["Nihal Pasham "] license = "Apache-2.0" repository = "https://github.com/NVIDIA/cuda-rust" -# Note: crates/rustc-codegen-cuda is NOT a workspace member because it requires -# special rustc nightly features and a different build process. +# Note: crates/rustc-codegen-cuda is NOT a workspace member because it links +# against rustc_private (via rustc-dev and RUSTC_BOOTSTRAP=1) and uses a +# separate build process. # Build and run examples with: cargo oxide run -# Examples are in: crates/rustc-codegen-cuda/examples/ +# Examples are in: crates/rustc-codegen-cuda/examples/ and cute-rs-preview/examples/ +# Complete model workloads are in: cute-rs-preview/models/ [workspace.dependencies] # Pliron MLIR framework, pinned to the published 0.20.0 release, retaining @@ -63,6 +71,12 @@ pliron-llvm = { version = "=0.20.0", default-features = false, features = [ # Internal crates cuda-target-spec = { path = "crates/cuda-target-spec" } llvm-export = { path = "crates/llvm-export" } +# Published printer; version-specific compatibility stays in the export adapter. +pliron-mlir-interop = { version = "=0.1.0" } +cuda-oxide-mlir-export = { path = "crates/cuda-oxide-mlir-export" } +cute-layout = { path = "cute-rs-preview/cute-layout" } +ctm-rs = { path = "cute-rs-preview/ctm-rs" } +dialect-cute = { path = "cute-rs-preview/dialect-cute" } dialect-mir = { path = "crates/dialect-mir" } dialect-iket = { path = "crates/dialect-iket" } dialect-ptx = { path = "crates/dialect-ptx" } @@ -110,6 +124,9 @@ ptx-parse = { path = "crates/ptx-parse" } ptx-schedule = { path = "crates/ptx-schedule" } # Common dependencies +# No default features: cuda-device is `no_std` and is compiled for the device, +# and only the marker traits are wanted there, not `std` or the derive macros. +bytemuck = { version = "1", default-features = false } thiserror = "2" linkme = "0.3" anyhow = "1.0" @@ -119,6 +136,7 @@ combine = "4.6.8" log = "0.4" half = "2.7" libloading = "0.8" +object = { version = "0.36", default-features = false, features = ["read", "std", "write"] } sha2 = "0.10" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/cuda-oxide/Justfile b/cuda-oxide/Justfile index ac37ec18d3..2690e9f273 100644 --- a/cuda-oxide/Justfile +++ b/cuda-oxide/Justfile @@ -7,6 +7,10 @@ set windows-shell := ["powershell.exe", "-NoLogo", "-Command"] # Use Bash on Unix set shell := ["bash", "-c"] +# Unlock unstable Cargo (`cargo-features = ["profile-rustflags"]`) and rustc +# (`-Z`, `#![feature(rustc_private)]`) gates on the pinned stable toolchain. +export RUSTC_BOOTSTRAP := "1" + # Format all Rust code (SIMT workspace, codegen, examples) fmt: cargo oxide fmt @@ -64,7 +68,8 @@ test: -p oxide-tile-export -p oxide-tilebc -p cuda-device -p libnvvm-sys -p nvjitlink-sys \ -p cuda-artifact-finalizer -p cargo-oxide \ -p dialect-iket -p iket-lower -p ptx-parse -p dialect-ptx \ - -p ptx-schedule + -p ptx-schedule -p cute-layout -p dialect-cute \ + -p cuda-oxide-mlir-export # `default = []`, but every consumer turns the object features on, and the # default set alone skips the eight ELF emit/extract tests. cargo test -p oxide-artifacts --all-targets --features object @@ -87,7 +92,8 @@ test: # Run the CUDA-toolkit-dependent packages (no driver required) test-cuda: cargo test --all-targets \ - -p cuda-oxide-codegen -p cuda-macros -p cuda-host + -p cuda-oxide-codegen -p cuda-macros -p cuda-host -p cute-rs -p ctm-rs + cargo test -p ctm-rs --all-targets --features host # Mirror unit-tests.yml's third job, `generated-intrinsics`: the three gates a # change under crates/cuda-intrinsics-gen has to pass. Nothing else here ran diff --git a/cuda-oxide/README.md b/cuda-oxide/README.md index 7d65a807f0..e872b2316e 100644 --- a/cuda-oxide/README.md +++ b/cuda-oxide/README.md @@ -31,7 +31,7 @@ Please see [CONTRIBUTING.md](../CONTRIBUTING.md) if you're interested in contrib ```rust use cuda_device::{cuda_module, kernel, thread, DisjointSlice}; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; // Device: generic kernel that applies any function to each element. @@ -95,7 +95,7 @@ lazy `DeviceOperation`, and execution happens when you call `.sync()` or `.await`. ```rust -use cuda_async::simt::device_operation::DeviceOperation; +use cuda_async::device_operation::DeviceOperation; // Assuming `module`, `input`, and `output` come from the cuda-async setup: let factor = 2.5f32; @@ -150,7 +150,7 @@ cargo oxide update ### Requirements - **cargo-oxide** — cargo subcommand that drives the build pipeline (`cargo oxide run`, `build`, `sanitize`, `debug`, etc.) -- **Rust nightly** with `rust-src` and `rustc-dev` and `llvm-tools` components (pinned in `rust-toolchain.toml`) +- **Rust 1.99.1 (stable)** with `rust-src`, `rustc-dev`, and `llvm-tools` components (pinned in `rust-toolchain.toml`) - **CUDA Toolkit** (13.0+, including the cuRAND headers; `libcurand-dev` on Ubuntu). The shared `cuda-bindings` crate loads `libcuda` at run time and needs a CUDA 13.x driver (R580+) - **Clang + libclang dev headers** (`clang-21` / `libclang-common-21-dev`) — needed by `bindgen` when building the host `cuda-bindings` crate - **Linux** (tested on Ubuntu 24.04) @@ -161,17 +161,17 @@ cargo oxide update From `cuda-oxide/`, `cargo oxide` works out of the box via a workspace alias. -For use outside the repo (your own projects), install it with the pinned nightly toolchain: +For use outside the repo (your own projects), install it with the pinned toolchain: ```bash -cargo +nightly-2026-10-11 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide +cargo +1.99.1 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide ``` On first run, `cargo-oxide` will automatically fetch and build the codegen backend. #### Nix (alternative) -If you have Nix with flakes enabled, `nix develop` from `cuda-oxide/` gives you a reproducible shell with CUDA 13, LLVM 22, Clang, and the pinned Rust nightly — no manual apt installs. The shellHook auto-discovers host NVIDIA drivers on NixOS and non-NixOS systems. +If you have Nix with flakes enabled, `nix develop` from `cuda-oxide/` gives you a reproducible shell with CUDA 13, LLVM 22, Clang, and the pinned Rust toolchain - no manual apt installs. The shellHook auto-discovers host NVIDIA drivers on NixOS and non-NixOS systems. ```bash # From the repository root: @@ -187,8 +187,8 @@ nix run 'github:NVIDIA/cuda-rust?dir=cuda-oxide#new' -- my-project ```bash # Toolchain installed automatically via rust-toolchain.toml # Manual install if needed: -rustup toolchain install nightly-2026-10-11 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-10-11 +rustup toolchain install 1.99.1 +rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain 1.99.1 ``` #### CUDA @@ -299,12 +299,15 @@ cargo oxide run gemm_sol_final | Crate | Description | |---------------------|---------------------------------------------------------------------------| | `cuda-device` | Device intrinsics (`thread::*`, `warp::*`, barriers) | +| `cute-layout` | CuTe-style layout algebra and coordinate mapping | +| `cute-rs` | Typed CuTe tensors, tiles, copies, pipelines, and MMA operations | +| `ctm-rs` | Typed instruction-level CUDA primitives | | `cuda-intrinsics` | Generated low-level CUDA intrinsic declarations | | `cuda-host` | Typed module loading, launch helpers, LTOIR loader | | `cuda-macros` | Proc macros (`#[cuda_module]`, `#[kernel]`, `gpu_printf!`) | | `cuda-bindings` | Raw `bindgen` FFI bindings to `cuda.h` (shared with cutile-rs) | -| `cuda-core` | Safe RAII wrappers (`CudaContext`, `DeviceBuffer`, ...); SIMT API under `cuda_core::simt` | -| `cuda-async` | Async layer (`DeviceOperation`, `DeviceBox`, ...); SIMT API under `cuda_async::simt` | +| `cuda-core` | Safe RAII wrappers (`CudaContext`, `DeviceBuffer`, ...); one host API shared by SIMT and Tile | +| `cuda-async` | Async layer (`DeviceOperation`, `DeviceBox`, ...); one runtime for Tile and SIMT | | `libnvvm-sys` | `dlopen` bindings to libNVVM (used by `cuda-host::ltoir`) | | `cuda-target-spec` | Shared CUDA target parsing and recorded LLVM PTX-floor policy | | `nvjitlink-sys` | `dlopen` bindings to nvJitLink (used by `cuda-host::ltoir`) | @@ -321,6 +324,8 @@ cargo oxide run gemm_sol_final | `dialect-iket` | pliron dialect modelling in-kernel event tracing | | `iket-lower` | `dialect-iket` profiles + instrumentation lowering | | `llvm-export` | pliron-llvm shim + textual `.ll` exporter | +| `dialect-cute` | Backend-neutral semantic IR for CuTe tensor programs | +| `cuda-oxide-mlir-export` | CUTLASS 4.7 mappings and profiles; upstream MLIR printing | | `dialect-nvvm` | pliron dialect modelling NVVM intrinsics | | `dialect-ptx` | pliron dialect modelling structured PTX | | `mir-transforms` | Optimization passes over the MIR dialect (loop unroll, ...) | diff --git a/cuda-oxide/THIRD_PARTY_NOTICES b/cuda-oxide/THIRD_PARTY_NOTICES index 9d953a7f40..df81ba851f 100644 --- a/cuda-oxide/THIRD_PARTY_NOTICES +++ b/cuda-oxide/THIRD_PARTY_NOTICES @@ -10,6 +10,57 @@ crates that cuda-oxide merely depends on through Cargo are not redistributed here and are therefore not listed below; their licenses are recorded in `dependency-licenses.csv` and enforced by `deny.toml`. +-------------------------------------------------------------------------------- +CUTLASS CuTeDSL block-scaled GEMM reference +-------------------------------------------------------------------------------- + +Location in cuda-oxide: cute-rs-preview/examples/blockscale_gemm_cute/reference/ +Upstream project: https://github.com/NVIDIA/cutlass +Upstream revision: v4.6.2 (6c65a175668952f09bcbf66cb97a8de1b734b4a0) +Copyright: (c) 2026 NVIDIA CORPORATION & AFFILIATES +License: BSD-3-Clause + +The unmodified files `dense_blockscaled_gemm_persistent_cooperative.py` and +`blockscaled_gemm_dispatch.py` come from +`examples/python/CuTeDSL/cute/blackwell_geforce/kernel/blockscaled_gemm/`. +Each retains the complete upstream copyright, license conditions, and disclaimer. +The adjacent `benchmark.py` and `cpp/` reference are cuda-oxide code under +Apache-2.0; they are not part of the vendored CUTLASS source. + +-------------------------------------------------------------------------------- +CUTLASS CuTeDSL FlashAttention v2 reference +-------------------------------------------------------------------------------- + +Location in cuda-oxide: cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/flash_attention_v2.py +Upstream project: https://github.com/NVIDIA/cutlass +Upstream revision: v4.7.0 (dcf215af68a2d08d305076c152a06f201728cd53) +Copyright: (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES +License: BSD-3-Clause + +The unmodified file comes from +`examples/python/CuTeDSL/cute/ampere/kernel/attention/flash_attention_v2.py`. +It retains the complete upstream copyright, license conditions, and disclaimer. +The adjacent launch, validation and benchmark helpers are cuda-oxide code under +Apache-2.0. The separately authored CuTeDSL decode reference is also cuda-oxide +code under Apache-2.0; it is not part of this vendored CUTLASS source. + +-------------------------------------------------------------------------------- +CUTLASS CuTeDSL datacenter Blackwell attention references +-------------------------------------------------------------------------------- + +Locations in cuda-oxide: cute-rs-preview/examples/sm100_fmha_prefill_cute/reference/ + cute-rs-preview/examples/sm100_mla_decode_fp16_cute/reference/ +Upstream project: https://github.com/NVIDIA/cutlass +Upstream revision: v4.7.0 (dcf215af68a2d08d305076c152a06f201728cd53) +Copyright: (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES +License: BSD-3-Clause + +The unchanged FMHA source, fmha_helpers.py and helpers/__init__.py, and the +unchanged MLA decode source and mla_helpers.py retain their complete upstream +BSD notices. Each reference/SOURCE.md records original paths and SHA-256 hashes. +The adjacent Rust examples and Python comparison adapters are cuda-oxide code +under Apache-2.0. + -------------------------------------------------------------------------------- 1. Rustlantis -------------------------------------------------------------------------------- diff --git a/cuda-oxide/crates/cargo-oxide/README.md b/cuda-oxide/crates/cargo-oxide/README.md index 2ed176b4ed..4ffdbedffc 100644 --- a/cuda-oxide/crates/cargo-oxide/README.md +++ b/cuda-oxide/crates/cargo-oxide/README.md @@ -10,10 +10,10 @@ Replaces the previous `xtask` pattern with a proper cargo subcommand that works **External users**: -Install with the project's pinned nightly toolchain: +Install with the project's pinned toolchain: ```bash -cargo +nightly-2026-10-11 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide +cargo +1.99.1 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide ``` On first run, `cargo-oxide` will automatically fetch and build the codegen backend if it's not already available. @@ -134,8 +134,12 @@ final-output mode, so it cannot be combined with `--emit-nvvm-ir`, ### `cargo oxide list [--json]` -Lists every example bundled with the cuda-oxide workspace, including its short -description and any hardware or toolkit requirements documented in its README. +Lists every example and model bundled with the cuda-oxide workspace, including +its short description and any hardware or toolkit requirements documented in +its README. Discovery includes `crates/rustc-codegen-cuda/examples/`, +`cute-rs-preview/examples/`, and `cute-rs-preview/models/` within the cuda-oxide component. +Names must be unique across these directories; `run` and `build` accept model +names in the same way as example names. ```bash cargo oxide list @@ -448,7 +452,7 @@ merged repository it also formats the parent shared-host workspace. With ### `cargo oxide doctor` -Validates that your environment is correctly set up: Rust nightly toolchain, +Validates that your environment is correctly set up: pinned Rust toolchain, CUDA headers (`cuda.h`), CUDA toolkit (`nvcc`, libNVVM, nvJitLink, libdevice), LLVM (`llc`), clang/libclang, the NVIDIA driver / GPU, and the codegen backend `.so`. Optional probes cover `cuda-gdb` (for `debug`) and @@ -506,11 +510,11 @@ Cargo.lock cuda-device = git+https://github.com/NVIDIA/cuda-rust.git# The kernels compile against the crates at that commit and the backend that lowers them is built from the same commit, so the two cannot drift. Changing the dependency (`cargo update`, a new `rev = ...`) rebuilds the cache on the -next build. The checkout's `rust-toolchain.toml` names the nightly that commit -needs; if the project pins a different one, the build stops before compiling -and tells you which channel to set. A path dependency on a local checkout -builds in place, like step 3. Projects with no cuda-oxide dependency fall back -to a one-time clone of `main`. The backend's build tree lives at +next build. The checkout's `rust-toolchain.toml` names the toolchain channel +that commit needs; if the project pins a different one, the build stops before +compiling and tells you which channel to set. A path dependency on a local +checkout builds in place, like step 3. Projects with no cuda-oxide dependency +fall back to a one-time clone of `main`. The backend's build tree lives at `~/.cargo/cuda-oxide/target` (shared by every commit's builds; delete it to reclaim the space). diff --git a/cuda-oxide/crates/cargo-oxide/src/backend.rs b/cuda-oxide/crates/cargo-oxide/src/backend.rs index 2e5ce0d5ff..ffee12453e 100644 --- a/cuda-oxide/crates/cargo-oxide/src/backend.rs +++ b/cuda-oxide/crates/cargo-oxide/src/backend.rs @@ -813,7 +813,7 @@ fn report_unhealable_toolchain_mismatch(current: &str, recorded: &str) { toolchain_identity_line(recorded) ); eprintln!( - "Your project's rust-toolchain.toml pins a different nightly than the \ + "Your project's rust-toolchain.toml pins a different toolchain than the \ cuda-oxide backend source it depends on. Align the pins (and run \ `cargo oxide update` after changing them), or point CUDA_OXIDE_BACKEND \ at a backend built with your project's toolchain." @@ -1056,6 +1056,7 @@ fn backend_build_command_in( cmd.env("CARGO_ENCODED_RUSTFLAGS", ""); cmd.env_remove("RUSTFLAGS"); cmd.env_remove("CARGO_BUILD_RUSTFLAGS"); + cmd.env("RUSTC_BOOTSTRAP", "1"); if let Some(path) = rustc_lib_path { cmd.env("LIBRARY_PATH", build_library_path(path)); diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/clean.rs b/cuda-oxide/crates/cargo-oxide/src/commands/clean.rs index 1e6683ddc0..0fea004e34 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/clean.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/clean.rs @@ -65,44 +65,46 @@ fn clean_workspace(ctx: &Context, summary: &mut CleanSummary) -> Result<(), Stri summary.removed_directories += 1; } + if ctx.oxide_workspace_root() != ctx.workspace_root + && remove_local_target(ctx.oxide_workspace_root())? + { + summary.removed_directories += 1; + } if remove_local_target(&ctx.codegen_crate)? { summary.removed_directories += 1; } - let entries = std::fs::read_dir(&ctx.examples_dir).map_err(|error| { - format!( - "could not read examples directory {}: {error}", - ctx.examples_dir.display() - ) - })?; - let mut example_dirs = Vec::new(); - - for entry in entries { - let entry = entry.map_err(|error| { + for root in example_roots(ctx) { + let entries = std::fs::read_dir(&root).map_err(|error| { format!( - "could not read an entry in {}: {error}", - ctx.examples_dir.display() + "could not read examples directory {}: {error}", + root.display() ) })?; - let file_type = entry.file_type().map_err(|error| { - format!( - "could not inspect example entry {}: {error}", - entry.path().display() - ) - })?; + for entry in entries { + let entry = entry.map_err(|error| { + format!("could not read an entry in {}: {error}", root.display()) + })?; - if !file_type.is_dir() { - continue; - } + let file_type = entry.file_type().map_err(|error| { + format!( + "could not inspect example entry {}: {error}", + entry.path().display() + ) + })?; + + if !file_type.is_dir() { + continue; + } - let example_dir = entry.path(); - if example_dir.join("Cargo.toml").is_file() { - example_dirs.push(example_dir); + let example_dir = entry.path(); + if example_dir.join("Cargo.toml").is_file() { + example_dirs.push(example_dir); + } } } - example_dirs.sort(); for example_dir in example_dirs { diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/codegen_env.rs b/cuda-oxide/crates/cargo-oxide/src/commands/codegen_env.rs index 4232a7bd96..1442a2ce61 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/codegen_env.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/codegen_env.rs @@ -4,6 +4,7 @@ */ use crate::backend; +use crate::cutlass_toolchain; use std::path::Path; use std::process::Command; @@ -465,6 +466,7 @@ fn apply_codegen_rustflags( append_full_debug_rustflags(&mut encoded, cmd, inherited_debug.as_deref()); cmd.env("CARGO_ENCODED_RUSTFLAGS", encoded) + .env("RUSTC_BOOTSTRAP", "1") .env_remove("RUSTFLAGS"); apply_profile_pins(cmd, profile); } @@ -552,6 +554,7 @@ pub fn has_configured_arch(ctx: &Context, cli_arch: Option<&str>) -> bool { } pub(super) fn apply_config_env(cmd: &mut Command, ctx: &Context) { + cmd.env("RUSTC_BOOTSTRAP", "1"); for (key, value) in &ctx.config.env { if matches!(key.as_str(), "RUSTFLAGS" | "CARGO_ENCODED_RUSTFLAGS") { continue; @@ -563,6 +566,61 @@ pub(super) fn apply_config_env(cmd: &mut Command, ctx: &Context) { cmd.env(key, value); } } + apply_managed_cutlass_compiler(cmd, ctx); +} + +fn apply_managed_cutlass_compiler(cmd: &mut Command, ctx: &Context) { + let backend = std::env::var_os(DEVICE_BACKEND_ENV) + .map(|value| value.as_encoded_bytes().to_vec()) + .or_else(|| { + project_config_env(ctx, DEVICE_BACKEND_ENV).map(|value| value.as_bytes().to_vec()) + }); + if !backend.as_deref().is_some_and(selector_is_cutlass_mlir) { + return; + } + + let explicitly_configured = std::env::var_os(cutlass_toolchain::CUTLASS_COMPILER_ENV).is_some() + || project_config_env(ctx, cutlass_toolchain::CUTLASS_COMPILER_ENV).is_some(); + if explicitly_configured { + return; + } + + match cutlass_toolchain::resolve_official() { + Ok(Some(paths)) => apply_resolved_cutlass_compiler(cmd, Some(&paths), false), + Ok(None) => { + eprintln!( + "Error: CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir needs the managed CUTLASS compiler" + ); + eprintln!("Install the pinned official release with:"); + eprintln!(" cargo oxide toolchain install cutlass"); + eprintln!( + "Or explicitly set {}.", + cutlass_toolchain::CUTLASS_COMPILER_ENV + ); + std::process::exit(2); + } + Err(error) => { + eprintln!("Error: managed CUTLASS compiler installation is invalid: {error}"); + eprintln!("Repair it with `cargo oxide toolchain install cutlass`."); + std::process::exit(2); + } + } +} + +pub(super) fn apply_resolved_cutlass_compiler( + cmd: &mut Command, + resolved: Option<&cutlass_toolchain::ResolvedCutlass>, + explicitly_configured: bool, +) { + if explicitly_configured { + return; + } + if let Some(paths) = resolved { + cmd.env( + cutlass_toolchain::CUTLASS_COMPILER_ENV, + &paths.compiler_library, + ); + } } pub(super) fn apply_common_codegen_env( diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/context.rs b/cuda-oxide/crates/cargo-oxide/src/commands/context.rs index 66fa592d43..2e45a019bc 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/context.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/context.rs @@ -41,6 +41,20 @@ pub struct Context { pub config: OxideConfig, } +impl Context { + /// The cuda-oxide component containing the selected backend. The checkout + /// root may instead be the parent stable workspace in the monorepo. + pub(super) fn oxide_workspace_root(&self) -> &Path { + if !self.is_workspace { + return &self.workspace_root; + } + self.codegen_crate + .parent() + .and_then(Path::parent) + .expect("workspace backend must live under crates/rustc-codegen-cuda") + } +} + /// Resolve the workspace root and backend, or exit with a helpful error. /// /// Supports two modes: @@ -489,26 +503,56 @@ fn parse_config_env( Ok(env) } -/// Resolve an example name to its directory path, or exit with a list of -/// available examples if not found. -pub(super) fn resolve_example_dir(ctx: &Context, example: &str) -> PathBuf { - let example_dir = ctx.examples_dir.join(example); - if !example_dir.exists() { - eprintln!("Error: Example not found: {}", example_dir.display()); - eprintln!(); - eprintln!("Available examples:"); - if let Ok(entries) = std::fs::read_dir(&ctx.examples_dir) { - let mut names: Vec<_> = entries - .flatten() - .filter(|e| e.path().is_dir()) - .map(|e| e.file_name().to_string_lossy().to_string()) - .collect(); - names.sort(); - for name in names { - eprintln!(" - {}", name); +/// Example and model roots belonging to the selected cuda-oxide component. +pub(super) fn example_roots(ctx: &Context) -> Vec { + let mut roots = vec![ctx.examples_dir.clone()]; + if ctx.is_workspace { + for relative in ["cute-rs-preview/examples", "cute-rs-preview/models"] { + let root = ctx.oxide_workspace_root().join(relative); + if root.is_dir() { + roots.push(root); } } + } + roots +} + +pub(super) fn resolve_example_dir(ctx: &Context, example: &str) -> PathBuf { + let roots = example_roots(ctx); + let matches: Vec<_> = roots + .iter() + .map(|root| root.join(example)) + .filter(|path| path.is_dir()) + .collect(); + if let [example_dir] = matches.as_slice() { + return example_dir.clone(); + } + if matches.len() > 1 { + eprintln!("Error: Example name is ambiguous: {example}"); + for path in matches { + eprintln!(" - {}", path.display()); + } + eprintln!("Example names must be unique across all configured roots."); std::process::exit(1); } - example_dir + + eprintln!("Error: Example not found: {example}"); + eprintln!(); + eprintln!("Available examples:"); + let mut names: Vec<_> = roots + .iter() + .filter_map(|root| std::fs::read_dir(root).ok()) + .flat_map(|entries| { + entries + .flatten() + .filter(|entry| entry.path().is_dir()) + .map(|entry| entry.file_name().to_string_lossy().to_string()) + }) + .collect(); + names.sort(); + names.dedup(); + for name in names { + eprintln!(" - {name}"); + } + std::process::exit(1); } diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/doctor.rs b/cuda-oxide/crates/cargo-oxide/src/commands/doctor.rs index 143d7e979b..cec6417c35 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/doctor.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/doctor.rs @@ -360,7 +360,7 @@ pub(super) fn backend_source_check( /// Validate the development environment. /// -/// Checks for: Rust nightly toolchain, `rust-toolchain.toml`, the codegen +/// Checks for: Rust toolchain (`rustc`), `rust-toolchain.toml`, the codegen /// backend `.so` (informational), CUDA headers (`cuda.h`), CUDA toolkit /// (`nvcc`, libNVVM, nvJitLink, libdevice), LLVM (`llc`), clang/libclang, /// the NVIDIA driver / GPU (informational), and optionally `cuda-gdb` / @@ -379,17 +379,11 @@ pub fn doctor(ctx: &Context) { let mut ok = true; // 1. Rust toolchain - print!("Rust nightly toolchain... "); + print!("Rust toolchain... "); match Command::new("rustc").args(["--version"]).output() { Ok(output) if output.status.success() => { let version = String::from_utf8_lossy(&output.stdout); - let version = version.trim(); - if version.contains("nightly") { - println!("✓ {}", version); - } else { - println!("✗ expected nightly, got: {}", version); - ok = false; - } + println!("✓ {}", version.trim()); } _ => { println!("✗ rustc not found"); diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/examples_list.rs b/cuda-oxide/crates/cargo-oxide/src/commands/examples_list.rs index 7eff89cb16..680bc2e624 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/examples_list.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/examples_list.rs @@ -4,7 +4,7 @@ */ use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use super::*; @@ -32,11 +32,13 @@ pub fn list_examples(ctx: &Context, json: bool) { if !ctx.is_workspace { eprintln!("Error: `cargo oxide list` must be run from inside a cuda-oxide checkout."); eprintln!(); - eprintln!("The command lists examples under crates/rustc-codegen-cuda/examples/."); + eprintln!( + "The command lists the standard and CuTe examples and models bundled with the workspace." + ); std::process::exit(1); } - let examples = discover_examples(&ctx.examples_dir).unwrap_or_else(|error| { + let examples = discover_examples_from_roots(&example_roots(ctx)).unwrap_or_else(|error| { eprintln!("Error: {error}"); std::process::exit(1); }); @@ -53,6 +55,26 @@ pub fn list_examples(ctx: &Context, json: bool) { print!("{output}"); } +/// Discover all configured example roots and reject duplicate names. +pub(super) fn discover_examples_from_roots(roots: &[PathBuf]) -> Result, String> { + let mut examples = Vec::new(); + for root in roots { + examples.extend(discover_examples(root)?); + } + + examples.sort_by(|left, right| left.name.cmp(&right.name)); + if let Some(pair) = examples + .windows(2) + .find(|pair| pair[0].name == pair[1].name) + { + return Err(format!( + "example name `{}` exists in more than one example root; names must be unique", + pair[0].name + )); + } + Ok(examples) +} + pub(super) fn discover_examples(examples_dir: &Path) -> Result, String> { let entries = fs::read_dir(examples_dir).map_err(|error| { format!( diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/fingerprint.rs b/cuda-oxide/crates/cargo-oxide/src/commands/fingerprint.rs index 03a9e4ce24..c37e7c3435 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/fingerprint.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/fingerprint.rs @@ -31,7 +31,7 @@ pub(super) fn passthrough_codegen_fingerprint( owner_filter, target_arch, materialization, - &inherited_process_env(), + &inherited_process_env_with_cutlass_compiler_identity(ctx), ) } @@ -163,7 +163,7 @@ pub(super) fn sanitize_codegen_fingerprint( detected_device_arch, ptx_dir, materialization, - &inherited_process_env(), + &inherited_process_env_with_cutlass_compiler_identity(ctx), ) } diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/fmt.rs b/cuda-oxide/crates/cargo-oxide/src/commands/fmt.rs index c1a4273efb..00bb589f16 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/fmt.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/fmt.rs @@ -28,11 +28,11 @@ use super::*; pub fn format_all(ctx: &Context, check: bool) { let mode = if check { "Checking" } else { "Formatting" }; let mut failed = false; + let oxide_root = ctx.oxide_workspace_root(); // In the merged repository, shared host crates live in the stable parent // workspace. Standalone cuda-oxide checkouts have no such parent manifest. - if let Some(shared_root) = ctx - .workspace_root + if let Some(shared_root) = oxide_root .parent() .filter(|root| root.join("Cargo.toml").is_file() && root.join("cuda-core").is_dir()) { @@ -43,7 +43,7 @@ pub fn format_all(ctx: &Context, check: bool) { } println!("📦 {} CUDA Oxide workspace...", mode); - if !run_cargo_fmt(&ctx.workspace_root, check) { + if !run_cargo_fmt(oxide_root, check) { failed = true; } @@ -74,13 +74,15 @@ pub fn format_all(ctx: &Context, check: bool) { // * `--all` also formats an example's path dependencies, which means // re-formatting the large shared workspaces once per example. let mut manifests = Vec::new(); - collect_example_manifests(&ctx.examples_dir, &mut manifests); + for root in example_roots(ctx) { + collect_example_manifests(&root, &mut manifests); + } manifests.sort(); for manifest in &manifests { let label = manifest .parent() - .and_then(|dir| dir.strip_prefix(&ctx.examples_dir).ok()) + .and_then(|dir| dir.strip_prefix(oxide_root).ok()) .unwrap_or(Path::new(".")) .display() .to_string(); @@ -142,6 +144,7 @@ fn run_cargo_fmt_manifest(manifest: &Path, check: bool) -> bool { } fn run_fmt_command(mut cmd: Command) -> bool { + cmd.env("RUSTC_BOOTSTRAP", "1"); match cmd.status() { Ok(status) => status.success(), Err(e) => { diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/host_cargo.rs b/cuda-oxide/crates/cargo-oxide/src/commands/host_cargo.rs index 36b1f36724..179ad7bb6f 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/host_cargo.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/host_cargo.rs @@ -367,6 +367,7 @@ fn cargo_package_manifest_matches(package: &serde_json::Value, manifest_path: &P fn cargo_metadata(manifest_dir: &Path) -> Result { let output = Command::new("cargo") + .env("RUSTC_BOOTSTRAP", "1") .args(["metadata", "--format-version=1", "--no-deps"]) .current_dir(manifest_dir) .output() diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs b/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs index 6e3496fe10..6af3ed318a 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs @@ -4,9 +4,11 @@ */ use std::collections::BTreeMap; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::process::Command; +use crate::cutlass_toolchain; + use super::*; /// Device debug-information policy requested on the command line. @@ -177,6 +179,79 @@ pub(super) fn inherited_process_env() -> BTreeMap> { .collect() } +pub(super) const DEVICE_BACKEND_ENV: &str = "CUDA_OXIDE_DEVICE_BACKEND"; +pub(super) const CUTLASS_DIGEST_FINGERPRINT_ENV: &str = + "CUDA_OXIDE_INTERNAL_CUTLASS_COMPILER_SHA256"; + +/// Include the selected CUTLASS compiler and its content identity in the +/// semantic environment seen by Cargo. +pub(super) fn inherited_process_env_with_cutlass_compiler_identity( + ctx: &Context, +) -> BTreeMap> { + let mut inherited = inherited_process_env(); + let inherited_compiler_path = + std::env::var_os(cutlass_toolchain::CUTLASS_COMPILER_ENV).map(PathBuf::from); + augment_env_with_cutlass_compiler_identity( + ctx, + &mut inherited, + inherited_compiler_path.as_deref(), + ); + inherited +} + +pub(super) fn augment_env_with_cutlass_compiler_identity( + ctx: &Context, + inherited: &mut BTreeMap>, + inherited_compiler_path: Option<&Path>, +) { + let backend = inherited + .get(DEVICE_BACKEND_ENV) + .map(Vec::as_slice) + .or_else(|| project_config_env(ctx, DEVICE_BACKEND_ENV).map(str::as_bytes)); + if !backend.is_some_and(selector_is_cutlass_mlir) { + return; + } + + let explicit_compiler = if inherited.contains_key(cutlass_toolchain::CUTLASS_COMPILER_ENV) { + inherited_compiler_path.map(Path::to_path_buf) + } else { + project_config_env(ctx, cutlass_toolchain::CUTLASS_COMPILER_ENV).map(PathBuf::from) + }; + if inherited.contains_key(cutlass_toolchain::CUTLASS_COMPILER_ENV) + || explicit_compiler.is_some() + { + let digest = match explicit_compiler.as_deref() { + Some(path) => cutlass_toolchain::sha256_file(path) + .unwrap_or_else(|error| format!("unavailable:{error}")), + None => "unavailable:non-unicode inherited compiler path".to_string(), + }; + inherited.insert( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_string(), + digest.into_bytes(), + ); + return; + } + + if let Ok(Some(paths)) = cutlass_toolchain::resolve_official() { + inherited.insert( + cutlass_toolchain::CUTLASS_COMPILER_ENV.to_string(), + paths + .compiler_library + .as_os_str() + .as_encoded_bytes() + .to_vec(), + ); + inherited.insert( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_string(), + paths.compiler_library_sha256.as_bytes().to_vec(), + ); + } +} + +pub(super) fn selector_is_cutlass_mlir(value: &[u8]) -> bool { + std::str::from_utf8(value).is_ok_and(|value| value.trim() == "cutlass-mlir") +} + fn cargo_passthrough_command( ctx: &Context, cargo_subcommand: CargoPassthroughSubcommand, diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs b/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs index db77c02d29..bcb0706940 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs @@ -14,10 +14,14 @@ use super::*; const GIT_REPO: &str = "https://github.com/NVIDIA/cuda-rust.git"; const RUST_TOOLCHAIN_TOML: &str = r#"[toolchain] -channel = "nightly-2026-10-11" +channel = "1.99.1" components = ["rust-src", "rustc-dev", "rust-analyzer", "clippy", "rustfmt", "llvm-tools"] "#; +const CARGO_CONFIG_TOML: &str = r#"[env] +RUSTC_BOOTSTRAP = "1" +"#; + const SCAFFOLD_GITIGNORE_EXTRA: &[&str] = &[ "/target/", "**/*.bc", // bitcode leftovers not in the clean suffix list @@ -52,6 +56,7 @@ pub(super) fn scaffold_gitignore() -> String { #[derive(Clone, Debug, Eq, PartialEq)] pub(super) struct ScaffoldFiles { pub(super) cargo_toml: String, + pub(super) cargo_config_toml: String, rust_toolchain_toml: String, pub(super) gitignore: String, pub(super) readme: String, @@ -146,9 +151,9 @@ cuda-core = {{ git = "{GIT_REPO}" }} fn scaffold_main_rs(async_mode: bool) -> String { if async_mode { - r#"use cuda_async::simt::device_context::{init_device_contexts, with_cuda_context}; -use cuda_async::simt::device_operation::DeviceOperation; -use cuda_core::simt::LaunchConfig; + r#"use cuda_async::device_context::{init_device_contexts, with_cuda_context}; +use cuda_async::device_operation::DeviceOperation; +use cuda_core::LaunchConfig; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; @@ -168,8 +173,8 @@ mod kernels { #[tokio::main] async fn main() -> Result<(), Box> { - use cuda_async::simt::device_box::DeviceBox; - use cuda_core::simt::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; + use cuda_async::device_box::DeviceBox; + use cuda_core::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; use std::mem; init_device_contexts(0, 1)?; @@ -190,9 +195,9 @@ async fn main() -> Result<(), Box> { memcpy_htod_async(b, b_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); stream.synchronize().unwrap(); ( - DeviceBox::<[f32]>::from_raw_parts(a, N, 0), - DeviceBox::<[f32]>::from_raw_parts(b, N, 0), - DeviceBox::<[f32]>::from_raw_parts(c, N, 0), + DeviceBox::<[f32]>::from_raw_parts(a, N, stream.clone()), + DeviceBox::<[f32]>::from_raw_parts(b, N, stream.clone()), + DeviceBox::<[f32]>::from_raw_parts(c, N, stream.clone()), ) } })?; @@ -299,6 +304,7 @@ fn main() -> Result<(), Box> { pub(super) fn scaffold_files(name: &str, async_mode: bool) -> ScaffoldFiles { ScaffoldFiles { cargo_toml: scaffold_cargo_toml(name, async_mode), + cargo_config_toml: CARGO_CONFIG_TOML.to_string(), rust_toolchain_toml: RUST_TOOLCHAIN_TOML.to_string(), gitignore: scaffold_gitignore(), readme: scaffold_readme(name, async_mode), @@ -319,10 +325,17 @@ pub fn scaffold_new(name: &str, async_mode: bool) { eprintln!("Error creating directory: {}", e); std::process::exit(1); }); + let cargo_dir = project_dir.join(".cargo"); + std::fs::create_dir_all(&cargo_dir).unwrap_or_else(|e| { + eprintln!("Error creating directory: {}", e); + std::process::exit(1); + }); let files = scaffold_files(name, async_mode); std::fs::write(project_dir.join("Cargo.toml"), files.cargo_toml) .expect("Failed to write Cargo.toml"); + std::fs::write(cargo_dir.join("config.toml"), files.cargo_config_toml) + .expect("Failed to write .cargo/config.toml"); std::fs::write( project_dir.join("rust-toolchain.toml"), files.rust_toolchain_toml, diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/tests.rs b/cuda-oxide/crates/cargo-oxide/src/commands/tests.rs index bf0551deb2..4fce532545 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/tests.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/tests.rs @@ -9,6 +9,7 @@ use std::path::{Path, PathBuf}; use std::process::Command; use super::*; +use crate::cutlass_toolchain; use std::ffi::OsStr; /// The concatenated sources of every module under `commands/`, standing in @@ -69,6 +70,40 @@ fn is_sha256(value: &str) -> bool { .all(|byte| byte.is_ascii_digit() || matches!(byte, b'a'..=b'f')) } +#[test] +fn managed_cutlass_compiler_sets_only_the_library_env() { + let resolved = cutlass_toolchain::ResolvedCutlass { + install_dir: PathBuf::from("/managed/cutlass"), + compiler_library: PathBuf::from("/managed/cutlass/lib/libCutlassCompiler.so"), + compiler_library_sha256: "fixture-digest".into(), + }; + let mut cmd = Command::new("cargo"); + apply_resolved_cutlass_compiler(&mut cmd, Some(&resolved), false); + assert_eq!( + command_env(&cmd, cutlass_toolchain::CUTLASS_COMPILER_ENV).as_deref(), + Some("/managed/cutlass/lib/libCutlassCompiler.so") + ); +} + +#[test] +fn explicit_cutlass_compiler_outranks_managed_resolution() { + let resolved = cutlass_toolchain::ResolvedCutlass { + install_dir: PathBuf::from("/managed/cutlass"), + compiler_library: PathBuf::from("/managed/cutlass/lib/libCutlassCompiler.so"), + compiler_library_sha256: "fixture-digest".into(), + }; + let mut cmd = Command::new("cargo"); + cmd.env( + cutlass_toolchain::CUTLASS_COMPILER_ENV, + "/explicit/libCutlassCompiler.so", + ); + apply_resolved_cutlass_compiler(&mut cmd, Some(&resolved), true); + assert_eq!( + command_env(&cmd, cutlass_toolchain::CUTLASS_COMPILER_ENV).as_deref(), + Some("/explicit/libCutlassCompiler.so") + ); +} + fn strings(args: &[&str]) -> Vec { args.iter().map(|arg| arg.to_string()).collect() } @@ -516,48 +551,69 @@ edition = "2024" } #[test] -fn workspace_clean_removes_root_backend_and_example_targets() { - let root = unique_temp_dir("cargo_oxide_clean_workspace"); - let codegen_crate = root.join("crates/rustc-codegen-cuda"); - let examples_dir = codegen_crate.join("examples"); - let example_dir = examples_dir.join("demo"); - - std::fs::create_dir_all(root.join("target/debug")).unwrap(); - std::fs::create_dir_all(codegen_crate.join("target/debug")).unwrap(); - std::fs::create_dir_all(example_dir.join("target/debug")).unwrap(); - - std::fs::write( - example_dir.join("Cargo.toml"), - r#" -[package] -name = "demo" -version = "0.1.0" -edition = "2024" -"#, - ) - .unwrap(); - - std::fs::write(example_dir.join("demo.ptx"), b"generated").unwrap(); - - let ctx = Context { - workspace_root: root.clone(), - codegen_crate: codegen_crate.clone(), - examples_dir, - backend_so: root.join("unused-backend.so"), - is_workspace: true, - config: OxideConfig::default(), - }; - - let summary = clean_context(&ctx).unwrap(); - - assert_eq!(summary.removed_directories, 3); - assert_eq!(summary.removed_files, 1); - assert!(!root.join("target").exists()); - assert!(!codegen_crate.join("target").exists()); - assert!(!example_dir.join("target").exists()); - assert!(!example_dir.join("demo.ptx").exists()); - - std::fs::remove_dir_all(root).unwrap(); +fn workspace_clean_covers_example_and_model_roots_in_flat_and_nested_layouts() { + for nested in [false, true] { + let root = unique_temp_dir("cargo_oxide_clean_workspace"); + let component = if nested { + root.join("cuda-oxide") + } else { + root.clone() + }; + let codegen_crate = component.join("crates/rustc-codegen-cuda"); + let examples_dir = codegen_crate.join("examples"); + let example_dir = examples_dir.join("demo"); + let cute_root = component.join("cute-rs-preview/examples"); + let cute_dir = cute_root.join("tiled"); + let models_root = component.join("cute-rs-preview/models"); + let model_dir = models_root.join("decoder"); + for dir in [ + &root, + &component, + &codegen_crate, + &example_dir, + &cute_dir, + &model_dir, + ] { + std::fs::create_dir_all(dir.join("target/debug")).unwrap(); + } + write_list_example(&examples_dir, "demo", None, None); + write_list_example(&cute_root, "tiled", None, None); + write_list_example(&models_root, "decoder", None, None); + std::fs::write(example_dir.join("demo.ptx"), b"generated").unwrap(); + std::fs::write(cute_dir.join("tiled.cubin"), b"generated").unwrap(); + std::fs::write(model_dir.join("decoder.cubin"), b"generated").unwrap(); + std::fs::write(cute_dir.join("other.cubin"), b"preserve").unwrap(); + let ctx = Context { + workspace_root: root.clone(), + codegen_crate: codegen_crate.clone(), + examples_dir, + backend_so: root.join("unused-backend.so"), + is_workspace: true, + config: OxideConfig::default(), + }; + let summary = clean_context(&ctx).unwrap(); + assert_eq!(summary.removed_directories, if nested { 6 } else { 5 }); + assert_eq!(summary.removed_files, 3); + for dir in [ + &root, + &component, + &codegen_crate, + &example_dir, + &cute_dir, + &model_dir, + ] { + assert!(!dir.join("target").exists()); + } + assert!(!example_dir.join("demo.ptx").exists()); + assert!(!cute_dir.join("tiled.cubin").exists()); + assert!(!model_dir.join("decoder.cubin").exists()); + assert_eq!( + std::fs::read(cute_dir.join("other.cubin")).unwrap(), + b"preserve" + ); + assert_eq!(clean_context(&ctx).unwrap(), CleanSummary::default()); + std::fs::remove_dir_all(root).unwrap(); + } } #[cfg(unix)] @@ -3556,6 +3612,117 @@ fn passthrough_fingerprint_tracks_output_affecting_settings() { ); } +#[test] +fn passthrough_fingerprint_tracks_the_managed_cutlass_library_digest() { + let ctx = test_context(OxideConfig::default()); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: Some("sm_120a"), + features: None, + cargo_target_dir: None, + device_codegen_crate: None, + device_cfgs: &[], + no_fmad: false, + unchecked_indexing: false, + materialize_cubin: false, + device_debug: DeviceDebug::Off, + debug_assertions: false, + }; + let mut environment = BTreeMap::from([ + (DEVICE_BACKEND_ENV.to_owned(), b"cutlass-mlir".to_vec()), + ( + cutlass_toolchain::CUTLASS_COMPILER_ENV.to_owned(), + b"/managed/libCutlassCompiler.so".to_vec(), + ), + ( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_owned(), + b"digest-a".to_vec(), + ), + ]); + let fingerprint = |environment: &BTreeMap>| { + passthrough_codegen_fingerprint_with_env( + &ctx, + &opts, + None, + Some("sm_120a"), + &MaterializationMode::default(), + environment, + ) + }; + environment.insert( + "CUDA_OXIDE_MLIR_PROFILE".into(), + b"cutlass-full-cute-mlir22".to_vec(), + ); + let full_profile = fingerprint(&environment); + environment.insert( + "CUDA_OXIDE_MLIR_PROFILE".into(), + b"cutlass-primitives-mlir22".to_vec(), + ); + let primitives_profile = fingerprint(&environment); + assert_ne!( + full_profile, primitives_profile, + "changing only the profile must rebuild device code" + ); + assert_eq!(primitives_profile, fingerprint(&environment)); + let before = fingerprint(&environment); + environment.insert( + CUTLASS_DIGEST_FINGERPRINT_ENV.to_owned(), + b"digest-b".to_vec(), + ); + + assert_ne!(before, fingerprint(&environment)); +} + +#[test] +fn passthrough_fingerprint_tracks_explicit_cutlass_library_content() { + let root = unique_temp_dir("cargo_oxide_explicit_cutlass_fingerprint"); + fs::create_dir_all(&root).unwrap(); + let compiler = root.join("libCutlassCompiler.so"); + fs::write(&compiler, b"compiler-a").unwrap(); + + let ctx = test_context(OxideConfig::default()); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: Some("sm_120a"), + features: None, + cargo_target_dir: None, + device_codegen_crate: None, + device_cfgs: &[], + no_fmad: false, + unchecked_indexing: false, + materialize_cubin: false, + device_debug: DeviceDebug::Off, + debug_assertions: false, + }; + let base_environment = BTreeMap::from([ + (DEVICE_BACKEND_ENV.to_owned(), b"cutlass-mlir".to_vec()), + ( + cutlass_toolchain::CUTLASS_COMPILER_ENV.to_owned(), + compiler.as_os_str().as_encoded_bytes().to_vec(), + ), + ]); + let fingerprint = |environment: &mut BTreeMap>| { + augment_env_with_cutlass_compiler_identity(&ctx, environment, Some(compiler.as_path())); + passthrough_codegen_fingerprint_with_env( + &ctx, + &opts, + None, + Some("sm_120a"), + &MaterializationMode::default(), + environment, + ) + }; + + let before = fingerprint(&mut base_environment.clone()); + fs::write(&compiler, b"compiler-b").unwrap(); + let after = fingerprint(&mut base_environment.clone()); + fs::remove_dir_all(&root).unwrap(); + + assert_ne!(before, after, "same-path compiler replacement must rebuild"); +} + #[test] fn passthrough_fingerprint_tracks_non_unicode_presence_switch_bytes() { let ctx = test_context(OxideConfig::default()); @@ -4726,6 +4893,100 @@ fn example_discovery_is_sorted_and_uses_manifest_fallback() { std::fs::remove_dir_all(root).unwrap(); } +#[test] +fn example_discovery_combines_roots_and_rejects_duplicate_names() { + let root = unique_temp_dir("cargo_oxide_list_multi_root"); + let standard = root.join("standard"); + let cute = root.join("cute"); + std::fs::create_dir_all(&standard).unwrap(); + std::fs::create_dir_all(&cute).unwrap(); + + write_list_example(&standard, "zeta", None, None); + write_list_example(&cute, "alpha", None, None); + let examples = discover_examples_from_roots(&[standard.clone(), cute.clone()]).unwrap(); + assert_eq!( + examples + .iter() + .map(|example| example.name.as_str()) + .collect::>(), + ["alpha", "zeta"] + ); + + write_list_example(&cute, "zeta", None, None); + let error = discover_examples_from_roots(&[standard, cute]).unwrap_err(); + assert!(error.contains("example name `zeta` exists in more than one example root")); + + std::fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn workspace_example_roots_follow_the_selected_backend_in_each_layout() { + for (nested, from_component) in [(false, false), (true, false), (true, true)] { + let root = unique_temp_dir("cargo_oxide_cute_example_root"); + let component = if nested { + root.join("cuda-oxide") + } else { + root.clone() + }; + let codegen = component.join("crates/rustc-codegen-cuda"); + let standard = codegen.join("examples"); + let cute = component.join("cute-rs-preview/examples"); + let models = component.join("cute-rs-preview/models"); + std::fs::create_dir_all(&standard).unwrap(); + std::fs::create_dir_all(&cute).unwrap(); + write_list_example(&standard, "native", None, None); + write_list_example(&cute, "tiled", None, None); + write_list_example(&models, "qwen_decoder_cute", None, None); + if nested { + // A leftover root tree must never shadow the selected component. + for path in ["cute-rs-preview/examples", "cute-rs-preview/models"] { + write_list_example(&root.join(path), "stale", None, None); + } + } + let mut ctx = test_context(OxideConfig::default()); + ctx.workspace_root = if from_component { + component + } else { + root.clone() + }; + ctx.codegen_crate = codegen; + ctx.examples_dir = standard.clone(); + ctx.is_workspace = true; + let examples = discover_examples_from_roots(&example_roots(&ctx)).unwrap(); + assert_eq!( + examples.iter().map(|e| e.name.as_str()).collect::>(), + ["native", "qwen_decoder_cute", "tiled"] + ); + assert_eq!(resolve_example_dir(&ctx, "tiled"), cute.join("tiled")); + assert_eq!( + resolve_example_dir(&ctx, "qwen_decoder_cute"), + models.join("qwen_decoder_cute") + ); + // Duplicate names still fail instead of choosing one source tree. + write_list_example(&models, "native", None, None); + assert!( + discover_examples_from_roots(&example_roots(&ctx)) + .unwrap_err() + .contains("more than one example root") + ); + std::fs::remove_dir_all(root).unwrap(); + } +} + +#[test] +fn standalone_example_roots_do_not_discover_neighboring_cute_trees() { + let root = unique_temp_dir("cargo_oxide_standalone_cute_root"); + for path in ["cute-rs-preview/examples", "cute-rs-preview/models"] { + std::fs::create_dir_all(root.join(path)).unwrap(); + } + let mut ctx = test_context(OxideConfig::default()); + ctx.workspace_root = root.clone(); + ctx.codegen_crate = root.clone(); + ctx.examples_dir = root.clone(); + assert_eq!(example_roots(&ctx).as_slice(), std::slice::from_ref(&root)); + std::fs::remove_dir_all(root).unwrap(); +} + #[test] fn example_discovery_keeps_examples_without_readmes() { let root = unique_temp_dir("cargo_oxide_list_missing_readme"); @@ -5387,16 +5648,15 @@ fn scaffold_async_template_keeps_async_deps_and_docs() { assert!(files.main_rs.contains("vecadd_async")); assert!(files.main_rs.contains("use cuda_host::cuda_module;")); assert!(!files.main_rs.contains("use cuda_device::{cuda_module")); - // Shared host runtime from the same git source as cuda-host; the SIMT - // surface lives under simt::. + // Shared host runtime from the same git source as cuda-host. assert!(files.cargo_toml.contains("cuda-async = { git = \"")); assert!(!files.cargo_toml.contains("cuda-async = \"")); assert!(!files.cargo_toml.contains("cuda-bindings")); - assert!(files.main_rs.contains("use cuda_core::simt::LaunchConfig;")); + assert!(files.main_rs.contains("use cuda_core::LaunchConfig;")); assert!( files .main_rs - .contains("use cuda_async::simt::device_operation::DeviceOperation;") + .contains("use cuda_async::device_operation::DeviceOperation;") ); } @@ -5480,6 +5740,8 @@ fn scaffold_names_the_canonical_repository() { .contains("git = \"https://github.com/NVIDIA/cuda-rust.git\"") ); assert!(!files.cargo_toml.contains("NVlabs")); + assert!(files.cargo_config_toml.contains("[env]")); + assert!(files.cargo_config_toml.contains("RUSTC_BOOTSTRAP = \"1\"")); } } diff --git a/cuda-oxide/crates/cargo-oxide/src/cutlass_toolchain.rs b/cuda-oxide/crates/cargo-oxide/src/cutlass_toolchain.rs new file mode 100644 index 0000000000..67ad0349f5 --- /dev/null +++ b/cuda-oxide/crates/cargo-oxide/src/cutlass_toolchain.rs @@ -0,0 +1,1033 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Managed installation of NVIDIA's public CUTLASS compiler bundle. +//! +//! The archive is a release artifact, not part of the CUDA Toolkit. Keep the +//! release URL and digest together here so an install can never silently drift +//! to a newer compiler. Publication is a same-filesystem directory rename: +//! builds either see the previous complete install or the new complete install, +//! never an extraction in progress. + +use sha2::Digest as _; +use std::collections::BTreeMap; +use std::ffi::OsString; +use std::fs::{self, File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Component, Path, PathBuf}; +use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +pub const CUTLASS_COMPILER_ENV: &str = "CUDA_OXIDE_CUTLASS_COMPILER"; + +const CUTLASS_VERSION: &str = "4.7.0"; +const CUTLASS_PLATFORM: &str = "x86_64-cu13"; +const ARCHIVE_NAME: &str = "cutlass-install-x86_64-cu13-4.7.0.tar.gz"; +const ARCHIVE_URL: &str = "https://github.com/NVIDIA/cutlass/releases/download/v4.7.0/cutlass-install-x86_64-cu13-4.7.0.tar.gz"; +const ARCHIVE_SHA256: &str = "32f0786c26ede0a5647fa25fc1b4a1429938f181ff28fd57b5632c0b00f0e24f"; +const COMPILER_LIBRARY_SHA256: &str = + "57df017e3585a10443c74c8b4cd99bda854242fb2f4c9534cf56d58c2c741628"; +const INSTALL_MANIFEST: &str = ".cuda-oxide-cutlass-toolchain"; +const INSTALL_MANIFEST_FORMAT: &str = "2"; +const LEGACY_INSTALL_MANIFEST_FORMAT: &str = "1"; + +static UNIQUE_COUNTER: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Copy, Debug)] +struct InstallSpec { + version: &'static str, + platform: &'static str, + archive_name: &'static str, + archive_url: &'static str, + archive_sha256: &'static str, + compiler_library_sha256: &'static str, +} + +const OFFICIAL: InstallSpec = InstallSpec { + version: CUTLASS_VERSION, + platform: CUTLASS_PLATFORM, + archive_name: ARCHIVE_NAME, + archive_url: ARCHIVE_URL, + archive_sha256: ARCHIVE_SHA256, + compiler_library_sha256: COMPILER_LIBRARY_SHA256, +}; + +/// Paths exported to compiler-backed builds after a verified installation. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ResolvedCutlass { + pub install_dir: PathBuf, + pub compiler_library: PathBuf, + pub compiler_library_sha256: String, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct InstallOutcome { + pub paths: ResolvedCutlass, + /// False when an already-complete pinned install was reused. + pub installed: bool, +} + +/// Install the one CUTLASS compiler release supported by this cargo-oxide. +pub fn install_official() -> Result { + ensure_supported_host()?; + let managed_root = managed_root()?; + let install_dir = install_dir(&managed_root, OFFICIAL); + + match resolve_at(&install_dir, OFFICIAL) { + Ok(Some(paths)) => { + return Ok(InstallOutcome { + paths, + installed: false, + }); + } + Ok(None) => {} + Err(error) => { + return Err(format!( + "the managed CUTLASS directory {} exists but is not a valid pinned installation: {error}\nmove or remove that directory, then rerun `cargo oxide toolchain install cutlass`", + install_dir.display() + )); + } + } + + let version_dir = install_dir.parent().ok_or_else(|| { + format!( + "could not determine the parent of managed install {}", + install_dir.display() + ) + })?; + fs::create_dir_all(version_dir).map_err(|error| { + format!( + "could not create managed CUTLASS directory {}: {error}", + version_dir.display() + ) + })?; + + let _lock = + InstallLock::acquire(version_dir.join(format!(".{}.install.lock", OFFICIAL.platform)))?; + + // A cooperating installer may have completed while this process waited + // for directory creation and lock acquisition. + match resolve_at(&install_dir, OFFICIAL) { + Ok(Some(paths)) => { + return Ok(InstallOutcome { + paths, + installed: false, + }); + } + Ok(None) => {} + Err(error) => { + return Err(format!( + "the managed CUTLASS directory {} became invalid: {error}", + install_dir.display() + )); + } + } + + let download_dir = TemporaryDirectory::create(version_dir, "download")?; + let archive = download_dir.path.join(OFFICIAL.archive_name); + eprintln!( + "NVIDIA's CUTLASS compiler bundle is proprietary software governed by the NVIDIA EULA:" + ); + eprintln!(" https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html"); + eprintln!( + "Downloading NVIDIA CUTLASS compiler {} ({})...", + OFFICIAL.version, OFFICIAL.platform + ); + download_archive(OFFICIAL.archive_url, &archive)?; + + eprintln!("Verifying SHA-256 {}...", OFFICIAL.archive_sha256); + let paths = install_verified_archive(OFFICIAL, &archive, &install_dir)?; + Ok(InstallOutcome { + paths, + installed: true, + }) +} + +/// Resolve the pinned managed install without downloading or mutating state. +/// +/// Unsupported hosts and machines without a resolvable Cargo home simply have +/// no managed installation. An existing but malformed install is an error so +/// a requested compiler backend can diagnose it instead of using partial files. +pub fn resolve_official() -> Result, String> { + if !supported_host() { + return Ok(None); + } + let Some(managed_root) = managed_root_optional() else { + return Ok(None); + }; + resolve_at(&install_dir(&managed_root, OFFICIAL), OFFICIAL) +} + +pub fn print_install_outcome(outcome: &InstallOutcome) { + if outcome.installed { + println!( + "Installed NVIDIA CUTLASS compiler {} ({})", + CUTLASS_VERSION, CUTLASS_PLATFORM + ); + } else { + println!( + "NVIDIA CUTLASS compiler {} ({}) is already installed", + CUTLASS_VERSION, CUTLASS_PLATFORM + ); + } + println!(" root: {}", outcome.paths.install_dir.display()); + println!(" compiler: {}", outcome.paths.compiler_library.display()); + println!(); + println!( + "Build and run commands will provide {CUTLASS_COMPILER_ENV} automatically when it is not explicitly configured." + ); +} + +fn supported_host() -> bool { + std::env::consts::OS == "linux" && std::env::consts::ARCH == "x86_64" +} + +fn ensure_supported_host() -> Result<(), String> { + if supported_host() { + Ok(()) + } else { + Err(format!( + "CUTLASS compiler {} is pinned to Linux x86_64 ({CUTLASS_PLATFORM}); this host is {} {}", + CUTLASS_VERSION, + std::env::consts::OS, + std::env::consts::ARCH + )) + } +} + +fn managed_root() -> Result { + managed_root_optional().ok_or_else(|| { + "could not determine the per-user cargo-oxide directory; set CARGO_HOME or HOME".to_string() + }) +} + +fn managed_root_optional() -> Option { + cargo_home_from( + std::env::var_os("CARGO_HOME"), + std::env::var_os("HOME"), + std::env::current_dir().ok().as_deref(), + ) + .map(|home| home.join("cuda-oxide/toolchains/cutlass")) +} + +fn cargo_home_from( + cargo_home: Option, + home: Option, + current_dir: Option<&Path>, +) -> Option { + let configured = cargo_home + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + .or_else(|| { + home.filter(|value| !value.is_empty()) + .map(PathBuf::from) + .map(|path| path.join(".cargo")) + })?; + if configured.is_absolute() { + Some(configured) + } else { + current_dir.map(|cwd| cwd.join(configured)) + } +} + +fn install_dir(managed_root: &Path, spec: InstallSpec) -> PathBuf { + managed_root.join(spec.version).join(spec.platform) +} + +fn download_archive(url: &str, destination: &Path) -> Result<(), String> { + let output = Command::new("curl") + .args([ + "--fail", + "--location", + "--proto", + "=https", + "--tlsv1.2", + "--retry", + "2", + "--retry-connrefused", + "--silent", + "--show-error", + "--user-agent", + concat!("cargo-oxide/", env!("CARGO_PKG_VERSION")), + "--output", + ]) + .arg(destination) + .arg(url) + .output() + .map_err(|error| { + format!( + "could not start `curl` to download {url}: {error}\ninstall curl and rerun `cargo oxide toolchain install cutlass`" + ) + })?; + if !output.status.success() { + return Err(format!( + "curl failed while downloading {url} (status {}):\n{}", + display_status(output.status.code()), + bounded_diagnostics(&output.stderr) + )); + } + let metadata = fs::metadata(destination).map_err(|error| { + format!( + "curl reported success but {} could not be inspected: {error}", + destination.display() + ) + })?; + if !metadata.is_file() || metadata.len() == 0 { + return Err(format!( + "curl reported success but downloaded an empty/non-file archive at {}", + destination.display() + )); + } + Ok(()) +} + +fn install_verified_archive( + spec: InstallSpec, + archive: &Path, + final_dir: &Path, +) -> Result { + let actual = sha256_file(archive)?; + if actual != spec.archive_sha256 { + return Err(format!( + "CUTLASS archive digest mismatch for {}\n expected: {}\n actual: {}\nrefusing to extract or publish the archive", + archive.display(), + spec.archive_sha256, + actual + )); + } + + if final_dir.exists() { + return Err(format!( + "refusing to overwrite existing managed CUTLASS directory {}", + final_dir.display() + )); + } + let parent = final_dir.parent().ok_or_else(|| { + format!( + "could not determine parent directory for {}", + final_dir.display() + ) + })?; + let staging = TemporaryDirectory::create(parent, "extract")?; + let extract_dir = staging.path.join("payload"); + fs::create_dir(&extract_dir).map_err(|error| { + format!( + "could not create extraction directory {}: {error}", + extract_dir.display() + ) + })?; + + validate_archive_paths(archive)?; + extract_archive(archive, &extract_dir)?; + let payload_root = find_payload_root(&extract_dir)?; + let discovered = discover_payload(&payload_root, spec)?; + write_install_manifest(&payload_root, spec, &discovered)?; + + fs::rename(&payload_root, final_dir).map_err(|error| { + format!( + "could not atomically publish CUTLASS compiler at {}: {error}", + final_dir.display() + ) + })?; + + match resolve_at(final_dir, spec) { + Ok(Some(paths)) => Ok(paths), + Ok(None) => Err(format!( + "published CUTLASS directory {} disappeared", + final_dir.display() + )), + Err(error) => { + let quarantine = staging.path.join("failed-publication"); + let _ = fs::rename(final_dir, &quarantine); + Err(format!( + "published CUTLASS directory failed validation and was withdrawn: {error}" + )) + } + } +} + +fn validate_archive_paths(archive: &Path) -> Result<(), String> { + let output = Command::new("tar") + .args(["--list", "--gzip", "--file"]) + .arg(archive) + .output() + .map_err(|error| { + format!( + "could not start `tar` to inspect {}: {error}\ninstall GNU tar and rerun the installer", + archive.display() + ) + })?; + if !output.status.success() { + return Err(format!( + "tar could not list verified archive {} (status {}):\n{}", + archive.display(), + display_status(output.status.code()), + bounded_diagnostics(&output.stderr) + )); + } + let listing = String::from_utf8(output.stdout).map_err(|_| { + format!( + "archive {} contains a non-UTF-8 path; refusing extraction", + archive.display() + ) + })?; + for entry in listing.lines().filter(|line| !line.is_empty()) { + let path = Path::new(entry); + if path.is_absolute() + || path.components().any(|component| { + matches!( + component, + Component::ParentDir | Component::RootDir | Component::Prefix(_) + ) + }) + { + return Err(format!( + "archive {} contains unsafe path {entry:?}; refusing extraction", + archive.display() + )); + } + } + Ok(()) +} + +fn extract_archive(archive: &Path, destination: &Path) -> Result<(), String> { + let output = Command::new("tar") + .args([ + "--extract", + "--gzip", + "--file", + ]) + .arg(archive) + .args([ + "--directory", + ]) + .arg(destination) + .args([ + "--no-same-owner", + "--no-same-permissions", + "--delay-directory-restore", + ]) + .output() + .map_err(|error| { + format!( + "could not start `tar` to extract {}: {error}\ninstall GNU tar and rerun the installer", + archive.display() + ) + })?; + if !output.status.success() { + return Err(format!( + "tar failed to extract verified archive {} (status {}):\n{}", + archive.display(), + display_status(output.status.code()), + bounded_diagnostics(&output.stderr) + )); + } + Ok(()) +} + +fn find_payload_root(extract_dir: &Path) -> Result { + let mut candidates = Vec::new(); + collect_payload_roots(extract_dir, 0, &mut candidates)?; + match candidates.as_slice() { + [root] => Ok(root.clone()), + [] => Err(format!( + "verified CUTLASS archive did not contain an exact libCutlassCompiler.so under {}", + extract_dir.display() + )), + _ => Err(format!( + "verified CUTLASS archive contained multiple compiler roots: {}", + candidates + .iter() + .map(|path| path.display().to_string()) + .collect::>() + .join(", ") + )), + } +} + +fn collect_payload_roots( + directory: &Path, + depth: usize, + candidates: &mut Vec, +) -> Result<(), String> { + if exact_compiler_library(directory).is_some() { + candidates.push(directory.to_path_buf()); + return Ok(()); + } + if depth >= 4 { + return Ok(()); + } + let mut children = fs::read_dir(directory) + .map_err(|error| format!("could not inspect {}: {error}", directory.display()))? + .collect::, _>>() + .map_err(|error| format!("could not inspect {}: {error}", directory.display()))?; + children.sort_by_key(|entry| entry.file_name()); + for child in children { + let file_type = child + .file_type() + .map_err(|error| format!("could not inspect {}: {error}", child.path().display()))?; + if file_type.is_dir() && !file_type.is_symlink() { + collect_payload_roots(&child.path(), depth + 1, candidates)?; + } + } + Ok(()) +} + +#[derive(Clone, Debug)] +struct PayloadPaths { + compiler_library: PathBuf, +} + +fn discover_payload(root: &Path, spec: InstallSpec) -> Result { + let compiler_library = exact_compiler_library(root).ok_or_else(|| { + format!( + "verified CUTLASS archive has no exact libCutlassCompiler.so under {}/lib or {}/lib64", + root.display(), + root.display() + ) + })?; + validate_regular_nonempty(root, &compiler_library, "libCutlassCompiler shared library")?; + let actual = sha256_file(&compiler_library)?; + if actual != spec.compiler_library_sha256 { + return Err(format!( + "CUTLASS compiler library digest mismatch for {}\n expected: {}\n actual: {}\nrefusing to publish the installation", + compiler_library.display(), + spec.compiler_library_sha256, + actual + )); + } + + Ok(PayloadPaths { compiler_library }) +} + +fn exact_compiler_library(root: &Path) -> Option { + for lib_dir in [root.join("lib"), root.join("lib64")] { + let exact = lib_dir.join("libCutlassCompiler.so"); + if exact.is_file() { + return Some(exact); + } + } + None +} + +fn validate_regular_nonempty( + root: &Path, + path: &Path, + label: &str, +) -> Result { + let unresolved_metadata = fs::symlink_metadata(path) + .map_err(|error| format!("could not inspect {label} {}: {error}", path.display()))?; + if !unresolved_metadata.file_type().is_file() { + return Err(format!( + "{label} {} is not a nonempty regular file", + path.display() + )); + } + let canonical_root = fs::canonicalize(root) + .map_err(|error| format!("could not resolve {}: {error}", root.display()))?; + let canonical_path = fs::canonicalize(path) + .map_err(|error| format!("could not resolve {label} {}: {error}", path.display()))?; + if !canonical_path.starts_with(&canonical_root) { + return Err(format!( + "{label} {} resolves outside the managed installation", + path.display() + )); + } + let metadata = fs::metadata(&canonical_path) + .map_err(|error| format!("could not inspect {label} {}: {error}", path.display()))?; + if !metadata.is_file() || metadata.len() == 0 { + return Err(format!( + "{label} {} is not a nonempty regular file", + path.display() + )); + } + Ok(metadata) +} + +fn write_install_manifest( + root: &Path, + spec: InstallSpec, + paths: &PayloadPaths, +) -> Result<(), String> { + let relative = |path: &Path| { + path.strip_prefix(root) + .map(path_to_manifest_value) + .map_err(|_| format!("{} is not under {}", path.display(), root.display())) + }; + let manifest = format!( + "format={INSTALL_MANIFEST_FORMAT}\nversion={}\nplatform={}\narchive-url={}\narchive-sha256={}\ncompiler-library={}\ncompiler-library-sha256={}\n", + spec.version, + spec.platform, + spec.archive_url, + spec.archive_sha256, + relative(&paths.compiler_library)?, + spec.compiler_library_sha256, + ); + let manifest_path = root.join(INSTALL_MANIFEST); + let mut file = File::create(&manifest_path).map_err(|error| { + format!( + "could not create install manifest {}: {error}", + manifest_path.display() + ) + })?; + file.write_all(manifest.as_bytes()).map_err(|error| { + format!( + "could not write install manifest {}: {error}", + manifest_path.display() + ) + })?; + file.sync_all().map_err(|error| { + format!( + "could not sync install manifest {}: {error}", + manifest_path.display() + ) + }) +} + +fn path_to_manifest_value(path: &Path) -> String { + path.components() + .filter_map(|component| match component { + Component::Normal(value) => Some(value.to_string_lossy()), + _ => None, + }) + .collect::>() + .join("/") +} + +fn resolve_at(root: &Path, spec: InstallSpec) -> Result, String> { + if !root.exists() { + return Ok(None); + } + if !root.is_dir() { + return Err(format!("{} is not a directory", root.display())); + } + let manifest_path = root.join(INSTALL_MANIFEST); + let contents = fs::read_to_string(&manifest_path).map_err(|error| { + format!( + "could not read completion manifest {}: {error}", + manifest_path.display() + ) + })?; + let fields = parse_manifest(&contents)?; + let format = manifest_field(&fields, "format")?; + if format != INSTALL_MANIFEST_FORMAT && format != LEGACY_INSTALL_MANIFEST_FORMAT { + return Err(format!( + "managed CUTLASS manifest has format={format:?}, expected {INSTALL_MANIFEST_FORMAT:?}" + )); + } + require_manifest_value(&fields, "version", spec.version)?; + require_manifest_value(&fields, "platform", spec.platform)?; + require_manifest_value(&fields, "archive-url", spec.archive_url)?; + require_manifest_value(&fields, "archive-sha256", spec.archive_sha256)?; + if format == INSTALL_MANIFEST_FORMAT { + require_manifest_value( + &fields, + "compiler-library-sha256", + spec.compiler_library_sha256, + )?; + } + + let compiler_library = manifest_path_join(root, manifest_field(&fields, "compiler-library")?)?; + + validate_regular_nonempty(root, &compiler_library, "libCutlassCompiler shared library")?; + let compiler_library_sha256 = sha256_file(&compiler_library)?; + if compiler_library_sha256 != spec.compiler_library_sha256 { + return Err(format!( + "managed CUTLASS compiler library digest mismatch for {}\n expected: {}\n actual: {}", + compiler_library.display(), + spec.compiler_library_sha256, + compiler_library_sha256 + )); + } + + Ok(Some(ResolvedCutlass { + install_dir: root.to_path_buf(), + compiler_library, + compiler_library_sha256, + })) +} + +fn parse_manifest(contents: &str) -> Result, String> { + let mut fields = BTreeMap::new(); + for (line_number, line) in contents.lines().enumerate() { + let (key, value) = line.split_once('=').ok_or_else(|| { + format!( + "invalid managed CUTLASS manifest line {}: expected key=value", + line_number + 1 + ) + })?; + if key.is_empty() || value.is_empty() { + return Err(format!( + "invalid managed CUTLASS manifest line {}: empty key or value", + line_number + 1 + )); + } + if fields.insert(key.to_string(), value.to_string()).is_some() { + return Err(format!( + "invalid managed CUTLASS manifest: duplicate field {key:?}" + )); + } + } + Ok(fields) +} + +fn manifest_field<'a>(fields: &'a BTreeMap, key: &str) -> Result<&'a str, String> { + fields + .get(key) + .map(String::as_str) + .ok_or_else(|| format!("managed CUTLASS manifest is missing {key:?}")) +} + +fn require_manifest_value( + fields: &BTreeMap, + key: &str, + expected: &str, +) -> Result<(), String> { + let actual = manifest_field(fields, key)?; + if actual == expected { + Ok(()) + } else { + Err(format!( + "managed CUTLASS manifest has {key}={actual:?}, expected {expected:?}" + )) + } +} + +fn manifest_path_join(root: &Path, value: &str) -> Result { + let relative = Path::new(value); + if relative.is_absolute() + || relative.as_os_str().is_empty() + || relative + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + return Err(format!( + "managed CUTLASS manifest contains unsafe relative path {value:?}" + )); + } + Ok(root.join(relative)) +} + +pub(crate) fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path) + .map_err(|error| format!("could not open {} for hashing: {error}", path.display()))?; + let mut hash = sha2::Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let count = file + .read(&mut buffer) + .map_err(|error| format!("could not hash {}: {error}", path.display()))?; + if count == 0 { + break; + } + hash.update(&buffer[..count]); + } + let digest: [u8; 32] = hash.finalize().into(); + let mut encoded = String::with_capacity(64); + use std::fmt::Write as _; + for byte in digest { + write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail"); + } + Ok(encoded) +} + +fn display_status(code: Option) -> String { + code.map(|code| code.to_string()) + .unwrap_or_else(|| "terminated by signal".to_string()) +} + +fn bounded_diagnostics(bytes: &[u8]) -> String { + const LIMIT: usize = 8 * 1024; + let bytes = if bytes.len() > LIMIT { + &bytes[bytes.len() - LIMIT..] + } else { + bytes + }; + let diagnostics = String::from_utf8_lossy(bytes); + let diagnostics = diagnostics.trim(); + if diagnostics.is_empty() { + "".to_string() + } else { + diagnostics.to_string() + } +} + +struct InstallLock { + path: PathBuf, +} + +impl InstallLock { + fn acquire(path: PathBuf) -> Result { + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + .map_err(|error| { + format!( + "could not acquire CUTLASS installer lock {}: {error}\nif no installer is running, remove this stale lock and retry", + path.display() + ) + })?; + writeln!(file, "pid={}", std::process::id()).map_err(|error| { + format!("could not write installer lock {}: {error}", path.display()) + })?; + Ok(Self { path }) + } +} + +impl Drop for InstallLock { + fn drop(&mut self) { + let _ = fs::remove_file(&self.path); + } +} + +struct TemporaryDirectory { + path: PathBuf, +} + +impl TemporaryDirectory { + fn create(parent: &Path, purpose: &str) -> Result { + for _ in 0..32 { + let unique = unique_suffix(); + let path = parent.join(format!( + ".{CUTLASS_PLATFORM}.{purpose}.{}.{unique}", + std::process::id() + )); + match fs::create_dir(&path) { + Ok(()) => return Ok(Self { path }), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => { + return Err(format!( + "could not create temporary installer directory {}: {error}", + path.display() + )); + } + } + } + Err(format!( + "could not allocate a unique installer directory under {}", + parent.display() + )) + } +} + +impl Drop for TemporaryDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + } +} + +fn unique_suffix() -> u128 { + let time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + time ^ u128::from(UNIQUE_COUNTER.fetch_add(1, Ordering::Relaxed)) +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TestDirectory(PathBuf); + + impl TestDirectory { + fn new(name: &str) -> Self { + let path = std::env::temp_dir().join(format!( + "cargo-oxide-cutlass-{name}-{}-{}", + std::process::id(), + unique_suffix() + )); + fs::create_dir(&path).unwrap(); + Self(path) + } + } + + impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn fixture_archive(root: &Path) -> PathBuf { + let source = root.join("source/cutlass-fixture"); + fs::create_dir_all(source.join("lib")).unwrap(); + fs::write( + source.join("lib/libCutlassCompiler.so"), + b"fixture library\n", + ) + .unwrap(); + let archive = root.join("fixture.tar.gz"); + let status = Command::new("tar") + .args(["--create", "--gzip", "--file"]) + .arg(&archive) + .args(["--directory"]) + .arg(root.join("source")) + .arg("cutlass-fixture") + .status() + .unwrap(); + assert!(status.success()); + archive + } + + fn fixture_spec(archive_digest: &'static str, library_digest: &'static str) -> InstallSpec { + InstallSpec { + version: "test-version", + platform: CUTLASS_PLATFORM, + archive_name: "fixture.tar.gz", + archive_url: "https://example.invalid/fixture.tar.gz", + archive_sha256: archive_digest, + compiler_library_sha256: library_digest, + } + } + + #[test] + fn cargo_home_prefers_explicit_value_and_absolutizes_relative_paths() { + let cwd = Path::new("/workspace"); + assert_eq!( + cargo_home_from( + Some(OsString::from("relative-cargo")), + Some(OsString::from("/home/user")), + Some(cwd), + ), + Some(PathBuf::from("/workspace/relative-cargo")) + ); + assert_eq!( + cargo_home_from(None, Some(OsString::from("/home/user")), Some(cwd)), + Some(PathBuf::from("/home/user/.cargo")) + ); + } + + #[test] + fn verified_fixture_is_published_and_resolved_without_network() { + let temp = TestDirectory::new("install"); + let archive = fixture_archive(&temp.0); + let digest = sha256_file(&archive).unwrap(); + let digest: &'static str = Box::leak(digest.into_boxed_str()); + let library_digest = sha256_file( + &temp + .0 + .join("source/cutlass-fixture/lib/libCutlassCompiler.so"), + ) + .unwrap(); + let library_digest: &'static str = Box::leak(library_digest.into_boxed_str()); + let spec = fixture_spec(digest, library_digest); + let final_dir = install_dir(&temp.0.join("managed"), spec); + fs::create_dir_all(final_dir.parent().unwrap()).unwrap(); + + let installed = install_verified_archive(spec, &archive, &final_dir).unwrap(); + let resolved = resolve_at(&final_dir, spec).unwrap().unwrap(); + assert_eq!(installed, resolved); + assert_eq!( + fs::read(&resolved.compiler_library).unwrap(), + b"fixture library\n" + ); + assert!( + resolved + .compiler_library + .ends_with("lib/libCutlassCompiler.so") + ); + assert!(final_dir.join(INSTALL_MANIFEST).is_file()); + } + + #[test] + fn resolver_rejects_library_replacement_at_the_same_managed_path() { + let temp = TestDirectory::new("replaced-library"); + let archive = fixture_archive(&temp.0); + let archive_digest = sha256_file(&archive).unwrap(); + let archive_digest: &'static str = Box::leak(archive_digest.into_boxed_str()); + let library_digest = sha256_file( + &temp + .0 + .join("source/cutlass-fixture/lib/libCutlassCompiler.so"), + ) + .unwrap(); + let library_digest: &'static str = Box::leak(library_digest.into_boxed_str()); + let spec = fixture_spec(archive_digest, library_digest); + let final_dir = install_dir(&temp.0.join("managed"), spec); + fs::create_dir_all(final_dir.parent().unwrap()).unwrap(); + let installed = install_verified_archive(spec, &archive, &final_dir).unwrap(); + + fs::write(&installed.compiler_library, b"replaced library\n").unwrap(); + let error = resolve_at(&final_dir, spec).unwrap_err(); + assert!(error.contains("library digest mismatch"), "{error}"); + assert!(error.contains(library_digest), "{error}"); + } + + #[cfg(unix)] + #[test] + fn resolver_rejects_a_symlinked_compiler_library() { + use std::os::unix::fs::symlink; + + let temp = TestDirectory::new("symlinked-library"); + let archive = fixture_archive(&temp.0); + let archive_digest = sha256_file(&archive).unwrap(); + let archive_digest: &'static str = Box::leak(archive_digest.into_boxed_str()); + let library_digest = sha256_file( + &temp + .0 + .join("source/cutlass-fixture/lib/libCutlassCompiler.so"), + ) + .unwrap(); + let library_digest: &'static str = Box::leak(library_digest.into_boxed_str()); + let spec = fixture_spec(archive_digest, library_digest); + let final_dir = install_dir(&temp.0.join("managed"), spec); + fs::create_dir_all(final_dir.parent().unwrap()).unwrap(); + let installed = install_verified_archive(spec, &archive, &final_dir).unwrap(); + let real_library = installed.compiler_library.with_extension("so.real"); + + fs::rename(&installed.compiler_library, &real_library).unwrap(); + symlink(&real_library, &installed.compiler_library).unwrap(); + + let error = resolve_at(&final_dir, spec).unwrap_err(); + assert!(error.contains("is not a nonempty regular file"), "{error}"); + } + + #[test] + fn digest_mismatch_never_extracts_or_publishes() { + let temp = TestDirectory::new("bad-digest"); + let archive = fixture_archive(&temp.0); + let spec = fixture_spec( + "0000000000000000000000000000000000000000000000000000000000000000", + "0000000000000000000000000000000000000000000000000000000000000000", + ); + let final_dir = install_dir(&temp.0.join("managed"), spec); + fs::create_dir_all(final_dir.parent().unwrap()).unwrap(); + + let error = install_verified_archive(spec, &archive, &final_dir).unwrap_err(); + assert!(error.contains("digest mismatch"), "{error}"); + assert!(error.contains("refusing to extract or publish"), "{error}"); + assert!(!final_dir.exists()); + } + + #[test] + fn resolver_rejects_manifest_path_escape() { + let temp = TestDirectory::new("escape"); + let root = temp.0.join("install"); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join(INSTALL_MANIFEST), + format!( + "format=2\nversion={}\nplatform={}\narchive-url={}\narchive-sha256={}\ncompiler-library=../outside\ncompiler-library-sha256={}\n", + OFFICIAL.version, + OFFICIAL.platform, + OFFICIAL.archive_url, + OFFICIAL.archive_sha256, + OFFICIAL.compiler_library_sha256, + ), + ) + .unwrap(); + + let error = resolve_at(&root, OFFICIAL).unwrap_err(); + assert!(error.contains("unsafe relative path"), "{error}"); + } + + #[test] + fn archive_path_check_rejects_parent_components() { + let error = manifest_path_join(Path::new("/managed"), "../escape").unwrap_err(); + assert!(error.contains("unsafe relative path"), "{error}"); + } +} diff --git a/cuda-oxide/crates/cargo-oxide/src/main.rs b/cuda-oxide/crates/cargo-oxide/src/main.rs index 9bdef87cfa..bfd4524b61 100644 --- a/cuda-oxide/crates/cargo-oxide/src/main.rs +++ b/cuda-oxide/crates/cargo-oxide/src/main.rs @@ -29,6 +29,7 @@ //! cargo oxide clean # remove local build outputs //! cargo oxide setup # explicitly build/install backend //! cargo oxide update # refresh cached backend (external) +//! cargo oxide toolchain install cutlass # install the pinned CUTLASS compiler //! ``` use clap::{CommandFactory, Parser, Subcommand, ValueEnum, error::ErrorKind}; @@ -38,6 +39,7 @@ mod artifact_identity; mod backend; mod backend_source; mod commands; +mod cutlass_toolchain; /// Top-level CLI structure parsed by clap. /// @@ -73,9 +75,9 @@ enum Commands { /// restating its alias, case and whitespace rules. #[command(name = "__debug-policy", hide = true)] DebugPolicy, - /// Build and run an example or project + /// Build and run an example, model, or project Run { - /// Example name (required in workspace, optional for standalone projects) + /// Example or model name (required in workspace, optional for standalone projects) example: Option, /// Generate NVVM IR (use with libNVVM -gen-lto) #[arg(long)] @@ -455,7 +457,7 @@ enum Commands { #[arg(long)] tui: bool, }, - /// List the examples bundled with the cuda-oxide workspace + /// List the examples and models bundled with the cuda-oxide workspace List { /// Emit stable machine-readable JSON #[arg(long)] @@ -520,6 +522,26 @@ enum Commands { #[arg(long)] force: bool, }, + /// Install compiler toolchains managed by cargo-oxide + Toolchain { + #[command(subcommand)] + command: ToolchainCommand, + }, +} + +#[derive(Subcommand)] +enum ToolchainCommand { + /// Install a pinned official toolchain release + Install { + #[arg(value_enum)] + tool: ManagedToolchain, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +enum ManagedToolchain { + /// NVIDIA CUTLASS compiler 4.7.0 for Linux x86_64 and CUDA 13 + Cutlass, } #[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] @@ -675,6 +697,10 @@ fn validate_materialization_cli(cli: &Cli) -> Result<(), String> { "--materialize-cubin cannot be used with update because update only refreshes the codegen backend" .to_string(), ), + Commands::Toolchain { .. } => Err( + "--materialize-cubin cannot be used with toolchain because toolchain manages compiler installations" + .to_string(), + ), Commands::MaterializerHandshake => Err( "--materialize-cubin cannot be passed to the internal materializer discovery helper" .to_string(), @@ -1190,6 +1216,17 @@ fn main() { let ctx = commands::resolve_passive_context(); commands::update(&ctx, force); } + Commands::Toolchain { command } => match command { + ToolchainCommand::Install { + tool: ManagedToolchain::Cutlass, + } => match cutlass_toolchain::install_official() { + Ok(outcome) => cutlass_toolchain::print_install_outcome(&outcome), + Err(error) => { + eprintln!("Error: could not install the CUTLASS compiler: {error}"); + std::process::exit(1); + } + }, + }, } } @@ -1215,7 +1252,9 @@ fn resolve_example_name(name: Option, ctx: &commands::Context, subcomman eprintln!(); eprintln!("Usage: cargo oxide {subcommand} "); eprintln!(); - eprintln!("Available examples are in crates/rustc-codegen-cuda/examples/"); + eprintln!( + "Run `cargo oxide list`, or browse crates/rustc-codegen-cuda/examples/, cute-rs-preview/examples/, and cute-rs-preview/models/" + ); std::process::exit(1); } @@ -1306,6 +1345,20 @@ mod tests { assert!(force); } + #[test] + fn toolchain_parser_accepts_pinned_cutlass_install() { + let cli = Cli::try_parse_from(["cargo-oxide", "toolchain", "install", "cutlass"]) + .expect("managed CUTLASS install command should parse"); + assert!(matches!( + cli.command, + Commands::Toolchain { + command: ToolchainCommand::Install { + tool: ManagedToolchain::Cutlass + } + } + )); + } + #[test] fn build_parser_preserves_nested_cargo_and_test_separators() { let args = strings(&[ diff --git a/cuda-oxide/crates/cuda-device/Cargo.toml b/cuda-oxide/crates/cuda-device/Cargo.toml index 00c2a1b1f0..fcbb2432b8 100644 --- a/cuda-oxide/crates/cuda-device/Cargo.toml +++ b/cuda-oxide/crates/cuda-device/Cargo.toml @@ -15,6 +15,11 @@ readme = "README.md" # it back on for all of them and defeat the feature. `cuda-host` requests # `host` explicitly, so any build that loads or launches a module still has it. cuda-macros = { workspace = true } +# Shared no_std marker for owned-buffer representation changes. Keeping it +# separate from cuda-core leaves kernel-only builds free of host dependencies. +cuda-repr = { version = "=0.5.0", path = "../../../cuda-repr" } +# Vector types implement Pod for safe cast_chunks targets. +bytemuck = { workspace = true } [dev-dependencies] trybuild = "1" diff --git a/cuda-oxide/crates/cuda-device/README.md b/cuda-oxide/crates/cuda-device/README.md index 3309d77ee3..3dad327a02 100644 --- a/cuda-oxide/crates/cuda-device/README.md +++ b/cuda-oxide/crates/cuda-device/README.md @@ -2,6 +2,10 @@ `#![no_std]` device-side intrinsics and abstractions for writing CUDA kernels in Rust. This crate provides everything that runs on the GPU: thread identification, memory abstractions, synchronization, warp primitives, tensor cores, TMA, atomics, and debug facilities. +The separate [`ctm-rs`](../../cute-rs-preview/ctm-rs) crate provides typed CUTLASS-style +PTX-level programming over the existing NVVM dialect. Its initial validated +subset targets two-CTA FP16/FP32 operations on SM100a and block-scaled MXFP4 operations on SM120a. + ```text User kernel code │ diff --git a/cuda-oxide/crates/cuda-device/src/atomic.rs b/cuda-oxide/crates/cuda-device/src/atomic.rs index 7433fec228..8fcc7c01e8 100644 --- a/cuda-oxide/crates/cuda-device/src/atomic.rs +++ b/cuda-oxide/crates/cuda-device/src/atomic.rs @@ -176,6 +176,16 @@ macro_rules! define_integer_atomic { // All operations go through atomic instructions emitted by the compiler. unsafe impl Sync for $Name {} + // SAFETY: `$Name` is `#[repr(transparent)]` over `UnsafeCell<$ty>`, + // which is itself `#[repr(transparent)]` over `$ty`, with no other + // fields, so it has the layout and the valid values of `$ty`. This + // lets a `DeviceBuffer<$ty>` be viewed as `DeviceBuffer<$Name>` with + // cuda-core's `DeviceBuffer::wrap_elem` and turned back with + // `peel_elem` so the host can read the result. Neither type needs + // destruction or owns resources. DeviceRepr exposes no reference + // conversions: &integer and &atomic have different aliasing rules. + unsafe impl cuda_repr::DeviceRepr<$ty> for $Name {} + impl $Name { /// Create a new atomic with the given initial value. pub const fn new(val: $ty) -> Self { @@ -372,6 +382,16 @@ macro_rules! define_float_atomic { // All operations go through atomic instructions emitted by the compiler. unsafe impl Sync for $Name {} + // SAFETY: `$Name` is `#[repr(transparent)]` over `UnsafeCell<$ty>`, + // which is itself `#[repr(transparent)]` over `$ty`, with no other + // fields, so it has the layout and the valid values of `$ty`. This + // lets a `DeviceBuffer<$ty>` be viewed as `DeviceBuffer<$Name>` with + // cuda-core's `DeviceBuffer::wrap_elem` and turned back with + // `peel_elem` so the host can read the result. Neither type needs + // destruction or owns resources. DeviceRepr exposes no reference + // conversions: &float and &atomic have different aliasing rules. + unsafe impl cuda_repr::DeviceRepr<$ty> for $Name {} + impl $Name { /// Create a new atomic with the given initial value. pub const fn new(val: $ty) -> Self { @@ -494,7 +514,7 @@ define_float_atomic! { /// /// `swap` lowers to hardware `atom.exch.b32`; this API does not expose /// `compare_exchange` for floats. `fetch_add` uses LLVM `atomicrmw fadd`. - /// With the pinned `nightly-2026-10-11` toolchain's LLVM 23 NVPTX defaults, + /// With the pinned `1.99.1` toolchain's LLVM 23 NVPTX defaults, /// global and generic `f32` adds use a compare-and-swap loop; shared `f32` /// and `f64` adds still use native instructions on supported targets. /// diff --git a/cuda-oxide/crates/cuda-device/src/barrier.rs b/cuda-oxide/crates/cuda-device/src/barrier.rs index b21907e0c4..82a8934c0c 100644 --- a/cuda-oxide/crates/cuda-device/src/barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/barrier.rs @@ -293,6 +293,21 @@ unsafe impl Send for ManagedBarrier {} // ============================================================================= impl ManagedBarrier { + /// Issue initialization without thread election, fences or synchronization. + /// + /// # Safety + /// Exactly one thread must initialize live shared storage. `count` must be + /// a valid hardware arrival count. Publish initialization with the required + /// fence and synchronization before any participant uses the barrier. + #[inline(always)] + pub unsafe fn init_instruction(self, count: u32) -> ManagedBarrier { + unsafe { mbarrier_init(self.ptr.cast_mut(), count) }; + ManagedBarrier { + ptr: self.ptr, + _state: PhantomData, + _kind: PhantomData, + } + } /// Create an Uninit barrier from an explicit static declaration. /// /// Wrap a `static mut Barrier` in the typestate wrapper. Each barrier @@ -422,6 +437,33 @@ impl ManagedBarrier { // ============================================================================= impl ManagedBarrier { + /// Adopt a participant's initialized barrier after caller synchronization. + /// + /// # Safety + /// `ptr` must identify an initialized barrier in live shared memory. All + /// participants must coordinate accesses and invalidation for its lifetime. + #[inline(always)] + pub unsafe fn adopt_ready(ptr: *mut Barrier) -> Self { + Self { + ptr, + _state: PhantomData, + _kind: PhantomData, + } + } + /// Issue invalidation without election or synchronization. + /// + /// # Safety + /// Exactly one thread must invalidate after all users and pending + /// asynchronous operations are finished. + #[inline(always)] + pub unsafe fn invalidate_instruction(self) -> ManagedBarrier { + unsafe { mbarrier_inval(self.ptr.cast_mut()) }; + ManagedBarrier { + ptr: self.ptr, + _state: PhantomData, + _kind: PhantomData, + } + } /// Get the raw pointer to the underlying barrier. /// /// Useful for interop with low-level APIs. diff --git a/cuda-oxide/crates/cuda-device/src/cluster.rs b/cuda-oxide/crates/cuda-device/src/cluster.rs index 07d6c93076..300156c827 100644 --- a/cuda-oxide/crates/cuda-device/src/cluster.rs +++ b/cuda-oxide/crates/cuda-device/src/cluster.rs @@ -168,6 +168,39 @@ pub fn cluster_idx() -> u32 { x + y * nx + z * nx * ny } +/// Get cluster's X index within the grid. +/// +/// This is the cluster's coordinate within the grid, analogous to how +/// `blockIdx.x` is the block's coordinate within the grid. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %clusterid.x` +#[inline(always)] +pub fn cluster_idx_x() -> u32 { + __cluster_idxX() +} + +/// Get cluster's Y index within the grid. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %clusterid.y` +#[inline(always)] +pub fn cluster_idx_y() -> u32 { + __cluster_idxY() +} + +/// Get cluster's Z index within the grid. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %clusterid.z` +#[inline(always)] +pub fn cluster_idx_z() -> u32 { + __cluster_idxZ() +} + /// Get total number of clusters in the grid. /// /// # PTX @@ -178,6 +211,39 @@ pub fn num_clusters() -> u32 { __cluster_grid_dimX() * __cluster_grid_dimY() * __cluster_grid_dimZ() } +/// Get number of clusters in the grid's X dimension. +/// +/// This is the grid's size in clusters, analogous to how `gridDim.x` is the +/// grid's size in blocks. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %nclusterid.x` +#[inline(always)] +pub fn cluster_grid_dim_x() -> u32 { + __cluster_grid_dimX() +} + +/// Get number of clusters in the grid's Y dimension. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %nclusterid.y` +#[inline(always)] +pub fn cluster_grid_dim_y() -> u32 { + __cluster_grid_dimY() +} + +/// Get number of clusters in the grid's Z dimension. +/// +/// # PTX +/// +/// Lowers to: `mov.u32 %r, %nclusterid.z` +#[inline(always)] +pub fn cluster_grid_dim_z() -> u32 { + __cluster_grid_dimZ() +} + // ============================================================================= // Derived Helpers // ============================================================================= diff --git a/cuda-oxide/crates/cuda-device/src/f32x2.rs b/cuda-oxide/crates/cuda-device/src/f32x2.rs index 99cacf76ec..5988465fde 100644 --- a/cuda-oxide/crates/cuda-device/src/f32x2.rs +++ b/cuda-oxide/crates/cuda-device/src/f32x2.rs @@ -10,76 +10,39 @@ //! element. Pack/unpack helpers and [`Float2`] wrappers provide typed sugar on //! top of that catalog-owned register representation. -use crate::{Float2, ptx_asm}; +use crate::Float2; include!("generated/f32x2.rs"); /// Pack two scalar f32 values into the PTX f32x2 register representation. #[inline(always)] pub fn pack_f32x2(lo: f32, hi: f32) -> u64 { - let result: u64; - unsafe { - ptx_asm!( - "mov.b64 %0, {%1, %2};", - out("=l") result, - in("f") lo, - in("f") hi, - options(register_only), - ); - } - result + u64::from(lo.to_bits()) | (u64::from(hi.to_bits()) << 32) } /// Return the low scalar lane from a packed PTX f32x2 register. #[inline(always)] pub fn unpack_f32x2_lo(value: u64) -> f32 { - let result: f32; - unsafe { - ptx_asm!( - "mov.b64 {%0, _}, %1;", - out("=f") result, - in("l") value, - options(register_only), - ); - } - result + f32::from_bits(value as u32) } /// Return the high scalar lane from a packed PTX f32x2 register. #[inline(always)] pub fn unpack_f32x2_hi(value: u64) -> f32 { - let result: f32; - unsafe { - ptx_asm!( - "mov.b64 {_, %0}, %1;", - out("=f") result, - in("l") value, - options(register_only), - ); - } - result + f32::from_bits((value >> 32) as u32) } /// Add the two scalar lanes of a packed PTX f32x2 register. #[inline(always)] pub fn horizontal_add_f32x2_packed(value: u64) -> f32 { - let result: f32; - unsafe { - ptx_asm!( - "{ .reg .f32 lo, hi; mov.b64 {lo, hi}, %1; add.f32 %0, lo, hi; }", - out("=f") result, - in("l") value, - options(register_only), - ); - } - result + unpack_f32x2_lo(value) + unpack_f32x2_hi(value) } /// Two-lane round-to-nearest-even fused multiply-add (`a * b + c`). #[inline(always)] pub fn fma_float2(a: Float2, b: Float2, c: Float2) -> Float2 { // PTX models f32x2 as one 64-bit register pair. Keep that representation - // at the asm boundary so ptxas can reuse the input/output pairs directly. + // at the native arithmetic boundary so ptxas can reuse the pairs directly. let result = fma_f32x2( pack_f32x2(a.x(), a.y()), pack_f32x2(b.x(), b.y()), @@ -94,3 +57,43 @@ pub fn add_float2(a: Float2, b: Float2) -> Float2 { let result = add_f32x2(pack_f32x2(a.x(), a.y()), pack_f32x2(b.x(), b.y())); Float2::new([unpack_f32x2_lo(result), unpack_f32x2_hi(result)]) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn packed_carrier_preserves_all_lane_bits() { + // Include signed zero, subnormal boundaries, infinities, and NaN payloads. + let patterns = [ + 0x0000_0000, + 0x8000_0000, + 0x0000_0001, + 0x007f_ffff, + 0x0080_0000, + 0x3f80_0000, + 0x7f80_0000, + 0xff80_0000, + 0x7fc1_2345, + 0x7f81_2345, + 0xffc5_4321, + ]; + for lo in patterns { + for hi in patterns { + let packed = pack_f32x2(f32::from_bits(lo), f32::from_bits(hi)); + assert_eq!(packed, u64::from(lo) | (u64::from(hi) << 32)); + assert_eq!(unpack_f32x2_lo(packed).to_bits(), lo); + assert_eq!(unpack_f32x2_hi(packed).to_bits(), hi); + } + } + } + + #[test] + fn horizontal_sum_uses_both_scalar_lanes() { + assert_eq!(horizontal_add_f32x2_packed(pack_f32x2(1.25, -2.5)), -1.25); + assert_eq!( + horizontal_add_f32x2_packed(pack_f32x2(-0.0, -0.0)).to_bits(), + 0x8000_0000 + ); + } +} diff --git a/cuda-oxide/crates/cuda-device/src/generated/async_copy.rs b/cuda-oxide/crates/cuda-device/src/generated/async_copy.rs index 7a0f911070..821e27ba83 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/async_copy.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/async_copy.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::async_copy` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/atomic.rs b/cuda-oxide/crates/cuda-device/src/generated/atomic.rs index 67ee992d23..54e963cc62 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/atomic.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::atomic` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/bf16.rs b/cuda-oxide/crates/cuda-device/src/generated/bf16.rs index ca33fc259d..85254d74f7 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/bf16.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/bf16.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::bf16` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs b/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs index d8a67a77ca..2f952e598e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::bf16x2` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/clc.rs b/cuda-oxide/crates/cuda-device/src/generated/clc.rs index 966eb38a39..d11649381e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/clc.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::clc` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/cluster_barrier.rs b/cuda-oxide/crates/cuda-device/src/generated/cluster_barrier.rs index a704ad9652..de73e4179b 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/cluster_barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::cluster` to keep public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/cluster_memory.rs b/cuda-oxide/crates/cuda-device/src/generated/cluster_memory.rs index 3c8fa79ec3..634f66471f 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/cluster_memory.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::cluster` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/cluster_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/cluster_sreg.rs index 7820114fbd..a3cb4344b0 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/cluster_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/cluster_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // This file is included inside `cuda_device::cluster`. diff --git a/cuda-oxide/crates/cuda-device/src/generated/convert.rs b/cuda-oxide/crates/cuda-device/src/generated/convert.rs index 35970bd5d8..f592b2299e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/convert.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/convert.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::convert`. diff --git a/cuda-oxide/crates/cuda-device/src/generated/counted_barrier.rs b/cuda-oxide/crates/cuda-device/src/generated/counted_barrier.rs index b941b8a5bb..6a6677919b 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/counted_barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/counted_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::barrier` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/debug_control.rs b/cuda-oxide/crates/cuda-device/src/generated/debug_control.rs index a54274bb6c..3d5bab10c8 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/debug_control.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::debug` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/debug_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/debug_sreg.rs index 1be0474951..2789f4f675 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/debug_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/debug_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::debug` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs b/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs index d93a6e5925..d39b60563a 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::dotprod` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/f16.rs b/cuda-oxide/crates/cuda-device/src/generated/f16.rs index 40c0334d6e..acdeaf41b5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f16.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f16.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::f16` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs b/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs index 4118c4cb83..fa908ce6b5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::f16x2` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs b/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs index ae23507b76..6769a157a2 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::f32x2` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/fence.rs b/cuda-oxide/crates/cuda-device/src/generated/fence.rs index 6c0eae9600..06ef909025 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/fence.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/fence.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // This file is included inside `cuda_device::fence` so existing paths stay stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/float.rs b/cuda-oxide/crates/cuda-device/src/generated/float.rs index 7438a673c3..3e97a682d8 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/float.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/float.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::float`. diff --git a/cuda-oxide/crates/cuda-device/src/generated/grid_dependency.rs b/cuda-oxide/crates/cuda-device/src/generated/grid_dependency.rs index 02fd87444a..6006b8ecbc 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/grid_dependency.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/grid_dependency.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::grid` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/grid_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/grid_sreg.rs index 6745481b5a..56d3944256 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/grid_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/grid_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::grid` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs b/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs index 4ff8b83d36..77feb7b3eb 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::i16x2` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/int.rs b/cuda-oxide/crates/cuda-device/src/generated/int.rs index d69f7999e5..d13952b077 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/int.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/int.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::int` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs b/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs index ed3eb25000..1645ec9922 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/mbarrier_basic.rs b/cuda-oxide/crates/cuda-device/src/generated/mbarrier_basic.rs index 3de723ea87..d7f7432f5f 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/mbarrier_basic.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::barrier` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/mbarrier_extended.rs b/cuda-oxide/crates/cuda-device/src/generated/mbarrier_extended.rs index eb5b4cc387..e019729ce6 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/mbarrier_extended.rs @@ -4,11 +4,23 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::barrier` to keep existing paths stable. +/// Acquires cluster-shared memory at cluster scope, without ordering global memory. +/// Lowers to `fence.acquire.sync_restrict::shared::cluster.cluster;`. +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract. +/// The caller must pair this acquire with a compatible cluster-scoped synchronization +/// protocol. This excludes global memory and does not substitute for a proxy fence. +#[inline(never)] +pub unsafe fn fence_acquire_shared_cluster_cluster() { + unreachable!("fence_acquire_shared_cluster_cluster called outside CUDA kernel context") +} + /// Releases mbarrier initialization at cluster scope. /// Lowers to `fence.mbarrier_init.release.cluster;`. /// @@ -64,6 +76,19 @@ pub unsafe fn mbarrier_arrive_cluster(remote_bar_addr: u64) { unreachable!("mbarrier_arrive_cluster called outside CUDA kernel context") } +/// Arrives at a remote cluster-shared barrier with relaxed memory ordering. +/// Lowers to `mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [

];`. +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract. +/// This orders barrier counts and execution only, not data accesses. The caller +/// must provide the proxy/TMEM fences required by its synchronization protocol. +#[inline(never)] +pub unsafe fn mbarrier_arrive_cluster_relaxed(remote_bar_addr: u64) { + let _ = remote_bar_addr; + unreachable!("mbarrier_arrive_cluster_relaxed called outside CUDA kernel context") +} + /// Arrives at a CTA-shared barrier and adds expected transaction bytes. /// Lowers to `mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 , [
], ;`. /// @@ -123,6 +148,19 @@ pub unsafe fn mbarrier_try_wait_parity_cluster(bar: *const Barrier, parity: u32) unreachable!("mbarrier_try_wait_parity_cluster called outside CUDA kernel context") } +/// Tests barrier parity at cluster scope with relaxed memory ordering. +/// Lowers to `mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 , [
], ;`. +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract. +/// This orders barrier counts and execution only, not data accesses. The caller +/// must provide the proxy/TMEM fences required by its synchronization protocol. +#[inline(never)] +pub unsafe fn mbarrier_try_wait_parity_cluster_relaxed(bar: *const Barrier, parity: u32) -> bool { + let _ = (bar, parity); + unreachable!("mbarrier_try_wait_parity_cluster_relaxed called outside CUDA kernel context") +} + /// Suspends the executing thread for approximately the requested nanoseconds. /// Lowers to `nanosleep.u32 ;`. /// diff --git a/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs b/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs index 12fabfbd99..90c6bdff83 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/prmt.rs b/cuda-oxide/crates/cuda-device/src/generated/prmt.rs index 601a4ab647..2c3ce2bedf 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/prmt.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::prmt` to keep the public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/register_control.rs b/cuda-oxide/crates/cuda-device/src/generated/register_control.rs index 8dfc5f6e85..9b7fa2f900 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/register_control.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/register_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::thread` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs b/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs index e6bcd21bd6..580351a3c3 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep public paths stable. @@ -2531,6 +2531,45 @@ pub unsafe fn mma_m16n8k4_f32_tf32(c: [f32; 4], a: [u32; 2], b: u32) -> [f32; 4] unreachable!("mma_m16n8k4_f32_tf32 called outside CUDA kernel context") } +/// Multiplies warp-distributed E2M1 fragments with UE8M0 block scales and adds an F32 accumulator. +/// Lowers to `mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0`. C, A, and B are this lane's fragments in PTX register order. +/// Requires PTX 8.7 and `sm_120a`. +/// +/// # Safety +/// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited. +/// `c`, `a`, and `b` must contain this lane's fragments in the documented PTX layout. +/// `scale_a` and `scale_b` contain this lane's packed scale data. +/// For `scale_vec::2X`, `byte_id_a` and `byte_id_b` must each be `0` or `2`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined. +/// See the [PTX MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma). +/// This register-only operation is not a memory fence. +#[allow(clippy::too_many_arguments)] +#[must_use] +#[inline(never)] +pub unsafe fn mma_m16n8k64_mxf4_f32( + c: [f32; 4], + a: [u32; 4], + b: [u32; 2], + scale_a: u32, + byte_id_a: u16, + thread_id_a: u16, + scale_b: u32, + byte_id_b: u16, + thread_id_b: u16, +) -> [f32; 4] { + let _ = ( + c, + a, + b, + scale_a, + byte_id_a, + thread_id_a, + scale_b, + byte_id_b, + thread_id_b, + ); + unreachable!("mma_m16n8k64_mxf4_f32 called outside CUDA kernel context") +} + /// Multiplies warp-distributed signed A and signed B INT4 fragments and adds a wrapping s32 accumulator. /// Lowers to `mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32`. C, A, and B are this lane's fragments in PTX register order. /// Requires PTX 7.0 and `sm_80+`. diff --git a/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs index 2eaceca7a3..068fc93ca0 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/shared_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::shared` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs b/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs index 7dfb3bd1ed..dae6ecf5e5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wmma` to keep public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs index 3634cf8a0e..d279384792 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // This file is included lexically inside `cuda_device::thread` so existing diff --git a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs index 2863fbf1f9..cba095423c 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::tcgen05` to keep public paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs b/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs index 4efa57fc4f..d2ae333e1a 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::tcgen05` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/tma.rs b/cuda-oxide/crates/cuda-device/src/generated/tma.rs index d716ca7d9e..828ddd5d95 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::tma` to keep its public API stable. @@ -237,11 +237,11 @@ pub unsafe fn cp_async_bulk_tensor_1d_g2s( /// /// # Safety /// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. -/// `barrier` must point into the issuing CTA's shared memory, as for the other -/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. #[inline(never)] pub unsafe fn cp_async_bulk_tensor_1d_g2s_cta( - dst: SharedPtr, + dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, barrier: *mut Barrier, @@ -284,11 +284,11 @@ pub unsafe fn cp_async_bulk_tensor_2d_g2s( /// /// # Safety /// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. -/// `barrier` must point into the issuing CTA's shared memory, as for the other -/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. #[inline(never)] pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta( - dst: SharedPtr, + dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, @@ -368,11 +368,11 @@ pub unsafe fn cp_async_bulk_tensor_3d_g2s( /// /// # Safety /// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. -/// `barrier` must point into the issuing CTA's shared memory, as for the other -/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. #[inline(never)] pub unsafe fn cp_async_bulk_tensor_3d_g2s_cta( - dst: SharedPtr, + dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, @@ -421,11 +421,11 @@ pub unsafe fn cp_async_bulk_tensor_4d_g2s( /// /// # Safety /// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. -/// `barrier` must point into the issuing CTA's shared memory, as for the other -/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. #[inline(never)] pub unsafe fn cp_async_bulk_tensor_4d_g2s_cta( - dst: SharedPtr, + dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, @@ -480,12 +480,12 @@ pub unsafe fn cp_async_bulk_tensor_5d_g2s( /// /// # Safety /// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. -/// `barrier` must point into the issuing CTA's shared memory, as for the other -/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. #[inline(never)] #[allow(clippy::too_many_arguments)] pub unsafe fn cp_async_bulk_tensor_5d_g2s_cta( - dst: SharedPtr, + dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, diff --git a/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs index 6630593fb4..f4eb6d8e04 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/warp_sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::warp` to keep existing paths stable. diff --git a/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs b/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs index 68d9c4d252..2824167161 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::wgmma` to keep its public API stable. diff --git a/cuda-oxide/crates/cuda-device/src/lib.rs b/cuda-oxide/crates/cuda-device/src/lib.rs index 67ee46790e..6bb9d75486 100644 --- a/cuda-oxide/crates/cuda-device/src/lib.rs +++ b/cuda-oxide/crates/cuda-device/src/lib.rs @@ -10,6 +10,8 @@ // both here and in downstream device crates. extern crate self as cuda_device; +pub use cuda_repr::DeviceRepr; + pub use cuda_macros::{ cluster_launch, constant, convergent, cooperative_launch, cuda_module, device, gpu_printf, kernel, launch_bounds, launch_contract, ptx_asm, pure, readonly, tile_export, @@ -82,7 +84,7 @@ pub use cusimd::{CuSimd, Float2, Float4, TmemRegs4, TmemRegs32}; pub use disjoint::{__LaunchContractDisjointSlice, __LaunchContractDisjointSliceAbi}; pub use disjoint::{DisjointSlice, SpaceLayout}; pub use fence::*; -pub use shared::{DynamicSharedArray, SharedArray, SharedPtr}; +pub use shared::{DynamicSharedArray, SharedArray}; pub use tcgen05::{ TensorMemoryHandle, TmemAddress, TmemDeallocated, TmemF32x4, TmemF32x32, TmemGuard, TmemReady, TmemUninit, diff --git a/cuda-oxide/crates/cuda-device/src/shared.rs b/cuda-oxide/crates/cuda-device/src/shared.rs index 790e24078d..2010b8b7c1 100644 --- a/cuda-oxide/crates/cuda-device/src/shared.rs +++ b/cuda-oxide/crates/cuda-device/src/shared.rs @@ -233,34 +233,6 @@ impl SharedArray { let _ = shared; unreachable!("SharedArray::as_raw_mut_ptr called outside CUDA kernel context") } - - /// Returns a [`SharedPtr`] to the first element of a shared array. - /// - /// Unlike [`Self::as_raw_mut_ptr`], which returns a generic pointer, the - /// result keeps the array's CTA-shared address space in its type, so it - /// stays a shared-memory pointer when it is passed to helper functions. - /// - /// # Safety - /// - /// `shared` must point to a `static mut SharedArray` in the current CUDA - /// kernel, normally obtained with `&raw mut`. The returned pointer is valid - /// only within that kernel. This requirement is what makes the type's - /// claim true: a `SharedPtr` built from any other pointer hands - /// [`SharedPtr::addr_u32`] an address outside the shared window, where the - /// conversion yields an unspecified value. Callers must keep concurrent - /// accesses disjoint and use the CUDA synchronization required before - /// reading data written by another thread. - /// - /// ```rust,ignore - /// static mut TILE: SharedArray = SharedArray::UNINIT; - /// let tile: SharedPtr = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; - /// ``` - #[inline(always)] - pub unsafe fn shared_ptr(shared: *mut Self) -> SharedPtr { - SharedPtr { - ptr: shared as *mut SharedSlot, - } - } } impl Index for SharedArray { @@ -548,24 +520,6 @@ impl DynamicSharedArray { let _ = byte_offset; unreachable!("DynamicSharedArray::offset called outside CUDA kernel context") } - - /// Like [`Self::get`], but returns a [`SharedPtr`], which keeps the - /// CTA-shared address space in its type. - #[inline(always)] - pub fn shared_ptr() -> SharedPtr { - SharedPtr { - ptr: Self::get() as *mut SharedSlot, - } - } - - /// Like [`Self::offset`], but returns a [`SharedPtr`], which keeps the - /// CTA-shared address space in its type. - #[inline(always)] - pub fn shared_ptr_at(byte_offset: usize) -> SharedPtr { - SharedPtr { - ptr: Self::offset(byte_offset) as *mut SharedSlot, - } - } } /// Convert a generic-address pointer into its raw `.shared` window offset. @@ -599,6 +553,20 @@ pub unsafe fn cvta_generic_to_shared_offset(ptr: *const u8) -> u64 { unreachable!("cvta_generic_to_shared_offset called outside CUDA kernel context") } +/// Convert a generic pointer to its cluster shared-memory window offset. +/// +/// Unlike CTA shared offsets, this retains the remote CTA rank encoded by mapa. +/// +/// # Safety +/// The pointer must refer to shared memory in the current cluster. The caller +/// must ensure its allocation and cluster lifetime remain valid. +// Keep this definition in `shared`: importer recognition uses the exact path. +#[inline(never)] +pub unsafe fn cvta_generic_to_cluster_shared_offset(ptr: *const u8) -> u64 { + let _ = ptr; + unreachable!("cvta_generic_to_cluster_shared_offset called outside CUDA kernel context") +} + /// Convert a generic address into its 32-bit `.shared::cta` state-space address. /// /// Unlike [`cvta_generic_to_shared_offset`], which returns the `u64` carrier @@ -619,136 +587,4 @@ pub unsafe fn cvta_generic_to_shared_u32(ptr: *const u8) -> u32 { unreachable!("cvta_generic_to_shared_u32 called outside CUDA kernel context") } -// ============================================================================ -// SharedPtr - a pointer into the issuing CTA's shared memory, typed as such -// ============================================================================ - -/// The pointee of [`SharedPtr`]'s field. A pointer to `SharedSlot` is a -/// pointer into CTA-shared memory (`addrspace(3)`) by its type, as a pointer -/// to [`SharedArray`] is; `repr(transparent)` gives it `T`'s size, so pointer -/// arithmetic strides by `T`. -#[doc(hidden)] -#[repr(transparent)] -pub struct SharedSlot(T); - -/// A raw pointer into the shared memory of the issuing CTA (`.shared::cta`). -/// -/// A `*mut T` in device code is a generic pointer, which may address global, -/// local or shared memory, including another CTA's shared memory in the -/// cluster. `SharedPtr` instead carries its address space in its type: the -/// compiler knows it is a CTA-shared pointer wherever it goes, including -/// through calls to helper functions, and emits the `.shared::cta` forms of -/// instructions that take it without converting a generic address. -/// -/// It is created from shared memory the CTA owns, with the `unsafe` -/// [`SharedArray::shared_ptr`] or the safe [`DynamicSharedArray::shared_ptr`], -/// or from a generic pointer with the `unsafe` [`SharedPtr::from_raw`]. Like a raw -/// pointer, it may dangle; dereferencing it goes through [`SharedPtr::as_ptr`] -/// and `unsafe` code. -/// -/// Pointers from [`crate::cluster::map_shared_rank`], which may name another -/// CTA's shared memory, are not `SharedPtr`s. -/// -/// # Example -/// -/// ```rust,ignore -/// static mut TILE: SharedArray = SharedArray::UNINIT; -/// -/// fn store(tile: SharedPtr, i: usize, value: f32) { -/// // `tile` is still a shared pointer here, although `store` is a -/// // separate function. -/// unsafe { tile.add(i).as_ptr().write(value) }; -/// } -/// -/// let tile = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; -/// store(tile, thread::threadIdx_x() as usize, 1.0); -/// ``` -#[repr(transparent)] -pub struct SharedPtr { - ptr: *mut SharedSlot, -} - -impl Clone for SharedPtr { - #[inline(always)] - fn clone(&self) -> Self { - *self - } -} - -impl Copy for SharedPtr {} - -impl SharedPtr { - /// Converts a generic pointer to a `SharedPtr`, as CUDA C++'s - /// `__cvta_generic_to_shared` or CuTe's `make_smem_ptr` do. - /// - /// # Safety - /// - /// `ptr` must point into the shared memory of the issuing CTA. A pointer - /// to global or local memory, or to another CTA's shared memory, gives an - /// unspecified address. - #[inline(always)] - pub unsafe fn from_raw(ptr: *mut T) -> Self { - Self { - ptr: ptr as *mut SharedSlot, - } - } - - /// The generic pointer to the same memory, for ordinary loads and stores - /// and for APIs that take raw pointers. - #[inline(always)] - pub fn as_ptr(self) -> *mut T { - self.ptr as *mut T - } - - /// The pointer's 32-bit `.shared::cta` address, as taken by the - /// `_shared_u32` forms of CTA-shared instructions. - #[inline(always)] - pub fn addr_u32(self) -> u32 { - // SAFETY: `self` points into the issuing CTA's shared memory. - unsafe { cvta_generic_to_shared_u32(self.as_ptr() as *const u8) } - } - - /// Reinterprets the pointee type, keeping the address space. - #[inline(always)] - pub fn cast(self) -> SharedPtr { - SharedPtr { - ptr: self.ptr as *mut SharedSlot, - } - } - - /// The pointer `count` elements of `T` further on, as `<*mut T>::add`. - /// - /// # Safety - /// - /// As for `<*mut T>::add`. - #[inline(always)] - pub unsafe fn add(self, count: usize) -> Self { - Self { - ptr: unsafe { self.ptr.add(count) }, - } - } - - /// The pointer `count` elements of `T` away, as `<*mut T>::offset`. - /// - /// # Safety - /// - /// As for `<*mut T>::offset`. - #[inline(always)] - pub unsafe fn offset(self, count: isize) -> Self { - Self { - ptr: unsafe { self.ptr.offset(count) }, - } - } - - /// The pointer `count` bytes further on, as `<*mut T>::byte_add`. - /// - /// # Safety - /// - /// As for `<*mut T>::byte_add`. - #[inline(always)] - pub unsafe fn byte_add(self, count: usize) -> Self { - unsafe { self.cast::().add(count).cast::() } - } -} - include!("generated/shared_sreg.rs"); diff --git a/cuda-oxide/crates/cuda-device/src/tcgen05.rs b/cuda-oxide/crates/cuda-device/src/tcgen05.rs index 3f6ea71d85..4254fb4551 100644 --- a/cuda-oxide/crates/cuda-device/src/tcgen05.rs +++ b/cuda-oxide/crates/cuda-device/src/tcgen05.rs @@ -331,7 +331,9 @@ impl TmemGuard { pub const fn n_cols(&self) -> u32 { N_COLS } +} +impl TmemGuard { /// Deallocate the TMEM. /// /// **All threads in the block should call this.** Only warp 0 performs diff --git a/cuda-oxide/crates/cuda-device/src/thread.rs b/cuda-oxide/crates/cuda-device/src/thread.rs index 0b3f33c435..ef0f01aaaf 100644 --- a/cuda-oxide/crates/cuda-device/src/thread.rs +++ b/cuda-oxide/crates/cuda-device/src/thread.rs @@ -1404,7 +1404,14 @@ pub fn __unroll_config() { // No runtime code is generated. } +/// Internal sentinel factor for `#[unroll(never)]`. +#[doc(hidden)] +pub const UNROLL_NEVER: u32 = u32::MAX - 1; + const fn validate_unroll_factor(factor: u32) { + if factor == UNROLL_NEVER { + return; + } assert!( factor == 0 || (factor >= 2 && factor <= 1024), "partial unroll factor must be in 2..=1024, or 0 for full unrolling" diff --git a/cuda-oxide/crates/cuda-device/src/tma.rs b/cuda-oxide/crates/cuda-device/src/tma.rs index fdbc2fb0be..87e12a00b0 100644 --- a/cuda-oxide/crates/cuda-device/src/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/tma.rs @@ -113,7 +113,6 @@ //! - **sm_120 (Blackwell)**: Latest TMA capabilities use crate::barrier::Barrier; -use crate::shared::SharedPtr; // ============================================================================= // TMA Descriptor Type @@ -127,8 +126,10 @@ use crate::shared::SharedPtr; /// /// # Size /// -/// - CUDA 12.0-12.x: 128 bytes, 64-byte aligned -/// - CUDA 13.0+: 128 bytes, 128-byte aligned +/// The kernel parameter representation is 128 bytes, 64-byte aligned. +/// The CUDA driver binding's `CUtensorMap` output type may require stronger +/// alignment (128 bytes in CUDA 13). Encode into that type, then copy its +/// opaque words here; do not cast this storage to the driver output type. /// /// # Safety /// @@ -143,6 +144,14 @@ pub struct TmaDescriptor { } impl TmaDescriptor { + /// Adopt the words of a successfully encoded CUDA tensor map. + /// + /// # Safety + /// `words` must come from the CUDA driver encoder. This method does not + /// validate descriptor contents or the referenced allocation's lifetime. + pub const unsafe fn from_raw_words(words: [u64; 16]) -> Self { + Self { _opaque: words } + } /// Create an uninitialized descriptor. /// /// # Safety diff --git a/cuda-oxide/crates/cuda-device/src/vector.rs b/cuda-oxide/crates/cuda-device/src/vector.rs index edb513e6c3..200ba123a7 100644 --- a/cuda-oxide/crates/cuda-device/src/vector.rs +++ b/cuda-oxide/crates/cuda-device/src/vector.rs @@ -67,9 +67,12 @@ //! taking `&[f32]`, and take the view **inside the kernel** with //! [`as_vectors`] / [`as_vectors_mut`]. This needs no host-side type at all. //! - Or reinterpret the flat buffer once at the boundary with `cuda-core`'s -//! `DeviceBuffer::cast_chunks`, using a host-defined `#[repr(C, align(16))]` -//! element that implements `DeviceCopy`, and launch a kernel that takes that -//! element type directly. +//! `DeviceBuffer::cast_chunks`, which requires `bytemuck::AnyBitPattern` of +//! its target. These types implement `bytemuck::Pod`, so +//! `floats.cast_chunks::()` works, and so does a host-defined +//! `#[repr(C, align(16))]` element with its own `Pod` impl. Launch a kernel +//! that takes the vector type directly, and cast back to the scalar type +//! before reading the result on the host. //! //! # Related "N x scalar" types //! @@ -124,6 +127,14 @@ macro_rules! define_vector { const LANES: usize = $lanes; } + // SAFETY: a public array of plain integers or floats whose size equals + // the alignment (checked below), so there is no padding and every bit + // pattern is a valid value. `Pod` is what lets cuda-core's + // `DeviceBuffer::cast_chunks` regroup a flat scalar buffer into this + // type. + unsafe impl bytemuck::Zeroable for $name {} + unsafe impl bytemuck::Pod for $name {} + impl $name { /// Build a vector from its lanes. #[must_use] diff --git a/cuda-oxide/crates/cuda-device/tests/atomic_device_repr.rs b/cuda-oxide/crates/cuda-device/tests/atomic_device_repr.rs new file mode 100644 index 0000000000..aba5a162c8 --- /dev/null +++ b/cuda-oxide/crates/cuda-device/tests/atomic_device_repr.rs @@ -0,0 +1,40 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![feature(f16)] + +use cuda_device::DeviceRepr; +use cuda_device::atomic::*; + +#[test] +fn atomic_representations_match_their_scalars() { + fn check, Inner>() { + assert_eq!(size_of::(), size_of::()); + assert_eq!(align_of::(), align_of::()); + assert!(!core::mem::needs_drop::()); + assert!(!core::mem::needs_drop::()); + } + + macro_rules! check_atomics { + ($($atomic:ty => $inner:ty),* $(,)?) => { + $(check::<$atomic, $inner>();)* + }; + } + check_atomics!( + DeviceAtomicU32 => u32, DeviceAtomicI32 => i32, + DeviceAtomicU64 => u64, DeviceAtomicI64 => i64, + DeviceAtomicF16 => f16, DeviceAtomicF32 => f32, DeviceAtomicF64 => f64, + BlockAtomicU32 => u32, BlockAtomicI32 => i32, + BlockAtomicU64 => u64, BlockAtomicI64 => i64, + BlockAtomicF16 => f16, BlockAtomicF32 => f32, BlockAtomicF64 => f64, + SystemAtomicU32 => u32, SystemAtomicI32 => i32, + SystemAtomicU64 => u64, SystemAtomicI64 => i64, + SystemAtomicF16 => f16, SystemAtomicF32 => f32, SystemAtomicF64 => f64, + ); +} + +#[test] +fn atomic_shared_reference_conversions_are_rejected() { + let t = trybuild::TestCases::new(); + t.compile_fail("tests/compile_fail/atomic_shared_references.rs"); +} diff --git a/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.rs b/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.rs new file mode 100644 index 0000000000..b6102824eb --- /dev/null +++ b/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.rs @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use bytemuck::TransparentWrapper; +use cuda_device::atomic::{DeviceAtomicF32, DeviceAtomicU32}; + +fn wrap(integer: &u32, float: &f32) { + // Shared scalar references do not permit atomic mutation. + let _ = DeviceAtomicU32::wrap_ref(integer); + let _ = DeviceAtomicF32::wrap_ref(float); +} + +fn peel(integer: &DeviceAtomicU32, float: &DeviceAtomicF32) { + // Atomic mutation must not coexist with ordinary shared references. + let _ = DeviceAtomicU32::peel_ref(integer); + let _ = DeviceAtomicF32::peel_ref(float); +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.stderr b/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.stderr new file mode 100644 index 0000000000..505f1d6a10 --- /dev/null +++ b/cuda-oxide/crates/cuda-device/tests/compile_fail/atomic_shared_references.stderr @@ -0,0 +1,81 @@ +error[E0599]: no associated function or constant named `wrap_ref` found for struct `DeviceAtomicU32` in the current scope + --> tests/compile_fail/atomic_shared_references.rs:9:30 + | +9 | let _ = DeviceAtomicU32::wrap_ref(integer); + | ^^^^^^^^ associated function or constant not found in `DeviceAtomicU32` + | +note: if you're trying to build a new `DeviceAtomicU32`, consider using `DeviceAtomicU32::new` which returns `DeviceAtomicU32` + --> src/atomic.rs + | + | pub const fn new(val: $ty) -> Self { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +... + | / define_integer_atomic! { + | | /// 32-bit unsigned atomic, **device scope** (`.gpu`). + | | pub struct DeviceAtomicU32(u32); + | | } + | |_- in this macro invocation + = note: this error originates in the macro `define_integer_atomic` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: no associated function or constant named `wrap_ref` found for struct `DeviceAtomicF32` in the current scope + --> tests/compile_fail/atomic_shared_references.rs:10:30 + | +10 | let _ = DeviceAtomicF32::wrap_ref(float); + | ^^^^^^^^ associated function or constant not found in `DeviceAtomicF32` + | +note: if you're trying to build a new `DeviceAtomicF32`, consider using `DeviceAtomicF32::new` which returns `DeviceAtomicF32` + --> src/atomic.rs + | + | pub const fn new(val: $ty) -> Self { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +... + | / define_float_atomic! { + | | /// 32-bit float atomic, **device scope** (`.gpu`). + | | /// + | | /// `swap` lowers to hardware `atom.exch.b32`; this API does not expose +... | + | | pub struct DeviceAtomicF32(f32); + | | } + | |_- in this macro invocation + = note: this error originates in the macro `define_float_atomic` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: no associated function or constant named `peel_ref` found for struct `DeviceAtomicU32` in the current scope + --> tests/compile_fail/atomic_shared_references.rs:15:30 + | +15 | let _ = DeviceAtomicU32::peel_ref(integer); + | ^^^^^^^^ associated function or constant not found in `DeviceAtomicU32` + | +note: if you're trying to build a new `DeviceAtomicU32`, consider using `DeviceAtomicU32::new` which returns `DeviceAtomicU32` + --> src/atomic.rs + | + | pub const fn new(val: $ty) -> Self { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +... + | / define_integer_atomic! { + | | /// 32-bit unsigned atomic, **device scope** (`.gpu`). + | | pub struct DeviceAtomicU32(u32); + | | } + | |_- in this macro invocation + = note: this error originates in the macro `define_integer_atomic` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: no associated function or constant named `peel_ref` found for struct `DeviceAtomicF32` in the current scope + --> tests/compile_fail/atomic_shared_references.rs:16:30 + | +16 | let _ = DeviceAtomicF32::peel_ref(float); + | ^^^^^^^^ associated function or constant not found in `DeviceAtomicF32` + | +note: if you're trying to build a new `DeviceAtomicF32`, consider using `DeviceAtomicF32::new` which returns `DeviceAtomicF32` + --> src/atomic.rs + | + | pub const fn new(val: $ty) -> Self { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +... + | / define_float_atomic! { + | | /// 32-bit float atomic, **device scope** (`.gpu`). + | | /// + | | /// `swap` lowers to hardware `atom.exch.b32`; this API does not expose +... | + | | pub struct DeviceAtomicF32(f32); + | | } + | |_- in this macro invocation + = note: this error originates in the macro `define_float_atomic` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-device/tests/compile_fail/swizzle_non_involutive.stderr b/cuda-oxide/crates/cuda-device/tests/compile_fail/swizzle_non_involutive.stderr index b5d29acc9d..599474d478 100644 --- a/cuda-oxide/crates/cuda-device/tests/compile_fail/swizzle_non_involutive.stderr +++ b/cuda-oxide/crates/cuda-device/tests/compile_fail/swizzle_non_involutive.stderr @@ -6,8 +6,6 @@ error[E0080]: evaluation panicked: Swizzle requires M + B <= S, otherwi | ::: src/swizzle.rs | - | impl Swizzle { - | /// Compile-time check that the swizzle is an involution. | const INVOLUTION: () = assert!( | ____________________________- | | M + B <= S, diff --git a/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_runtime_mismatch.stderr b/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_runtime_mismatch.stderr index fb4421dbec..4ca8d8067f 100644 --- a/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_runtime_mismatch.stderr +++ b/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_runtime_mismatch.stderr @@ -13,17 +13,3 @@ note: method defined here | | pub fn get_mut<'kernel>(&mut self, idx: ThreadIndex<'kernel, IndexSpace>) -> Option<&mut T> { | ^^^^^^^ -note: the `get_mut` call is resolved to the method in `cuda_device::DisjointSlice<'a, T, IndexSpace>`, shadowing the method of the same name on the inherent impl for `cuda_device::DisjointSlice<'a, T, cuda_device::Runtime2DIndex>` - --> tests/compile_fail/thread_index_runtime_mismatch.rs:18:17 - | -18 | let _ = out.get_mut(idx); - | ^^^^^^^ refers to `cuda_device::DisjointSlice::<'a, T, IndexSpace>::get_mut` - = note: additionally, there is 1 other available method that isn't in scope -help: you might have meant to call one of the other methods; you can use the fully-qualified path to call one of them explicitly - | -18 - let _ = out.get_mut(idx); -18 + let _ = DisjointSlice::get_mut(&mut out, idx); - | -18 - let _ = out.get_mut(idx); -18 + let _ = <_ as std::slice::SliceIndex<_>>::get_mut(out, idx); - | diff --git a/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_stride_mismatch.stderr b/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_stride_mismatch.stderr index 96a0802bce..7eed21fb06 100644 --- a/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_stride_mismatch.stderr +++ b/cuda-oxide/crates/cuda-device/tests/compile_fail/thread_index_stride_mismatch.stderr @@ -13,17 +13,3 @@ note: method defined here | | pub fn get_mut<'kernel>(&mut self, idx: ThreadIndex<'kernel, IndexSpace>) -> Option<&mut T> { | ^^^^^^^ -note: the `get_mut` call is resolved to the method in `cuda_device::DisjointSlice<'a, T, IndexSpace>`, shadowing the method of the same name on the inherent impl for `cuda_device::DisjointSlice<'a, T, cuda_device::Runtime2DIndex>` - --> tests/compile_fail/thread_index_stride_mismatch.rs:11:17 - | -11 | let _ = out.get_mut(idx); - | ^^^^^^^ refers to `cuda_device::DisjointSlice::<'a, T, IndexSpace>::get_mut` - = note: additionally, there is 1 other available method that isn't in scope -help: you might have meant to call one of the other methods; you can use the fully-qualified path to call one of them explicitly - | -11 - let _ = out.get_mut(idx); -11 + let _ = DisjointSlice::get_mut(&mut out, idx); - | -11 - let _ = out.get_mut(idx); -11 + let _ = <_ as std::slice::SliceIndex<_>>::get_mut(out, idx); - | diff --git a/cuda-oxide/crates/cuda-host/Cargo.toml b/cuda-oxide/crates/cuda-host/Cargo.toml index e24c5fc2ba..3906514968 100644 --- a/cuda-oxide/crates/cuda-host/Cargo.toml +++ b/cuda-oxide/crates/cuda-host/Cargo.toml @@ -9,6 +9,7 @@ license.workspace = true repository.workspace = true [dependencies] +cuda-device = { workspace = true } half = { workspace = true } # `host` is a cuda-macros default, but name it: cuda-host is precisely the # crate that needs the generated host surface, and it must keep it even if @@ -32,7 +33,6 @@ async = ["dep:cuda-async", "cuda-macros/async"] [dev-dependencies] libloading = { workspace = true } -cuda-device = { workspace = true } # The crates.io release cuda-core re-exports bundle types from; tests # construct payloads with it so they use the same crate instance. oxide-artifacts = { workspace = true } diff --git a/cuda-oxide/crates/cuda-host/README.md b/cuda-oxide/crates/cuda-host/README.md index 7d921a8bb6..ef8ee57bbe 100644 --- a/cuda-oxide/crates/cuda-host/README.md +++ b/cuda-oxide/crates/cuda-host/README.md @@ -8,7 +8,7 @@ keep `#[kernel]` on the actual GPU entry points, then load the embedded device artifact as a typed Rust value. ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; @@ -206,7 +206,7 @@ Enable the `async` feature to generate async launch methods. They use the same scalar mapping, but take no stream argument: ```rust -use cuda_async::simt::device_operation::DeviceOperation; +use cuda_async::device_operation::DeviceOperation; let module = kernels::load_async(0)?; let launch = unsafe { @@ -222,7 +222,7 @@ launch.sync()?; ``` For async launches, device-slice parameters accept either `DeviceBuffer` or -`cuda_async::simt::device_box::DeviceBox<[T]>`. The mutable +`cuda_async::device_box::DeviceBox<[T]>`. The mutable `AsyncKernelLaunchBuilder` collects arguments and options. Finalizing it with a raw configuration is unsafe and produces an immutable `AsyncKernelLaunch<'_>`; geometry cannot be changed after that point. Rust keeps referenced buffers and @@ -361,5 +361,5 @@ specific 8x8 tile arrangements: - [cuda-device](../cuda-device/) -- device-side intrinsics - [cuda-macros](../cuda-macros/) -- proc-macro implementations -- [cuda-core](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-core) -- shared CUDA driver API crate; `DeviceBuffer`, `simt::LaunchConfig` -- [cuda-async](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-async) -- shared async crate; the SIMT model lives under `cuda_async::simt` +- [cuda-core](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-core) -- shared CUDA driver API crate; `DeviceBuffer`, `LaunchConfig` +- [cuda-async](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-async) -- shared async crate; one runtime for Tile and SIMT diff --git a/cuda-oxide/crates/cuda-host/src/entry_registry.rs b/cuda-oxide/crates/cuda-host/src/entry_registry.rs index 657da3159a..b12885c3af 100644 --- a/cuda-oxide/crates/cuda-host/src/entry_registry.rs +++ b/cuda-oxide/crates/cuda-host/src/entry_registry.rs @@ -195,7 +195,7 @@ fn type_id_divergence_message( format!( "generic kernel PTX entry `{requested}` was not found, but the loaded module DOES \ contain `{base}` under a different type-identity hash: {found}. The host \ - (core::intrinsics::reflection::type_id) and the device backend (tcx.type_id_hash) disagreed on \ + (core::intrinsics::type_id) and the device backend (tcx.type_id_hash) disagreed on \ the kernel's `_TID_` name, which means the TypeId-to-u128 extraction in \ cuda_host::type_id is no longer value-correct on this toolchain. Host and device \ code must be built by the same pinned nightly; re-verify the TypeId contract in \ @@ -329,10 +329,7 @@ mod tests { "{message}" ); // The message must identify the two hash producers and the remedy. - assert!( - message.contains("core::intrinsics::reflection::type_id"), - "{message}" - ); + assert!(message.contains("core::intrinsics::type_id"), "{message}"); assert!(message.contains("tcx.type_id_hash"), "{message}"); assert!(message.contains("same pinned nightly"), "{message}"); assert!(message.contains("compute_kernel_export_name"), "{message}"); diff --git a/cuda-oxide/crates/cuda-host/src/launch.rs b/cuda-oxide/crates/cuda-host/src/launch.rs index c3a2266731..4ce744a60f 100644 --- a/cuda-oxide/crates/cuda-host/src/launch.rs +++ b/cuda-oxide/crates/cuda-host/src/launch.rs @@ -164,7 +164,7 @@ pub trait CudaKernel { /// The bound on `ptx_name` is intentionally just whatever the kernel /// itself declared — typically `Copy` on the value-passed generics. No /// `'static` is added: the helper `cuda_host::type_id_u128` uses -/// `core::intrinsics::reflection::type_id`, which is bound `T: ?Sized` rather than +/// `core::intrinsics::type_id`, which is bound `T: ?Sized` rather than /// `T: 'static`, so closures that borrow non-`'static` data still go /// through the typed launch path (just like they did under the legacy /// `type_name`-based scheme). @@ -466,6 +466,14 @@ pub fn push_kernel_row_width_device_slice( /// `UnsafeCell`-based or atomic element types under their synchronization /// contract. /// +/// `storage()` must return the allocation covering this exact view, with the +/// same CUDA context and shared access tracker as every other safe alias. +/// When it returns `None`, the implementation must independently retain the +/// allocation and exclude conflicting accesses until GPU completion, even +/// when the launch future is dropped or forgotten. Returning `false` from +/// `shared_access_writes()` additionally promises that the shared reference +/// cannot write any device memory in this range. +/// /// Implementing this trait is unsafe because generated launch methods trust /// the pointer and length without inspecting the allocation: /// @@ -491,6 +499,20 @@ pub unsafe trait KernelSliceArg { /// Returns the number of `T` elements in the allocation. fn len(&self) -> usize; + /// An independently retained allocation and access tracker for this argument. + /// Implementations returning None must themselves keep memory and exclusion + /// valid through GPU completion, including dropped and forgotten futures. + fn storage(&self) -> Option { + None + } + + /// Shared references to interior-mutable elements (including device atomics) + /// can write. Such arguments require exclusive tracking against ordinary + /// reads/copies; the conservative default also covers unknown custom types. + fn shared_access_writes(&self) -> bool { + true + } + /// Returns true when the allocation has zero elements. fn is_empty(&self) -> bool { self.len() == 0 @@ -514,6 +536,13 @@ pub unsafe trait KernelSliceArgMut: KernelSliceArg {} // SAFETY: DeviceBuffer owns the reported allocation and keeps its CUDA context // alive; its pointer and length accessors describe that allocation exactly. unsafe impl KernelSliceArg for cuda_core::DeviceBuffer { + fn storage(&self) -> Option { + Some(self.storage().clone()) + } + fn shared_access_writes(&self) -> bool { + self.shared_access_writes() + } + type Elem = T; fn cu_deviceptr(&self) -> cuda_core::sys::CUdeviceptr { @@ -533,7 +562,15 @@ unsafe impl KernelSliceArgMut for cuda_core::DeviceBuffer {} #[cfg(feature = "async")] // SAFETY: DeviceBox owns the reported allocation and its raw constructor // requires the pointer, element count, and device ordinal to be truthful. -unsafe impl KernelSliceArg for cuda_async::simt::device_box::DeviceBox<[T]> { +unsafe impl KernelSliceArg for cuda_async::device_box::DeviceBox<[T]> { + fn shared_access_writes(&self) -> bool { + self.shared_access_writes() + } + + fn storage(&self) -> Option { + Some(self.storage().clone()) + } + type Elem = T; fn cu_deviceptr(&self) -> cuda_core::sys::CUdeviceptr { @@ -548,7 +585,7 @@ unsafe impl KernelSliceArg for cuda_async::simt::device_box::DeviceBox< #[cfg(feature = "async")] // SAFETY: &mut DeviceBox provides exclusive host authority to launch device // writes through this adapter for the duration of the operation. -unsafe impl KernelSliceArgMut for cuda_async::simt::device_box::DeviceBox<[T]> {} +unsafe impl KernelSliceArgMut for cuda_async::device_box::DeviceBox<[T]> {} #[cfg(feature = "async")] // SAFETY: Arc only extends the lifetime of B and delegates both truthful @@ -557,6 +594,13 @@ unsafe impl KernelSliceArg for Arc where B: KernelSliceArg + ?Sized, { + fn storage(&self) -> Option { + self.as_ref().storage() + } + fn shared_access_writes(&self) -> bool { + self.as_ref().shared_access_writes() + } + type Elem = B::Elem; fn cu_deviceptr(&self) -> cuda_core::sys::CUdeviceptr { @@ -568,26 +612,45 @@ where } } +#[cfg(feature = "async")] +// SAFETY: views retain their parent's storage and preserve typed byte extents; +// launch retention protects the allocation independently of the Rust borrow. +unsafe impl KernelSliceArg for cuda_core::DeviceBufferView<'_, T> { + type Elem = T; + fn cu_deviceptr(&self) -> cuda_core::sys::CUdeviceptr { + self.cu_deviceptr() + } + fn len(&self) -> usize { + self.len() + } + fn storage(&self) -> Option { + Some(self.storage().clone()) + } + fn shared_access_writes(&self) -> bool { + self.shared_access_writes() + } +} + #[doc(hidden)] #[cfg(feature = "async")] pub fn new_async_kernel_launch_builder<'a>( func: cuda_core::CudaFunction, -) -> cuda_async::simt::launch::AsyncKernelLaunchBuilder<'a> { - cuda_async::simt::launch::AsyncKernelLaunchBuilder::new(Arc::new(func)) +) -> cuda_async::launch::AsyncKernelLaunchBuilder<'a> { + cuda_async::launch::AsyncKernelLaunchBuilder::new(Arc::new(func)) } #[doc(hidden)] #[cfg(feature = "async")] pub fn new_owned_async_kernel_launch( - launch: cuda_async::simt::launch::AsyncKernelLaunch<'static>, + launch: cuda_async::launch::AsyncKernelLaunch<'static>, resources: R, -) -> cuda_async::simt::launch::OwnedAsyncKernelLaunch { - cuda_async::simt::launch::OwnedAsyncKernelLaunch::new(launch, resources) +) -> cuda_async::launch::OwnedAsyncKernelLaunch { + cuda_async::launch::OwnedAsyncKernelLaunch::new(launch, resources) } /// Async operation carrying a validated, kernel-branded launch. /// -/// The underlying [`cuda_async::simt::launch::AsyncKernelLaunch`] is immutable after +/// The underlying [`cuda_async::launch::AsyncKernelLaunch`] is immutable after /// its builder is finalized. This wrapper additionally carries the prepared /// contract and checks the stream selected by the async scheduler against the /// prepared function's CUDA context immediately before submission. @@ -596,7 +659,7 @@ pub struct PreparedAsyncKernelLaunch<'a, Contract> where Contract: cuda_core::KernelLaunchContract, { - launch: cuda_async::simt::launch::AsyncKernelLaunch<'a>, + launch: cuda_async::launch::AsyncKernelLaunch<'a>, prepared: cuda_core::PreparedLaunch, } @@ -619,7 +682,7 @@ where R: Send, Contract: cuda_core::KernelLaunchContract, { - launch: cuda_async::simt::launch::OwnedAsyncKernelLaunch, + launch: cuda_async::launch::OwnedAsyncKernelLaunch, prepared: cuda_core::PreparedLaunch, } @@ -645,7 +708,7 @@ where #[doc(hidden)] #[cfg(feature = "async")] pub unsafe fn new_prepared_async_kernel_launch<'a, Contract>( - launch: cuda_async::simt::launch::AsyncKernelLaunch<'a>, + launch: cuda_async::launch::AsyncKernelLaunch<'a>, prepared: cuda_core::PreparedLaunch, ) -> PreparedAsyncKernelLaunch<'a, Contract> where @@ -663,7 +726,7 @@ where #[doc(hidden)] #[cfg(feature = "async")] pub unsafe fn new_prepared_owned_async_kernel_launch( - launch: cuda_async::simt::launch::OwnedAsyncKernelLaunch, + launch: cuda_async::launch::OwnedAsyncKernelLaunch, prepared: cuda_core::PreparedLaunch, ) -> PreparedOwnedAsyncKernelLaunch where @@ -674,7 +737,7 @@ where } #[cfg(feature = "async")] -impl<'a, Contract> cuda_async::simt::device_operation::DeviceOperation +impl<'a, Contract> cuda_async::device_operation::DeviceOperation for PreparedAsyncKernelLaunch<'a, Contract> where Contract: cuda_core::KernelLaunchContract, @@ -683,76 +746,80 @@ where unsafe fn execute( self, - context: &cuda_async::simt::device_operation::ExecutionContext, - ) -> Result<(), cuda_async::simt::error::DeviceError> { + context: &cuda_async::device_operation::ExecutionContext, + ) -> Result<(), cuda_async::error::DeviceError> { self.prepared .validate_stream(context.get_cuda_stream()) - .map_err(|error| cuda_async::simt::error::DeviceError::Launch(error.to_string()))?; - unsafe { - cuda_async::simt::device_operation::DeviceOperation::execute(self.launch, context) - } + .map_err(|error| cuda_async::error::DeviceError::Launch(error.to_string()))?; + unsafe { cuda_async::device_operation::DeviceOperation::execute(self.launch, context) } } } +#[cfg(feature = "async")] +impl cuda_async::device_operation::GraphNode + for PreparedAsyncKernelLaunch<'_, Contract> +{ +} + +#[cfg(feature = "async")] +impl cuda_async::device_operation::GraphNode + for PreparedOwnedAsyncKernelLaunch +{ +} + #[cfg(feature = "async")] impl<'a, Contract> IntoFuture for PreparedAsyncKernelLaunch<'a, Contract> where Contract: cuda_core::KernelLaunchContract, { - type Output = Result<(), cuda_async::simt::error::DeviceError>; - type IntoFuture = cuda_async::simt::device_future::DeviceFuture<(), Self>; + type Output = Result<(), cuda_async::error::DeviceError>; + type IntoFuture = cuda_async::device_future::DeviceFuture<(), Self>; fn into_future(self) -> Self::IntoFuture { - match cuda_async::simt::device_context::with_default_device_policy(|policy| { - cuda_async::simt::scheduling_policies::SchedulingPolicy::schedule(policy, self) + match cuda_async::device_context::with_default_device_policy(|policy| { + cuda_async::device_operation::DeviceOp::schedule(self, policy) }) { Ok(Ok(future)) => future, - Ok(Err(error)) | Err(error) => { - cuda_async::simt::device_future::DeviceFuture::failed(error) - } + Ok(Err(error)) | Err(error) => cuda_async::device_future::DeviceFuture::failed(error), } } } #[cfg(feature = "async")] -impl cuda_async::simt::device_operation::DeviceOperation +impl cuda_async::device_operation::DeviceOperation for PreparedOwnedAsyncKernelLaunch where - R: Send + 'static, + R: Send, Contract: cuda_core::KernelLaunchContract, { type Output = R; unsafe fn execute( self, - context: &cuda_async::simt::device_operation::ExecutionContext, - ) -> Result { + context: &cuda_async::device_operation::ExecutionContext, + ) -> Result { self.prepared .validate_stream(context.get_cuda_stream()) - .map_err(|error| cuda_async::simt::error::DeviceError::Launch(error.to_string()))?; - unsafe { - cuda_async::simt::device_operation::DeviceOperation::execute(self.launch, context) - } + .map_err(|error| cuda_async::error::DeviceError::Launch(error.to_string()))?; + unsafe { cuda_async::device_operation::DeviceOperation::execute(self.launch, context) } } } #[cfg(feature = "async")] impl IntoFuture for PreparedOwnedAsyncKernelLaunch where - R: Send + 'static, + R: Send, Contract: cuda_core::KernelLaunchContract, { - type Output = Result; - type IntoFuture = cuda_async::simt::device_future::DeviceFuture; + type Output = Result; + type IntoFuture = cuda_async::device_future::DeviceFuture; fn into_future(self) -> Self::IntoFuture { - match cuda_async::simt::device_context::with_default_device_policy(|policy| { - cuda_async::simt::scheduling_policies::SchedulingPolicy::schedule(policy, self) + match cuda_async::device_context::with_default_device_policy(|policy| { + cuda_async::device_operation::DeviceOp::schedule(self, policy) }) { Ok(Ok(future)) => future, - Ok(Err(error)) | Err(error) => { - cuda_async::simt::device_future::DeviceFuture::failed(error) - } + Ok(Err(error)) | Err(error) => cuda_async::device_future::DeviceFuture::failed(error), } } } @@ -760,7 +827,7 @@ where #[doc(hidden)] #[cfg(feature = "async")] pub fn set_async_kernel_cluster_dim( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, cluster_dim: (u32, u32, u32), ) { launch.set_cluster_dim(cluster_dim); @@ -769,7 +836,7 @@ pub fn set_async_kernel_cluster_dim( #[doc(hidden)] #[cfg(feature = "async")] pub fn set_async_kernel_cooperative( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, cooperative: bool, ) { launch.set_cooperative(cooperative); @@ -778,7 +845,7 @@ pub fn set_async_kernel_cooperative( #[doc(hidden)] #[cfg(feature = "async")] pub fn push_async_kernel_scalar<'a, T: KernelScalar + 'a>( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'a>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'a>, value: T, ) { if std::mem::size_of::() != 0 { @@ -796,11 +863,14 @@ pub fn push_async_kernel_scalar<'a, T: KernelScalar + 'a>( #[cfg(feature = "async")] #[track_caller] pub fn push_async_read_only_device_slice( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, buffer: &B, ) where B: KernelSliceArg + ?Sized, { + if let Some(storage) = buffer.storage() { + launch.retain_buffer(storage, buffer.shared_access_writes()); + } let len = buffer.len(); launch.push_scalar_arg(kernel_slice_device_ptr::( buffer.cu_deviceptr(), @@ -819,11 +889,14 @@ pub fn push_async_read_only_device_slice( #[cfg(feature = "async")] #[track_caller] pub fn push_async_writable_device_slice( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, buffer: &mut B, ) where B: KernelSliceArgMut + ?Sized, { + if let Some(storage) = buffer.storage() { + launch.retain_buffer(storage, true); + } let len = buffer.len(); launch.push_scalar_arg(kernel_slice_device_ptr::( buffer.cu_deviceptr(), @@ -845,9 +918,10 @@ pub fn push_async_writable_device_slice( #[cfg(feature = "async")] #[track_caller] pub fn push_async_row_width_device_slice( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, bound: RowWidth<'_, T>, ) { + launch.retain_buffer(bound.buffer.storage().clone(), true); let (ptr, len, width) = row_width_device_buffer_arg(bound); launch.push_scalar_arg(ptr); launch.push_scalar_arg(len); @@ -864,11 +938,14 @@ pub fn push_async_row_width_device_slice( #[cfg(feature = "async")] #[track_caller] pub fn push_async_owned_row_width_device_slice( - launch: &mut cuda_async::simt::launch::AsyncKernelLaunchBuilder<'_>, + launch: &mut cuda_async::launch::AsyncKernelLaunchBuilder<'_>, bound: &mut RowWidthOwned, ) where B: KernelSliceArgMut, { + if let Some(storage) = bound.buffer.storage() { + launch.retain_buffer(storage, true); + } let len = bound.buffer.len(); launch.push_scalar_arg(kernel_slice_device_ptr::( bound.buffer.cu_deviceptr(), @@ -883,15 +960,15 @@ pub fn push_async_owned_row_width_device_slice( pub fn load_cuda_module_from_async_context( device_id: usize, f: F, -) -> Result +) -> Result where F: FnOnce(&Arc) -> Result, { - cuda_async::simt::device_context::with_cuda_context(device_id, |ctx| f(ctx))?.map_err(|error| { + cuda_async::device_context::with_cuda_context(device_id, |ctx| f(ctx))?.map_err(|error| { if let crate::EmbeddedModuleError::Driver(error) = error { - cuda_async::simt::error::DeviceError::Driver(error) + cuda_async::error::DeviceError::Driver(error) } else { - cuda_async::simt::error::DeviceError::KernelCache(error.to_string()) + cuda_async::error::DeviceError::KernelCache(error.to_string()) } }) } @@ -904,10 +981,10 @@ where pub fn load_kernel_module_async( name: &str, device_id: usize, -) -> Result, cuda_async::simt::error::DeviceError> { - cuda_async::simt::device_context::with_cuda_context(device_id, |ctx| { +) -> Result, cuda_async::error::DeviceError> { + cuda_async::device_context::with_cuda_context(device_id, |ctx| { crate::ltoir::load_kernel_module(ctx, name) - .map_err(|error| cuda_async::simt::error::DeviceError::KernelCache(error.to_string())) + .map_err(|error| cuda_async::error::DeviceError::KernelCache(error.to_string())) })? } diff --git a/cuda-oxide/crates/cuda-host/src/lib.rs b/cuda-oxide/crates/cuda-host/src/lib.rs index 1fa90924a7..cfdaea6eed 100644 --- a/cuda-oxide/crates/cuda-host/src/lib.rs +++ b/cuda-oxide/crates/cuda-host/src/lib.rs @@ -3,17 +3,17 @@ * SPDX-License-Identifier: Apache-2.0 */ -// Required for `type_id_u128`'s use of `core::intrinsics::reflection::type_id`. The +// Required for `type_id_u128`'s use of `core::intrinsics::type_id`. The // intrinsic is the only way to obtain the same 128-bit hash the backend // uses while keeping the bound at `T: ?Sized` (stable `TypeId::of` would // force `T: 'static` on every kernel marker — see `type_id.rs` for why // that would silently reject non-`'static` borrowing closures). // // We accept the `internal_features` warning here: cuda-oxide already ships -// `rustc_codegen_cuda` against `rustc_private` and pins a nightly -// toolchain, so the broader project is firmly inside rustc's internal API -// surface. Anyone trying to lift this helper to a stable crate will hit -// the same gate and have to make the same trade-off there. +// `rustc_codegen_cuda` against `rustc_private` and pins its Rust toolchain +// with `RUSTC_BOOTSTRAP=1`, so the broader project is firmly inside rustc's +// internal API surface. Anyone trying to lift this helper to a standalone +// crate will hit the same gate and have to make the same trade-off there. #![feature(core_intrinsics)] #![allow(internal_features)] @@ -51,7 +51,7 @@ //! ```ignore //! use cuda_device::{kernel, thread, DisjointSlice}; //! use cuda_host::cuda_module; -//! use cuda_core::simt::LaunchConfig; +//! use cuda_core::LaunchConfig; //! use cuda_core::{CudaContext, DeviceBuffer}; //! //! #[cuda_module] @@ -122,16 +122,12 @@ pub use launch::{ set_async_kernel_cluster_dim, set_async_kernel_cooperative, }; -/// The shared async crate, re-exported whole. Its root is cutile's Tile API; -/// the SIMT surface cuda-host builds on is `cuda_host::cuda_async::simt::*`. -/// Same-named root items such as `DeviceError` or `init_device_contexts` are -/// the Tile ones and do not interoperate with the generated launch methods. +/// The shared async runtime crate, re-exported whole for convenience. +/// Both Tile and SIMT launch methods interoperate through this root surface. #[cfg(feature = "async")] pub use cuda_async; #[cfg(feature = "async")] -pub use cuda_async::simt::launch::{ - AsyncKernelLaunch, AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch, -}; +pub use cuda_async::launch::{AsyncKernelLaunch, AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; pub use embedded::{ EmbeddedModuleError, load_all_ptx_bundles_merged, load_embedded_module, @@ -154,7 +150,7 @@ pub use cuda_macros::{cuda_launch, cuda_module}; /// Re-export of [`cuda_macros::cuda_launch_async`]. /// -/// Builds a lazy `cuda_async::simt::launch::AsyncKernelLaunch`. Raw launch +/// Builds a lazy `cuda_async::launch::AsyncKernelLaunch`. Raw launch /// configuration is not tied to the kernel's indexing assumptions, so the /// macro must be called inside `unsafe`. Stream assignment is deferred to the /// scheduling policy -- call `.sync()` to block or `.await` to suspend. diff --git a/cuda-oxide/crates/cuda-host/src/type_id.rs b/cuda-oxide/crates/cuda-host/src/type_id.rs index e9aa917b55..6d480a4bf4 100644 --- a/cuda-oxide/crates/cuda-host/src/type_id.rs +++ b/cuda-oxide/crates/cuda-host/src/type_id.rs @@ -10,7 +10,7 @@ //! API would force a `T: 'static` bound on the kernel marker, which would in //! turn reject perfectly valid non-`'static` borrowing closures (e.g. a kernel //! launcher capturing `&[f32]` from a stack frame the caller keeps alive -//! across the launch). The `core::intrinsics::reflection::type_id` form has bound +//! across the launch). The `core::intrinsics::type_id` form has bound //! `T: ?Sized` — i.e. no `'static` requirement — and produces the exact same //! 128-bit value that `tcx.type_id_hash` does for that type, because both go //! through the same `erase_and_anonymize_regions` + stable-hash pipeline. @@ -20,10 +20,10 @@ //! every monomorphized type and const argument, so one hash covers the complete //! specialization without inventing a parallel encoding for const values. //! -//! Framing note for future contributors: `core::intrinsics::reflection::type_id` is an +//! Framing note for future contributors: `core::intrinsics::type_id` is an //! internal API and requires `#![feature(core_intrinsics)]` on the owning //! crate. cuda-oxide already ships against `rustc_private` and pins a -//! nightly toolchain, so this is inside our existing risk surface — but the +//! toolchain, so this is inside our existing risk surface, but the //! helper cannot be lifted into a stable-feeling utility crate without //! re-introducing the feature gate there. @@ -54,7 +54,7 @@ static GENERIC_KERNEL_NAMES: OnceLock /// caller (the borrow must outlive `stream.synchronize()`). #[inline] pub fn type_id_u128() -> u128 { - let id = const { core::intrinsics::reflection::type_id::() }; + let id = const { core::intrinsics::type_id::() }; unsafe { core::mem::transmute::(id) } } diff --git a/cuda-oxide/crates/cuda-host/tests/cuda_module_api.rs b/cuda-oxide/crates/cuda-host/tests/cuda_module_api.rs index e41e9c70d7..8c700ff37e 100644 --- a/cuda-oxide/crates/cuda-host/tests/cuda_module_api.rs +++ b/cuda-oxide/crates/cuda-host/tests/cuda_module_api.rs @@ -5,12 +5,12 @@ #![allow(dead_code, non_upper_case_globals, clippy::needless_lifetimes)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaStream, DeviceBuffer, LaunchConfig1D}; #[cfg(feature = "async")] -use cuda_host::cuda_async::simt::device_box::DeviceBox; +use cuda_host::cuda_async::device_box::DeviceBox; #[cfg(feature = "async")] -use cuda_host::cuda_async::simt::device_operation::DeviceOperation; +use cuda_host::cuda_async::device_operation::DeviceOperation; use cuda_host::cuda_module; use cuda_macros::{cooperative_launch, kernel, launch_contract}; diff --git a/cuda-oxide/crates/cuda-host/tests/shared_launch_lifetimes.rs b/cuda-oxide/crates/cuda-host/tests/shared_launch_lifetimes.rs new file mode 100644 index 0000000000..01d2d3cdd6 --- /dev/null +++ b/cuda-oxide/crates/cuda-host/tests/shared_launch_lifetimes.rs @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![cfg(feature = "async")] + +use cuda_async::{ + device_operation::DeviceOperation, launch::AsyncKernelLaunchBuilder, + scheduling_policies::SingleStream, +}; +use cuda_core::{Context, DeviceBuffer, LaunchConfig}; +use std::{ + future::Future, + pin::Pin, + sync::{Arc, mpsc}, + task::{Context as TaskContext, Waker}, + time::Duration, +}; + +const PTX: &str = ".version 7.0\n.target sm_52\n.address_size 64\n.visible .entry use_slice(.param .u64 ptr, .param .u64 len) { ret; }"; + +struct Gate(mpsc::Sender<()>); +impl Drop for Gate { + fn drop(&mut self) { + let _ = self.0.send(()); + } +} + +fn pending_launch_keeps_access(forget: bool, atomic: bool) { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + let other = context.new_stream().unwrap(); + let function = Arc::new( + context + .load_module_from_ptx_src(PTX) + .unwrap() + .load_function("use_slice") + .unwrap(), + ); + let buffer = DeviceBuffer::from_host(&stream, &[1u32; 32]).unwrap(); + let storage = buffer.storage().clone(); + let mut builder = AsyncKernelLaunchBuilder::new(function); + if atomic { + let buffer = buffer.wrap_elem::(); + cuda_host::push_async_read_only_device_slice(&mut builder, &buffer); + drop(buffer); + } else { + let mut buffer = buffer; + cuda_host::push_async_writable_device_slice(&mut builder, &mut buffer); + drop(buffer); + } + let (tx, rx) = mpsc::channel(); + stream + .launch_host_function(move || { + let _ = rx.recv_timeout(Duration::from_secs(10)); + }) + .unwrap(); + let gate = Gate(tx); + // SAFETY: PTX declares exactly the pointer/length packet above and no + // memory accesses. The descriptors still model a writable kernel argument, + // so this test can observe exclusion without deliberately racing GPU memory. + let launch = unsafe { + builder.finalize_unchecked(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + }; + let mut future = launch + .schedule(&SingleStream::from_stream(stream.clone())) + .unwrap(); + assert!( + Pin::new(&mut future) + .poll(&mut TaskContext::from_waker(Waker::noop())) + .is_pending() + ); + if forget { + std::mem::forget(future); + } else { + drop(future); + } + assert!( + storage.acquire(&other, false, false).is_err(), + "early future release lost a writable argument's access lease" + ); + drop(gate); + stream.synchronize().unwrap(); + if forget { + assert!( + storage.acquire(&other, false, false).is_err(), + "a forgotten submission conservatively retains its access leases" + ); + } else { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while storage.acquire(&other, false, false).is_err() { + assert!( + std::time::Instant::now() < deadline, + "completed work failed to release exclusion" + ); + std::thread::sleep(Duration::from_millis(5)); + } + } +} + +#[test] +#[ignore = "requires a live CUDA GPU"] +fn dropped_borrowed_launch_retains_buffers_and_accesses() { + pending_launch_keeps_access(false, false); +} +#[test] +#[ignore = "requires a live CUDA GPU"] +fn forgotten_borrowed_launch_retains_buffers_and_accesses() { + pending_launch_keeps_access(true, false); +} +#[test] +#[ignore = "requires a live CUDA GPU"] +fn shared_atomic_arguments_exclude_ordinary_reads_after_drop() { + pending_launch_keeps_access(false, true); +} + +#[test] +#[ignore = "requires a live CUDA GPU"] +fn device_box_preserves_concurrent_read_permissions() { + let context = Context::new(0).unwrap(); + let stream = context.new_stream().unwrap(); + let other = context.new_stream().unwrap(); + let function = Arc::new( + context + .load_module_from_ptx_src(PTX) + .unwrap() + .load_function("use_slice") + .unwrap(), + ); + let buffer = cuda_async::device_box::DeviceBox::from_buffer( + DeviceBuffer::from_host(&stream, &[1u32; 32]).unwrap(), + ); + let raw_ptr = unsafe { cuda_core::memory::malloc_async(stream.cu_stream(), 32 * 4).unwrap() }; + stream.synchronize().unwrap(); + let raw_box = unsafe { + cuda_async::device_box::DeviceBox::<[u32]>::from_raw_parts(raw_ptr, 32, stream.clone()) + }; + assert!(!raw_box.shared_access_writes()); + let make_launch = || { + let mut builder = AsyncKernelLaunchBuilder::new(function.clone()); + cuda_host::push_async_read_only_device_slice(&mut builder, &buffer); + cuda_host::push_async_read_only_device_slice(&mut builder, &raw_box); + // SAFETY: the PTX's pointer/length signature matches the first descriptor; + // the no-op kernel never dereferences the modeled shared slices. + unsafe { builder.finalize_unchecked(LaunchConfig::for_num_elems(32)) } + }; + let (tx, rx) = mpsc::channel(); + stream + .launch_host_function(move || { + let _ = rx.recv_timeout(Duration::from_secs(10)); + }) + .unwrap(); + let gate = Gate(tx); + let mut first = make_launch() + .schedule(&SingleStream::from_stream(stream.clone())) + .unwrap(); + assert!( + Pin::new(&mut first) + .poll(&mut TaskContext::from_waker(Waker::noop())) + .is_pending() + ); + // A read on another stream must succeed while the first read is gated. + // Treating DeviceBox's ordinary u32 elements as interior-mutable would fail. + make_launch().sync_on(&other).unwrap(); + drop(gate); + cuda_async::futures::executor::block_on(first).unwrap(); +} diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs index 52398b3382..a12f4846d2 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/matrix.rs @@ -175,6 +175,7 @@ pub enum RegisterMmaKind { Standard, F8f6f4, Mxf8f6f4, + Mxf4, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs index 78d47948ab..7c6edea917 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/memory.rs @@ -12,11 +12,14 @@ pub enum MbarrierExtendedOperation { ArriveExpectTxCta, ArriveExpectTxCluster, ArriveRemoteCluster, + ArriveRemoteClusterRelaxed, TryWaitTokenCta, TryWaitParityCta, TryWaitParityCluster, + TryWaitParityClusterRelaxed, FenceProxyAsyncSharedCta, FenceMbarrierInitReleaseCluster, + FenceAcquireSharedClusterCluster, FenceProxyAsyncGenericReleaseSharedCtaCluster, FenceProxyAsyncGenericAcquireSharedClusterCluster, Nanosleep, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs index 2a7dd95c58..be1b534f97 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs @@ -168,6 +168,8 @@ pub struct StmatrixAdmissionVariant { pub struct MbarrierExtendedAdmission { pub llvm_evidence_profile: String, pub libnvvm_evidence_profile: String, + #[serde(default)] + pub relaxed_libnvvm_evidence_profile: Option, pub runtime_validation: RuntimeValidation, #[serde(rename = "variant")] pub variants: Vec, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/probe.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/probe.rs index a4e5520c43..eb1f387362 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/probe.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/probe.rs @@ -1799,6 +1799,22 @@ fn parse_ptxas_version(stdout: &str, stderr: &str) -> Result { Ok(version.to_owned()) } +fn llvm_major_minor_series(version: &str) -> Option<&str> { + let rest = version.strip_prefix("LLVM version ")?; + let mut parts = rest.split(['.', '-']); + let major = parts.next()?; + let minor = parts.next()?; + if major.is_empty() + || !major.bytes().all(|b| b.is_ascii_digit()) + || minor.is_empty() + || !minor.bytes().all(|b| b.is_ascii_digit()) + { + return None; + } + let prefix_len = "LLVM version ".len() + major.len() + 1 + minor.len(); + Some(&version[..prefix_len]) +} + fn validate_backend_identity( expected_version: &str, expected_commit: &RustcCommit, @@ -1807,8 +1823,12 @@ fn validate_backend_identity( let ProbeBackendIdentity::SelectedEvidence(actual) = actual else { return Ok(()); }; + let same_version_or_series = actual.tool.version == expected_version + || (llvm_major_minor_series(&actual.tool.version).is_some() + && llvm_major_minor_series(&actual.tool.version) + == llvm_major_minor_series(expected_version)); ensure!( - actual.tool.version == expected_version, + same_version_or_series, "rust-toolchain llc version mismatch: selected evidence records {expected_version:?}, found {:?}; use an explicit `--llc` only for a comparison probe", actual.tool.version ); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/collector_targets.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/collector_targets.rs index 213bec7e94..7c5e957878 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/collector_targets.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/collector_targets.rs @@ -191,6 +191,7 @@ pub(super) fn generated_intrinsic_variant(record: &CatalogIntrinsic) -> String { RegisterMmaKind::Standard => "GeneratedRegisterMmaKind::Standard", RegisterMmaKind::F8f6f4 => "GeneratedRegisterMmaKind::F8f6f4", RegisterMmaKind::Mxf8f6f4 => "GeneratedRegisterMmaKind::Mxf8f6f4", + RegisterMmaKind::Mxf4 => "GeneratedRegisterMmaKind::Mxf4", }; let accumulator = match mma.accumulator { RegisterMmaAccumulator::F16 => "GeneratedRegisterMmaAccumulator::F16", @@ -571,7 +572,7 @@ impl GeneratedIntrinsicTarget { replace_exact_render_fragment( &mut output, "pub enum GeneratedRegisterMmaAccumulator { F32, F64, S32 }", - "pub enum GeneratedRegisterMmaOperation { Multiply, AndPopc, XorPopc }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F16, F32, F64, S32 }", + "pub enum GeneratedRegisterMmaOperation { Multiply, AndPopc, XorPopc }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4, Mxf4 }\n#[derive(Debug, Clone, Copy, PartialEq, Eq)]\npub enum GeneratedRegisterMmaAccumulator { F16, F32, F64, S32 }", ); replace_exact_render_fragment( &mut output, @@ -762,7 +763,7 @@ impl GeneratedIntrinsicTarget { replace_exact_render_fragment( &mut output, " };\n let accumulator_matches = match accumulator {", - " };\n let operation_matches = match mma_operation {\n GeneratedRegisterMmaOperation::Multiply => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::Multiply,\n GeneratedRegisterMmaOperation::AndPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::AndPopc,\n GeneratedRegisterMmaOperation::XorPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::XorPopc,\n };\n let kind_matches = match kind {\n GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard,\n GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4,\n GeneratedRegisterMmaKind::Mxf8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf8f6f4,\n };\n let accumulator_matches = match accumulator {", + " };\n let operation_matches = match mma_operation {\n GeneratedRegisterMmaOperation::Multiply => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::Multiply,\n GeneratedRegisterMmaOperation::AndPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::AndPopc,\n GeneratedRegisterMmaOperation::XorPopc => op.operation_or_multiply(ctx) == RegisterMmaOperationAttr::XorPopc,\n };\n let kind_matches = match kind {\n GeneratedRegisterMmaKind::Standard => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Standard,\n GeneratedRegisterMmaKind::F8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::F8f6f4,\n GeneratedRegisterMmaKind::Mxf8f6f4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf8f6f4,\n GeneratedRegisterMmaKind::Mxf4 => op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf4,\n };\n let accumulator_matches = match accumulator {", ); replace_exact_render_fragment( &mut output, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs index a396dc2042..c5575349fa 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs @@ -8,10 +8,10 @@ use crate::model::{ ClusterBarrierOrdering, ClusterMemoryOperation, CpAsyncCachePolicy, CpAsyncControlOperation, CpAsyncMbarrierOperation, CpAsyncSourceSize, DebugControlOperation, MbarrierBasicOperation, MbarrierExtendedOperation, PackedAluFormat, PackedAtomicFormat, PackedConversionSourceFormat, - RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaOverflow, ScalarMathFormat, - SparseMmaAdapter, SparseMmaCompatibilitySource, SparseMmaOverflow, Tcgen05LdShape, - Tcgen05MmaForm, Tcgen05MmaSelectorLayout, Tcgen05Operation, TmaAdapter, TmaOperation, - WgmmaControlMode, + RegisterMmaAdapter, RegisterMmaCompatibilitySource, RegisterMmaKind, RegisterMmaOverflow, + ScalarMathFormat, SparseMmaAdapter, SparseMmaCompatibilitySource, SparseMmaOverflow, + Tcgen05LdShape, Tcgen05MmaForm, Tcgen05MmaSelectorLayout, Tcgen05Operation, TmaAdapter, + TmaOperation, WgmmaControlMode, }; use crate::render::common::{hardware_target_label, llvm, rust_header, source_label}; use crate::render::families::{ @@ -102,10 +102,12 @@ pub(super) fn render_compat_register_mma(catalog: &CatalogFile, hash: &str) -> S "/// `c`, `a`, and `b` must contain this lane's fragments in the documented PTX layout.\n", ); if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { - output.push_str( - "/// `scale_a` and `scale_b` contain this lane's packed scale data.\n\ - /// For `scale_vec::1X`, `byte_id_a` and `byte_id_b` must be in `0..=3`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined.\n", - ); + output.push_str("/// `scale_a` and `scale_b` contain this lane's packed scale data.\n"); + output.push_str(if mma.kind == Some(RegisterMmaKind::Mxf4) { + "/// For `scale_vec::2X`, `byte_id_a` and `byte_id_b` must each be `0` or `2`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined.\n" + } else { + "/// For `scale_vec::1X`, `byte_id_a` and `byte_id_b` must be in `0..=3`, `thread_id_a` in `0..=1`, and `thread_id_b` in `0..=3`; other values make the PTX operation undefined.\n" + }); } writeln!( output, @@ -942,7 +944,7 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { ); if operation.is_cta_g2s() { output.push_str( - "/// `barrier` must point into the issuing CTA's shared memory, as for the other\n/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.\n", + "/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other\n/// `mbarrier` operations.\n", ); } } else if is_s2g { @@ -971,15 +973,8 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { output.push_str("#[inline(never)]\n"); if is_g2s { let dimensions = dimensions.unwrap(); - // The CTA-local forms take a typed CTA-shared destination; the - // raw ABI keeps `*mut u8`, as `*mut Barrier` keeps `*mut u64`. - let destination = if operation.is_cta_g2s() { - "dst: SharedPtr" - } else { - "dst: *mut u8" - }; let mut arguments = vec![ - destination.to_owned(), + "dst: *mut u8".to_owned(), "tensor_map: *const TmaDescriptor".to_owned(), ]; let mut values = vec!["dst".to_owned(), "tensor_map".to_owned()]; @@ -1874,6 +1869,16 @@ pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) output.push_str( "/// The caller must satisfy the barrier, scope, and memory-ordering contract.\n", ); + if matches!( + contract.operation, + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed + | MbarrierExtendedOperation::TryWaitParityClusterRelaxed + ) { + output.push_str("/// This orders barrier counts and execution only, not data accesses. The caller\n/// must provide the proxy/TMEM fences required by its synchronization protocol.\n"); + } + if contract.operation == MbarrierExtendedOperation::FenceAcquireSharedClusterCluster { + output.push_str("/// The caller must pair this acquire with a compatible cluster-scoped synchronization\n/// protocol. This excludes global memory and does not substitute for a proxy fence.\n"); + } output.push_str("#[inline(never)]\n"); match contract.operation { MbarrierExtendedOperation::ArriveExpectTxCta @@ -1886,7 +1891,8 @@ pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) .unwrap(); output.push_str(" let _ = (bar, bytes);\n"); } - MbarrierExtendedOperation::ArriveRemoteCluster => { + MbarrierExtendedOperation::ArriveRemoteCluster + | MbarrierExtendedOperation::ArriveRemoteClusterRelaxed => { writeln!( output, "pub unsafe fn {}(remote_bar_addr: u64) {{", @@ -1905,7 +1911,8 @@ pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) output.push_str(" let _ = (bar, token);\n"); } MbarrierExtendedOperation::TryWaitParityCta - | MbarrierExtendedOperation::TryWaitParityCluster => { + | MbarrierExtendedOperation::TryWaitParityCluster + | MbarrierExtendedOperation::TryWaitParityClusterRelaxed => { writeln!( output, "pub unsafe fn {}(bar: *const Barrier, parity: u32) -> bool {{", @@ -1916,6 +1923,7 @@ pub(super) fn render_compat_mbarrier_extended(catalog: &CatalogFile, hash: &str) } MbarrierExtendedOperation::FenceProxyAsyncSharedCta | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => { writeln!(output, "pub unsafe fn {}() {{", record.rust.name).unwrap(); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs index 58a7cd672b..9347f32f74 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/mma.rs @@ -515,7 +515,7 @@ pub enum RegisterMmaOperationAttr { Multiply, AndPopc, XorPopc } #[pliron_attr(name = "nvvm.register_mma_kind", format, verifier = "succ")] #[derive(PartialEq, Eq, Clone, Debug, Hash)] -pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 } +pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4, Mxf4 } #[pliron_attr(name = "nvvm.register_mma_accumulator", format, verifier = "succ")] #[derive(PartialEq, Eq, Clone, Debug, Hash)] diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs index 587f2ccb2a..db02c0af44 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/warp_sync.rs @@ -784,7 +784,8 @@ fn verify_mbarrier_extended( "PointerU32ToU64", " pub fn build(ctx: &mut Context, barrier: Value, bytes: Value) -> Ptr {\n let token_ty = IntegerType::get(ctx, 64, Signedness::Unsigned);\n Operation::new(ctx, Self::get_concrete_op_info(), vec![token_ty.into()], vec![barrier, bytes], vec![], 0)\n }\n", ), - MbarrierExtendedOperation::ArriveRemoteCluster => ( + MbarrierExtendedOperation::ArriveRemoteCluster + | MbarrierExtendedOperation::ArriveRemoteClusterRelaxed => ( 1, 0, "RawU64ToVoid", @@ -797,7 +798,8 @@ fn verify_mbarrier_extended( " pub fn build(ctx: &mut Context, barrier: Value, token: Value) -> Ptr {\n let result_ty = IntegerType::get(ctx, 1, Signedness::Signless);\n Operation::new(ctx, Self::get_concrete_op_info(), vec![result_ty.into()], vec![barrier, token], vec![], 0)\n }\n", ), MbarrierExtendedOperation::TryWaitParityCta - | MbarrierExtendedOperation::TryWaitParityCluster => ( + | MbarrierExtendedOperation::TryWaitParityCluster + | MbarrierExtendedOperation::TryWaitParityClusterRelaxed => ( 2, 1, "PointerU32ToI1", @@ -805,6 +807,7 @@ fn verify_mbarrier_extended( ), MbarrierExtendedOperation::FenceProxyAsyncSharedCta | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => ( 0, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/mma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/mma.rs index 87ddd80577..8f2e5ae26b 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/mma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/mma.rs @@ -157,6 +157,7 @@ pub(in crate::render) fn register_mma_attr_variants( RegisterMmaKind::Standard => "RegisterMmaKindAttr::Standard", RegisterMmaKind::F8f6f4 => "RegisterMmaKindAttr::F8f6f4", RegisterMmaKind::Mxf8f6f4 => "RegisterMmaKindAttr::Mxf8f6f4", + RegisterMmaKind::Mxf4 => "RegisterMmaKindAttr::Mxf4", }; let accumulator = match mma.accumulator { RegisterMmaAccumulator::F16 => "RegisterMmaAccumulatorAttr::F16", diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/raw_abi.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/raw_abi.rs index 20fac09b2e..48eb860bdf 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/raw_abi.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/raw_abi.rs @@ -7,8 +7,8 @@ use crate::model::{ CatalogFile, ClusterBarrierOrdering, ClusterMemoryOperation, CpAsyncControlOperation, CpAsyncMbarrierOperation, CpAsyncSourceSize, ExecutionControlOperation, LdmatrixParticipation, LdmatrixShape, MbarrierBasicOperation, MbarrierExtendedAdapter, RegisterMmaAdapter, - RegisterMmaOverflow, SparseMmaOverflow, Tcgen05LdShape, Tcgen05Operation, TmaAdapter, - WarpShuffleValueKind, WgmmaControlMode, + RegisterMmaKind, RegisterMmaOverflow, SparseMmaOverflow, Tcgen05LdShape, Tcgen05Operation, + TmaAdapter, WarpShuffleValueKind, WgmmaControlMode, }; use crate::render::common::{hardware_target_label, rust_header, source_label}; use crate::render::families::{ @@ -215,9 +215,13 @@ pub(super) fn render_raw_abi(catalog: &CatalogFile, hash: &str) -> Result String { RuntimeValidation::Executed => "executed on a GPU", }; if mma.adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 { + let selectors = if mma.kind == Some(RegisterMmaKind::Mxf4) { + "For `scale_vec::2X`, byte selectors must each be `0` or `2`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`." + } else { + "For `scale_vec::1X`, byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`." + }; writeln!( output, - "- `{}` takes C, A, B, scale-A data/selectors, and scale-B data/selectors in PTX operand order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. For `scale_vec::1X`, byte selectors must be in `0..=3`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`. Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", + "- `{}` takes C, A, B, scale-A data/selectors, and scale-B data/selectors in PTX operand order, performs {operation}, and lowers to one convergent, register-only `{}` instruction. {selectors} Every non-exited warp lane must execute the same instruction and qualifiers. Integer overflow is {overflow}; runtime validation is {runtime}.", record.id, expected_ptx_head(record), ) diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs index b45c3afbae..5a6037138a 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mma.rs @@ -22,6 +22,36 @@ use crate::render::families::{ use crate::util::read_json; use std::path::Path; +#[test] +fn dense_mxf4_k64_rendering_preserves_registers_and_scale_selectors() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve_candidate( + &repo_root, + "mma_m16n8k64_mxf4_f32", + "test-llvm", + &"a".repeat(64), + "sm_120a", + "+ptx87", + ) + .unwrap() + .catalog; + let record = &catalog.intrinsics[0]; + assert_eq!(intrinsic_marker(&catalog, record), "v1:i1039"); + assert_eq!( + register_mma_template(record), + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7}, $14, {$15, $16}, $17, {$18, $19};" + ); + assert_eq!( + register_mma_constraints(record), + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h" + ); + let compatibility = render_compat_register_mma(&catalog, "test-hash"); + assert!(compatibility.contains("pub unsafe fn mma_m16n8k64_mxf4_f32(")); + let importer = render_importer(&catalog, "test-hash"); + assert!(importer.contains("set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::Mxf4)")); + assert!(generated_intrinsic_variant(record).contains("kind: GeneratedRegisterMmaKind::Mxf4")); +} + #[test] fn stmatrix_rendering_preserves_all_four_public_and_backend_contracts() { let catalog = catalog_with_stmatrix(); @@ -345,9 +375,9 @@ fn register_mma_rendering_preserves_apis_order_convergence_and_variants() { let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); let catalog = crate::resolve::resolve(&repo_root).unwrap(); validate_renderable(&catalog).unwrap(); - assert_eq!(catalog.intrinsics.len(), 1038); + assert_eq!(catalog.intrinsics.len(), 1042); let records: Vec<_> = register_mmas(&catalog).collect(); - assert_eq!(records.len(), 154); + assert_eq!(records.len(), 155); let generated_records = records .iter() .copied() @@ -364,7 +394,7 @@ fn register_mma_rendering_preserves_apis_order_convergence_and_variants() { == RegisterMmaCompatibilitySource::ExistingStub }) .collect::>(); - assert_eq!(generated_records.len(), 149); + assert_eq!(generated_records.len(), 150); assert_eq!(existing_records.len(), 5); let raw = render_raw_abi(&catalog, "test-hash").unwrap(); @@ -376,7 +406,7 @@ fn register_mma_rendering_preserves_apis_order_convergence_and_variants() { assert!(raw.contains("Signed accumulator overflow clamps")); let compatibility = render_compat_register_mma(&catalog, "test-hash"); - assert_eq!(compatibility.matches("pub unsafe fn ").count(), 149); + assert_eq!(compatibility.matches("pub unsafe fn ").count(), 150); for record in generated_records { let argument_names: &[&str] = if record.register_mma.as_ref().unwrap().adapter == RegisterMmaAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32 @@ -425,7 +455,7 @@ fn register_mma_rendering_preserves_apis_order_convergence_and_variants() { assert!(dialect.contains("RegisterMmaOperationAttr::Multiply")); assert!(dialect.contains("RegisterMmaOperationAttr::AndPopc")); assert!(dialect.contains("RegisterMmaOperationAttr::XorPopc")); - assert!(dialect.contains("pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4 }")); + assert!(dialect.contains("pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4, Mxf4 }")); assert!(dialect.contains("kind_or_inferred")); assert!(dialect.contains("RegisterMmaAccumulatorAttr::F16")); assert!(dialect.contains("operation_or_multiply")); @@ -590,7 +620,9 @@ fn register_mma_rendering_preserves_apis_order_convergence_and_variants() { assert!(targets.contains("GeneratedRegisterMmaOperation::Multiply")); assert!(targets.contains("GeneratedRegisterMmaOperation::AndPopc")); assert!(targets.contains("GeneratedRegisterMmaOperation::XorPopc")); - assert!(targets.contains("pub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4 }")); + assert!( + targets.contains("pub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4, Mxf4 }") + ); assert!(targets.contains("kind: GeneratedRegisterMmaKind::Standard")); assert!(targets.contains("kind: GeneratedRegisterMmaKind::F8f6f4")); assert!(targets.contains("kind: GeneratedRegisterMmaKind::Mxf8f6f4")); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs index dec130798d..5dac3bd1aa 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs @@ -420,11 +420,11 @@ fn tma_rendering_preserves_api_and_injects_backend_defaults() { "pub unsafe fn cp_async_bulk_tensor_reduce_xor_im2col_5d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" )); assert!(compatibility.contains( - "pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(dst: SharedPtr, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier)" + "pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier)" )); for dimensions in 1..=5 { assert!(compatibility.contains(&format!( - "pub unsafe fn cp_async_bulk_tensor_{dimensions}d_g2s_cta(dst: SharedPtr," + "pub unsafe fn cp_async_bulk_tensor_{dimensions}d_g2s_cta(dst: *mut u8," ))); assert!(compatibility.contains(&format!( "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2(" diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs index 7604c3dfc1..ccd4071399 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/warp_memory.rs @@ -817,13 +817,16 @@ fn extended_mbarrier_rendering_preserves_all_manual_contracts() { let catalog = crate::resolve::resolve(&repo_root).unwrap(); validate_renderable(&catalog).unwrap(); let records = mbarrier_extended(&catalog).collect::>(); - assert_eq!(records.len(), 11); + assert_eq!(records.len(), 14); let compatibility = render_compat_mbarrier_extended(&catalog, "test-hash"); for signature in [ "pub unsafe fn mbarrier_arrive_expect_tx(bar: *const Barrier, _tx_count: u32, bytes: u32) -> u64", "pub unsafe fn mbarrier_arrive_expect_tx_cluster(bar: *const Barrier, _tx_count: u32, bytes: u32) -> u64", "pub unsafe fn mbarrier_arrive_cluster(remote_bar_addr: u64)", + "pub unsafe fn mbarrier_arrive_cluster_relaxed(remote_bar_addr: u64)", + "pub unsafe fn mbarrier_try_wait_parity_cluster_relaxed(bar: *const Barrier, parity: u32) -> bool", + "pub unsafe fn fence_acquire_shared_cluster_cluster()", "pub unsafe fn mbarrier_try_wait(bar: *const Barrier, token: u64) -> bool", "pub unsafe fn mbarrier_try_wait_parity(bar: *const Barrier, parity: u32) -> bool", "pub unsafe fn mbarrier_try_wait_parity_cluster(bar: *const Barrier, parity: u32) -> bool", @@ -886,6 +889,7 @@ fn extended_mbarrier_rendering_preserves_all_manual_contracts() { assert!(importer.contains("MbarrierTryWaitSharedOp::build(ctx, barrier, token)")); assert!(importer.contains("MbarrierTryWaitParitySharedOp::build(ctx, barrier, parity)")); assert!(importer.contains("MbarrierTryWaitParityClusterOp::build(ctx, barrier, parity)")); + assert!(importer.contains("FenceAcquireSharedClusterClusterOp::build(ctx)")); assert!(importer.contains("FenceProxyAsyncSharedCtaOp::build(ctx)")); assert!(importer.contains("FenceMbarrierInitReleaseClusterOp::build(ctx)")); assert!(importer.contains("NanosleepOp::build(ctx, ns)")); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs index defa726c97..62a5fd3e9a 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs @@ -1110,6 +1110,14 @@ pub(super) fn validate_renderable(catalog: &CatalogFile) -> Result<()> { MbarrierExtendedOperation::ArriveRemoteCluster, MbarrierExtendedAdapter::RawClusterAddressToVoid, ) => record.id == "mbarrier_arrive_cluster", + ( + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed, + MbarrierExtendedAdapter::RawClusterAddressToVoid, + ) => record.id == "mbarrier_arrive_cluster_relaxed", + ( + MbarrierExtendedOperation::TryWaitParityClusterRelaxed, + MbarrierExtendedAdapter::PointerParityToPredicate, + ) => record.id == "mbarrier_try_wait_parity_cluster_relaxed", ( MbarrierExtendedOperation::TryWaitTokenCta, MbarrierExtendedAdapter::PointerTokenToPredicate, @@ -1140,6 +1148,10 @@ pub(super) fn validate_renderable(catalog: &CatalogFile) -> Result<()> { MbarrierExtendedAdapter::ZeroOperandsToVoid, ) => record.id == "fence_proxy_async_generic_acquire_shared_cluster_cluster", + ( + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster, + MbarrierExtendedAdapter::ZeroOperandsToVoid, + ) => record.id == "fence_acquire_shared_cluster_cluster", ( MbarrierExtendedOperation::Nanosleep, MbarrierExtendedAdapter::NanosecondsToVoid, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs index 50800ed4e4..ebdab99e4d 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mbarrier_extended.rs @@ -203,6 +203,50 @@ pub(in crate::resolve) fn mbarrier_extended_recipe( ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive", summary: "Arrives at a remote cluster-shared barrier through its raw address.", }, + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed => MbarrierExtendedRecipe { + operation, + abi_id: "i1040", + id: "mbarrier_arrive_cluster_relaxed", + operation_key: "barrier.mbarrier.arrive.shared.cluster.relaxed.cluster.raw_address", + source_record: None, + llvm_symbol: None, + ptx_native_instruction: Some("mbarrier.arrive.relaxed.cluster.shared::cluster.b64"), + rust_arguments: &["u64"], + rust_result: "()", + must_use: false, + dialect_op_type: "MbarrierArriveClusterRelaxedOp", + dialect_op_name: "nvvm.mbarrier_arrive_cluster_relaxed", + dialect_operands: &["i64"], + dialect_results: &[], + llvm_arguments: &[], + llvm_results: &[], + llvm_properties: &[], + adapter: MbarrierExtendedAdapter::RawClusterAddressToVoid, + source_contract: MbarrierExtendedSourceContract::PtxNativeRawClusterAddress, + execution_scope: "cluster", + minimum_ptx: "8.6", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction( + &[ + "mbarrier", + "arrive", + "relaxed", + "cluster", + "shared::cluster", + "b64", + ], + vec![ + OperandPattern::Exact { value: "_".into() }, + OperandPattern::Address, + ], + ), + inline_ptx: "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [$0];", + inline_constraints: "l,~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: mbarrier.arrive", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-arrive", + summary: "Arrives at a remote cluster-shared barrier with relaxed memory ordering.", + }, MbarrierExtendedOperation::TryWaitTokenCta => MbarrierExtendedRecipe { operation, abi_id: "i0309", @@ -325,6 +369,60 @@ pub(in crate::resolve) fn mbarrier_extended_recipe( ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait", summary: "Tests barrier parity with cluster-scope acquire ordering.", }, + MbarrierExtendedOperation::TryWaitParityClusterRelaxed => MbarrierExtendedRecipe { + operation, + abi_id: "i1041", + id: "mbarrier_try_wait_parity_cluster_relaxed", + operation_key: "barrier.mbarrier.try_wait.parity.shared.cta.relaxed.cluster", + source_record: Some( + "int_nvvm_mbarrier_try_wait_parity_relaxed_scope_cluster_space_cta", + ), + llvm_symbol: Some("llvm.nvvm.mbarrier.try.wait.parity.relaxed.scope.cluster.space.cta"), + ptx_native_instruction: None, + rust_arguments: &["*const u64", "u32"], + rust_result: "bool", + must_use: true, + dialect_op_type: "MbarrierTryWaitParityClusterRelaxedOp", + dialect_op_name: "nvvm.mbarrier_try_wait_parity_cluster_relaxed", + dialect_operands: &["ptr", "i32"], + dialect_results: &["i1"], + llvm_arguments: &["shared_ptr", "i32"], + llvm_results: &["i1"], + llvm_properties: &[ + "IntrArgMemOnly", + "IntrConvergent", + "IntrNoCallback", + "IntrReadMem", + "NoCapture", + ], + adapter: MbarrierExtendedAdapter::PointerParityToPredicate, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cluster", + minimum_ptx: "8.6", + minimum_sm: "sm_90", + ptx_result: "bool", + expected_ptx: instruction( + &[ + "mbarrier", + "try_wait", + "parity", + "relaxed", + "cluster", + "shared::cta", + "b64", + ], + vec![ + OperandPattern::Register, + OperandPattern::Address, + OperandPattern::Register, + ], + ), + inline_ptx: "{ .reg .pred %p0; mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", + inline_constraints: "=r,l,r,~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: mbarrier.test_wait / mbarrier.try_wait", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait", + summary: "Tests barrier parity at cluster scope with relaxed memory ordering.", + }, MbarrierExtendedOperation::FenceProxyAsyncSharedCta => MbarrierExtendedRecipe { operation, abi_id: "i0312", @@ -481,6 +579,45 @@ pub(in crate::resolve) fn mbarrier_extended_recipe( summary: "Acquires cluster-shared async-proxy writes through the generic proxy.", } } + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster => MbarrierExtendedRecipe { + operation, + abi_id: "i1042", + id: "fence_acquire_shared_cluster_cluster", + operation_key: "fence.acquire.sync_restrict.shared_cluster.cluster", + source_record: Some("int_nvvm_fence_acquire_sync_restrict_space_cluster_scope_cluster"), + llvm_symbol: Some("llvm.nvvm.fence.acquire.sync_restrict.space.cluster.scope.cluster"), + ptx_native_instruction: None, + rust_arguments: &[], + rust_result: "()", + must_use: false, + dialect_op_type: "FenceAcquireSharedClusterClusterOp", + dialect_op_name: "nvvm.fence_acquire_shared_cluster_cluster", + dialect_operands: &[], + dialect_results: &[], + llvm_arguments: &[], + llvm_results: &[], + llvm_properties: &["IntrNoCallback"], + adapter: MbarrierExtendedAdapter::ZeroOperandsToVoid, + source_contract: MbarrierExtendedSourceContract::LlvmImported, + execution_scope: "cluster", + minimum_ptx: "8.6", + minimum_sm: "sm_90", + ptx_result: "()", + expected_ptx: instruction( + &[ + "fence", + "acquire", + "sync_restrict::shared::cluster", + "cluster", + ], + vec![], + ), + inline_ptx: "fence.acquire.sync_restrict::shared::cluster.cluster;", + inline_constraints: "~{memory}", + ptx_isa_section: "Parallel Synchronization and Communication Instructions: membar / fence", + ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar", + summary: "Acquires cluster-shared memory at cluster scope, without ordering global memory.", + }, MbarrierExtendedOperation::Nanosleep => MbarrierExtendedRecipe { operation, abi_id: "i0316", @@ -560,6 +697,9 @@ pub(in crate::resolve) fn expand_mbarrier_extended_admission( MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, MbarrierExtendedOperation::Nanosleep, + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed, + MbarrierExtendedOperation::TryWaitParityClusterRelaxed, + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster, ]); let actual_operations = admission .variants @@ -658,7 +798,12 @@ pub(in crate::resolve) fn expand_mbarrier_extended_admission( admission.llvm_evidence_profile.clone() } IntrinsicBackend::LibNvvm => { - admission.libnvvm_evidence_profile.clone() + if matches!(recipe.operation, MbarrierExtendedOperation::ArriveRemoteClusterRelaxed | MbarrierExtendedOperation::TryWaitParityClusterRelaxed | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster) { + admission.relaxed_libnvvm_evidence_profile.as_ref() + .unwrap_or(&admission.libnvvm_evidence_profile).clone() + } else { + admission.libnvvm_evidence_profile.clone() + } } }, targets: None, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs index b76e4112bd..f74d0e4c76 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/validate.rs @@ -31,8 +31,8 @@ pub(in crate::resolve) fn validate_register_mma_policy( Some(RegisterMmaKind::F8f6f4) => { return validate_register_mma_f8f6f4_policy(policy, declaration, mma); } - Some(RegisterMmaKind::Mxf8f6f4) => { - return validate_register_mma_mxf8f6f4_policy(policy, declaration, mma); + Some(RegisterMmaKind::Mxf8f6f4 | RegisterMmaKind::Mxf4) => { + return validate_register_mma_block_scaled_policy(policy, declaration, mma); } None if register_mma_f8f6f4_element_name(mma.a_element).is_some() || register_mma_f8f6f4_element_name(mma.b_element).is_some() => @@ -559,20 +559,77 @@ pub(in crate::resolve) fn validate_register_mma_f8f6f4_policy( Ok(()) } -pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( +pub(in crate::resolve) fn validate_register_mma_block_scaled_policy( policy: &OverlayIntrinsic, declaration: &ImportedIntrinsic, mma: &RegisterMma, ) -> Result<()> { let a = register_mma_f8f6f4_element_name(mma.a_element) - .with_context(|| format!("{} has a non-mxf8f6f4 A format", policy.id))?; + .with_context(|| format!("{} has a non-block-scaled A format", policy.id))?; let b = register_mma_f8f6f4_element_name(mma.b_element) - .with_context(|| format!("{} has a non-mxf8f6f4 B format", policy.id))?; - let expected_id = format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"); - let expected_source = - format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"); - let expected_symbol = - format!("llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0"); + .with_context(|| format!("{} has a non-block-scaled B format", policy.id))?; + let is_mxf4 = mma.kind == Some(RegisterMmaKind::Mxf4); + let (expected_id, expected_source, expected_symbol, operation_key, shape, targets, modifiers) = + if is_mxf4 { + ensure!( + a == "e2m1" && b == "e2m1", + "{} MXF4 requires E2M1 operands", + policy.id + ); + ( + "mma_m16n8k64_mxf4_f32".to_owned(), + "int_nvvm_mma_block_scale_m16n8k64_row_col_mxf4_scale_2x_f32_e2m1_e2m1_f32_ue8m0" + .to_owned(), + "llvm.nvvm.mma.block.scale.m16n8k64.row.col.mxf4.scale.2x.f32.e2m1.e2m1.f32.ue8m0" + .to_owned(), + "matrix.mma.m16n8k64.row.col.kind_mxf4.scale_vec_2x.f32.e2m1.e2m1.f32.ue8m0" + .to_owned(), + RegisterMmaShape::M16n8k64, + "sm_120a", + vec![ + "sync", + "aligned", + "m16n8k64", + "row", + "col", + "kind::mxf4", + "block_scale", + "scale_vec::2X", + "f32", + "e2m1", + "e2m1", + "f32", + "ue8m0", + ], + ) + } else { + ( + format!("mma_m16n8k32_mxf8f6f4_f32_{a}_{b}"), + format!("int_nvvm_mma_block_scale_m16n8k32_row_col_mxf8f6f4_f32_{a}_{b}_f32_ue8m0"), + format!( + "llvm.nvvm.mma.block.scale.m16n8k32.row.col.mxf8f6f4.f32.{a}.{b}.f32.ue8m0" + ), + format!( + "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" + ), + RegisterMmaShape::M16n8k32, + REGISTER_MMA_F8F6F4_TARGETS, + vec![ + "sync", + "aligned", + "m16n8k32", + "row", + "col", + "kind::mxf8f6f4", + "block_scale", + "f32", + a, + b, + "f32", + "ue8m0", + ], + ) + }; let rust_arguments = [ "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16", ]; @@ -587,15 +644,12 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( let results = ["f32", "f32", "f32", "f32"]; ensure!( policy.id == expected_id - && policy.operation_key - == format!( - "matrix.mma.m16n8k32.row.col.kind_mxf8f6f4.scale_vec_1x.f32.{a}.{b}.f32.ue8m0" - ) + && policy.operation_key == operation_key && policy.source.is_none() && policy.source_record.as_deref() == Some(expected_source.as_str()) && policy.llvm_symbol.as_deref() == Some(expected_symbol.as_str()) && policy.resolved_llvm_symbol.is_none(), - "{} dense mxf8f6f4 MMA identity changed", + "{} dense block-scaled MMA identity changed", policy.id ); ensure!( @@ -608,7 +662,7 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && policy.safe_allowlist_reason.is_none() && policy.public_rust_path == format!("cuda_intrinsics::matrix::{expected_id}") && policy.compatibility_rust_paths == [format!("cuda_device::wmma::{expected_id}")], - "{} must preserve the unsafe must-use dense mxf8f6f4 API", + "{} must preserve the unsafe must-use dense block-scaled API", policy.id ); ensure!( @@ -619,8 +673,11 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && policy.llvm_arguments == llvm_arguments && policy.llvm_results == results && policy.ptx_result == "[f32; 4]" - && mma.shape == RegisterMmaShape::M16n8k32 - && mma.kind == Some(RegisterMmaKind::Mxf8f6f4) + && mma.shape == shape + && matches!( + mma.kind, + Some(RegisterMmaKind::Mxf8f6f4 | RegisterMmaKind::Mxf4) + ) && mma.operation == RegisterMmaOperation::Multiply && mma.accumulator == RegisterMmaAccumulator::F32 && mma.a_layout == RegisterMmaLayout::Row @@ -632,7 +689,7 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && mma.compatibility_source == RegisterMmaCompatibilitySource::GeneratedStub && mma.runtime_validation == RuntimeValidation::Unexecuted && policy.lowering == "generated_register_mma", - "{} dense mxf8f6f4 carrier or lowering changed", + "{} dense block-scaled carrier or lowering changed", policy.id ); ensure!( @@ -642,8 +699,8 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && policy.execution_scope == "warp" && policy.minimum_ptx == "8.7" && policy.minimum_sm.is_none() - && policy.targets == REGISTER_MMA_F8F6F4_TARGETS, - "{} dense mxf8f6f4 effects or exact target set changed", + && policy.targets == targets, + "{} dense block-scaled effects or exact target set changed", policy.id ); ensure!( @@ -651,12 +708,12 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && policy.ptx_isa_section == "9.7.15.5.14 Multiply-and-Accumulate Instruction: mma" && policy.ptx_isa_url == "https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma", - "{} dense mxf8f6f4 PTX provenance changed", + "{} dense block-scaled PTX provenance changed", policy.id ); let [selection] = declaration.selections.as_slice() else { bail!( - "{} must retain exactly one imported dense mxf8f6f4 instruction selection", + "{} must retain exactly one imported dense block-scaled instruction selection", policy.id ); }; @@ -666,30 +723,14 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( && selection_matches_policy(policy, selection)? && selection.predicates == ["Subtarget->hasMMABlockScale()"] && selection.constraints.is_empty(), - "{} imported dense mxf8f6f4 declaration or selection changed", + "{} imported dense block-scaled declaration or selection changed", policy.id ); ensure!( policy.expected_ptx == InstructionPattern { mnemonic: "mma".into(), - modifiers: [ - "sync", - "aligned", - "m16n8k32", - "row", - "col", - "kind::mxf8f6f4", - "block_scale", - "f32", - a, - b, - "f32", - "ue8m0", - ] - .into_iter() - .map(Into::into) - .collect(), + modifiers: modifiers.into_iter().map(Into::into).collect(), operands: vec![ OperandPattern::RegisterList { length: 4 }, OperandPattern::RegisterList { length: 4 }, @@ -701,7 +742,7 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( OperandPattern::RegisterList { length: 2 }, ], }, - "{} expected dense mxf8f6f4 PTX changed", + "{} expected dense block-scaled PTX changed", policy.id ); let backend_pairs: BTreeSet<_> = policy @@ -709,28 +750,35 @@ pub(in crate::resolve) fn validate_register_mma_mxf8f6f4_policy( .iter() .map(|lowering| (lowering.backend, lowering.mechanism)) .collect(); + let expected_backend_pairs = if is_mxf4 { + BTreeSet::from([( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + )]) + } else { + BTreeSet::from([ + ( + IntrinsicBackend::LlvmNvptx, + BackendLoweringMechanism::InlinePtx, + ), + ( + IntrinsicBackend::LibNvvm, + BackendLoweringMechanism::InlinePtx, + ), + ]) + }; ensure!( - policy.backend_lowerings.len() == 2 - && backend_pairs - == BTreeSet::from([ - ( - IntrinsicBackend::LlvmNvptx, - BackendLoweringMechanism::InlinePtx, - ), - ( - IntrinsicBackend::LibNvvm, - BackendLoweringMechanism::InlinePtx, - ), - ]) + policy.backend_lowerings.len() == expected_backend_pairs.len() + && backend_pairs == expected_backend_pairs && policy.backend_lowerings.iter().all(|lowering| { lowering.targets.is_none() && lowering.minimum_ptx.is_none() && lowering.minimum_sm.is_none() && !lowering.evidence_profile.trim().is_empty() }), - "{} must inherit the exact reviewed target set on both inline-PTX routes", + "{} must inherit the exact reviewed target set on its inline-PTX routes", policy.id ); - ensure_no_other_family_contract(policy, "dense mxf8f6f4 register MMA")?; + ensure_no_other_family_contract(policy, "dense block-scaled register MMA")?; Ok(()) } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/policy.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/policy.rs index 7313e7ca91..e17efc8fc8 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/policy.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/policy.rs @@ -572,6 +572,7 @@ pub(super) fn validate_policy( | MbarrierExtendedOperation::ArriveExpectTxCluster | MbarrierExtendedOperation::TryWaitParityCta | MbarrierExtendedOperation::TryWaitParityCluster + | MbarrierExtendedOperation::TryWaitParityClusterRelaxed ) }) => { diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs index 777126a5df..750bbcf4a4 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/targets.rs @@ -7,7 +7,8 @@ use crate::model::{ CatalogHardwareAlternative, CatalogHardwareTarget, CatalogTargetAlternative, CatalogTargetContract, CatalogTargetRequirement, ClcOperation, ExecutionControlOperation, IntrinsicBackend, LdmatrixElement, LdmatrixShape, OverlayIntrinsic, PtxVersion, - SparseMmaAccumulator, TargetContract, TargetSelectorBinding, Tcgen05Operation, TmaOperation, + RegisterMmaKind, SparseMmaAccumulator, TargetContract, TargetSelectorBinding, Tcgen05Operation, + TmaOperation, }; use anyhow::{Context, Result, bail, ensure}; @@ -352,7 +353,17 @@ pub(super) fn validate_selected_target_predicates( let (minimum_ptx_matches, target_matches) = match policy.family.as_str() { "register_mma" => ( overlay_ptx == 87, - policy.targets == "sm_120a|sm_120f|sm_121a|sm_121f" && policy.minimum_sm.is_none(), + policy.targets + == if policy + .register_mma + .as_ref() + .is_some_and(|mma| mma.kind == Some(RegisterMmaKind::Mxf4)) + { + "sm_120a" + } else { + "sm_120a|sm_120f|sm_121a|sm_121f" + } + && policy.minimum_sm.is_none(), ), "sparse_mma" => ( overlay_ptx == 87, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs index d4710b6630..457b16f77b 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/fixtures.rs @@ -1885,6 +1885,9 @@ pub(super) fn test_mbarrier_extended_admission() -> MbarrierExtendedAdmission { MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster, MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster, MbarrierExtendedOperation::Nanosleep, + MbarrierExtendedOperation::ArriveRemoteClusterRelaxed, + MbarrierExtendedOperation::TryWaitParityClusterRelaxed, + MbarrierExtendedOperation::FenceAcquireSharedClusterCluster, ] .map(|operation| crate::model::MbarrierExtendedAdmissionVariant { abi_id: mbarrier_extended_recipe(operation).abi_id.into(), @@ -1894,6 +1897,7 @@ pub(super) fn test_mbarrier_extended_admission() -> MbarrierExtendedAdmission { MbarrierExtendedAdmission { llvm_evidence_profile: "llvm-test".into(), libnvvm_evidence_profile: "libnvvm-test".into(), + relaxed_libnvvm_evidence_profile: None, runtime_validation: RuntimeValidation::Unexecuted, variants, } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs index 1fce9ebfdd..c9ac59f5c8 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/misc_families.rs @@ -252,16 +252,83 @@ fn active_debug_control_sources_parse_and_prove_both_backend_routes() { } } +#[test] +fn relaxed_cluster_arrival_and_wait_preserve_distinct_memory_contracts() { + let arrival: MbarrierExtendedOperation = + serde_json::from_str("\"arrive_remote_cluster_relaxed\"") + .expect("relaxed remote arrival is admitted"); + let wait: MbarrierExtendedOperation = + serde_json::from_str("\"try_wait_parity_cluster_relaxed\"") + .expect("relaxed cluster wait is admitted"); + let arrival = mbarrier_extended_recipe(arrival); + let wait = mbarrier_extended_recipe(wait); + assert_eq!(arrival.abi_id, "i1040"); + assert_eq!(wait.abi_id, "i1041"); + assert_eq!(arrival.id, "mbarrier_arrive_cluster_relaxed"); + assert_eq!(wait.id, "mbarrier_try_wait_parity_cluster_relaxed"); + assert_eq!(arrival.rust_arguments, ["u64"]); + assert_eq!(arrival.rust_result, "()"); + assert_eq!(arrival.execution_scope, "cluster"); + assert_eq!( + arrival.inline_ptx, + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [$0];" + ); + assert_eq!(wait.execution_scope, "cluster"); + assert!( + wait.inline_ptx + .contains("mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64") + ); + assert_eq!(arrival.minimum_ptx, "8.6"); + assert_eq!(wait.minimum_ptx, "8.6"); + let existing = mbarrier_extended_recipe(MbarrierExtendedOperation::ArriveRemoteCluster); + assert_eq!(existing.abi_id, "i0308"); + assert_eq!( + existing.inline_ptx, + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [$0];" + ); + let existing = mbarrier_extended_recipe(MbarrierExtendedOperation::TryWaitParityCluster); + assert_eq!(existing.abi_id, "i0311"); + assert!( + existing + .inline_ptx + .contains("try_wait.parity.acquire.cluster") + ); +} + +#[test] +fn restricted_acquire_fence_orders_cluster_shared_memory_only() { + let operation: MbarrierExtendedOperation = + serde_json::from_str("\"fence_acquire_shared_cluster_cluster\"") + .expect("restricted shared-memory acquire is admitted"); + let fence = mbarrier_extended_recipe(operation); + assert_eq!(fence.abi_id, "i1042"); + assert_eq!(fence.id, "fence_acquire_shared_cluster_cluster"); + assert_eq!(fence.rust_arguments, [] as [&str; 0]); + assert_eq!(fence.rust_result, "()"); + assert_eq!(fence.execution_scope, "cluster"); + assert_eq!(fence.minimum_ptx, "8.6"); + assert_eq!(fence.minimum_sm, "sm_90"); + assert_eq!( + fence.llvm_symbol, + Some("llvm.nvvm.fence.acquire.sync_restrict.space.cluster.scope.cluster") + ); + assert_eq!( + fence.inline_ptx, + "fence.acquire.sync_restrict::shared::cluster.cluster;" + ); +} + #[test] fn compact_extended_mbarrier_admission_preserves_all_manual_contracts() { let records = expand_mbarrier_extended_admission(&test_mbarrier_extended_admission()).unwrap(); - assert_eq!(records.len(), 11); + assert_eq!(records.len(), 14); assert_eq!( records .iter() .map(|record| record.abi_id.as_str()) .collect::>(), (306..=316) + .chain([1040, 1041, 1042]) .map(|id| format!("i{id:04}")) .collect::>() ); diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs index 14f5cd2c1e..9d08c0e7a3 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/mma.rs @@ -25,6 +25,68 @@ use crate::resolve::families::*; use crate::resolve::guards::*; use crate::resolve::overlay::*; +#[test] +fn dense_mxf4_k64_preserves_selector_abi_and_rejects_contract_drift() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let (overlay, _) = + read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); + let policy = overlay + .intrinsics + .iter() + .find(|record| record.id == "mma_m16n8k64_mxf4_f32") + .expect("canonical dense MXF4 K64 admission"); + let imported: ImportedFile = read_json(&repo_root.join("intrinsics/imported.json")).unwrap(); + let declaration = imported + .intrinsics + .iter() + .find(|record| Some(&record.source_record) == policy.source_record.as_ref()) + .unwrap(); + validate_register_mma_policy(policy, declaration).unwrap(); + resolve_candidate( + &repo_root, + &policy.id, + "test-llvm", + &"a".repeat(64), + "sm_120a", + "+ptx87", + ) + .unwrap(); + assert_eq!(policy.abi_id, "i1039"); + assert_eq!( + policy.rust_arguments, + [ + "[f32; 4]", "[u32; 4]", "[u32; 2]", "u32", "u16", "u16", "u32", "u16", "u16" + ] + ); + assert_eq!( + policy.expected_ptx.modifiers, + [ + "sync", + "aligned", + "m16n8k64", + "row", + "col", + "kind::mxf4", + "block_scale", + "scale_vec::2X", + "f32", + "e2m1", + "e2m1", + "f32", + "ue8m0" + ] + ); + let mut wrong_shape = policy.clone(); + wrong_shape.register_mma.as_mut().unwrap().shape = RegisterMmaShape::M16n8k32; + assert!(validate_register_mma_policy(&wrong_shape, declaration).is_err()); + let mut wrong_selector = policy.clone(); + wrong_selector.rust_arguments[4] = "u32".into(); + assert!(validate_register_mma_policy(&wrong_selector, declaration).is_err()); + let mut wrong_scale = policy.clone(); + wrong_scale.expected_ptx.modifiers[7] = "scale_vec::4X".into(); + assert!(validate_register_mma_policy(&wrong_scale, declaration).is_err()); +} + #[test] fn stmatrix_admission_is_closed_and_uses_schema_35() { let shard = |schema| { @@ -1134,7 +1196,7 @@ fn pinned_register_mma_records_match_the_closed_recipes_and_fail_closed() { .iter() .filter(|record| record.family == "register_mma") .collect(); - assert_eq!(records.len(), 154); + assert_eq!(records.len(), 155); let dense_f8f6f4_records = records .iter() diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs index 8be4c95782..a402496aba 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/overlay_schema.rs @@ -48,8 +48,8 @@ fn overlay_manifest_loads_sorted_family_shards() { let (overlay, hash) = read_overlay(&repo_root, &repo_root.join("intrinsics/overlay.toml")).unwrap(); assert_eq!(overlay.schema, OVERLAY_SCHEMA); - assert_eq!(overlay.shards.len(), 67); - assert_eq!(overlay.intrinsics.len(), 1038); + assert_eq!(overlay.shards.len(), 68); + assert_eq!(overlay.intrinsics.len(), 1042); assert_eq!( overlay .intrinsics @@ -136,7 +136,7 @@ fn overlay_manifest_loads_sorted_family_shards() { .iter() .filter(|record| record.family == "register_mma") .count(), - 154 + 155 ); assert_eq!( overlay diff --git a/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs b/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs index 30caf5b4fe..a208567ecb 100644 --- a/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs +++ b/cuda-oxide/crates/cuda-intrinsics/src/generated/abi_v1.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Raw ABI functions recognized by cuda-oxide. @@ -3696,6 +3696,21 @@ pub fn i0820(_arg0: f32) -> f32 { ) } +/// Acquires cluster-shared memory at cluster scope, without ordering global memory. +/// +/// Catalog ID: `fence_acquire_shared_cluster_cluster`. Source: LLVM `llvm.nvvm.fence.acquire.sync_restrict.space.cluster.scope.cluster` from `int_nvvm_fence_acquire_sync_restrict_space_cluster_scope_cluster`; expects PTX `fence.acquire.sync_restrict::shared::cluster.cluster;`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// This is a convergent, side-effecting operation with a compiler memory clobber. +/// The call site must satisfy the scope and proxy-ordering contract named by the operation. +#[inline(never)] +pub unsafe fn i1042() -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042` executed outside device compilation" + ) +} + /// Releases mbarrier initialization at cluster scope. /// /// Catalog ID: `fence_mbarrier_init_release_cluster`. Source: LLVM `llvm.nvvm.fence.mbarrier_init.release.cluster` from `int_nvvm_fence_mbarrier_init_release_cluster`; expects PTX `fence.mbarrier_init.release.cluster;`. @@ -5252,6 +5267,21 @@ pub unsafe fn i0308(_arg0: u64) -> () { ) } +/// Arrives at a remote cluster-shared barrier with relaxed memory ordering. +/// +/// Catalog ID: `mbarrier_arrive_cluster_relaxed`. Source: PTX-native `mbarrier.arrive.relaxed.cluster.shared::cluster.b64`; expects PTX `mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// This is a convergent, side-effecting operation with a compiler memory clobber. +/// `_arg0` must be a live remote shared-memory barrier address obtained through cluster address mapping. +#[inline(never)] +pub unsafe fn i1040(_arg0: u64) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040` executed outside device compilation" + ) +} + /// Arrives at a CTA-shared barrier and adds expected transaction bytes. /// /// Catalog ID: `mbarrier_arrive_expect_tx`. Source: LLVM `llvm.nvvm.mbarrier.arrive.expect.tx.scope.cta.space.cta` from `int_nvvm_mbarrier_arrive_expect_tx_scope_cta_space_cta`; expects PTX `mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 , [
], ;`. @@ -5399,6 +5429,22 @@ pub unsafe fn i0311(_arg0: *const u64, _arg1: u32) -> bool { ) } +/// Tests barrier parity at cluster scope with relaxed memory ordering. +/// +/// Catalog ID: `mbarrier_try_wait_parity_cluster_relaxed`. Source: LLVM `llvm.nvvm.mbarrier.try.wait.parity.relaxed.scope.cluster.space.cta` from `int_nvvm_mbarrier_try_wait_parity_relaxed_scope_cluster_space_cta`; expects PTX `mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 , [
], ;`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// This is a convergent, side-effecting operation with a compiler memory clobber. +/// `_arg0` must point to a live initialized barrier. `_arg1` must be the expected phase parity. +#[must_use] +#[inline(never)] +pub unsafe fn i1041(_arg0: *const u64, _arg1: u32) -> bool { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041` executed outside device compilation" + ) +} + /// Computes extended min for bf16 values. /// /// Catalog ID: `min_bf16`. Source: LLVM `llvm.nvvm.fmin.bf16` from `int_nvvm_fmin_bf16`; expects PTX `min.bf16 , , ;`. @@ -8381,6 +8427,37 @@ pub unsafe fn i0520(_arg0: [f32; 4], _arg1: [u32; 2], _arg2: u32) -> [f32; 4] { ) } +/// Multiplies warp-distributed E2M1 fragments with UE8M0 block scales and adds an F32 accumulator. +/// +/// Catalog ID: `mma_m16n8k64_mxf4_f32`. Source: LLVM `llvm.nvvm.mma.block.scale.m16n8k64.row.col.mxf4.scale.2x.f32.e2m1.e2m1.f32.ue8m0` from `int_nvvm_mma_block_scale_m16n8k64_row_col_mxf4_scale_2x_f32_e2m1_e2m1_f32_ue8m0`; expects PTX `mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 , , , , , , , ;`. +/// Available on `sm_120a` targets from PTX 8.7. +/// +/// # Safety +/// All 32 warp lanes must execute the same `mma.sync.aligned` instruction with the same qualifiers, and no lane may have exited. +/// `_arg0`, `_arg1`, and `_arg2` must contain this lane's C, A, and B fragments in the documented PTX layout. +/// The operation is register-only and is not a memory fence. +/// `_arg3` and `_arg6` contain this lane's packed A and B scale data. `_arg4`/`_arg5` and `_arg7`/`_arg8` are the corresponding byte/thread selectors. +/// For `scale_vec::2X`, A and B byte selectors must each be `0` or `2`, the A thread selector in `0..=1`, and the B thread selector in `0..=3`; other values make the PTX operation undefined. +/// See the [PTX MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma). +#[allow(clippy::too_many_arguments)] +#[must_use] +#[inline(never)] +pub unsafe fn i1039( + _arg0: [f32; 4], + _arg1: [u32; 4], + _arg2: [u32; 2], + _arg3: u32, + _arg4: u16, + _arg5: u16, + _arg6: u32, + _arg7: u16, + _arg8: u16, +) -> [f32; 4] { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039` executed outside device compilation" + ) +} + /// Multiplies warp-distributed signed A and signed B INT4 fragments and adds a wrapping s32 accumulator. /// /// Catalog ID: `mma_m16n8k64_s32_s4`. Source: LLVM `llvm.nvvm.mma.m16n8k64.row.col.s4` from `int_nvvm_mma_m16n8k64_row_col_s4`; expects PTX `mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 , , , ;`. diff --git a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs index 0dbdbcf1bc..5a6623e110 100644 --- a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs +++ b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. #[doc(hidden)] @@ -58,6 +58,9 @@ pub mod barrier { pub use crate::__cuda_oxide_intrinsic_abi_v1::i0885 as barrier_cta_arrive; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0886 as barrier_cta_arrive_aligned; pub use crate::__cuda_oxide_intrinsic_abi_v1::i1017 as mbarrier_arrive_no_complete; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1040 as mbarrier_arrive_cluster_relaxed; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1041 as mbarrier_try_wait_parity_cluster_relaxed; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1042 as fence_acquire_shared_cluster_cluster; } /// Generated `bf16` intrinsic source API. @@ -683,6 +686,7 @@ pub mod matrix { pub use crate::__cuda_oxide_intrinsic_abi_v1::i1027 as mma_sp_m16n8k64_f32_e4m3_e5m2_f32; pub use crate::__cuda_oxide_intrinsic_abi_v1::i1028 as mma_sp_m16n8k64_f32_e5m2_e4m3_f32; pub use crate::__cuda_oxide_intrinsic_abi_v1::i1029 as mma_sp_m16n8k64_f32_e5m2_e5m2_f32; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1039 as mma_m16n8k64_mxf4_f32; } /// Generated `prmt` intrinsic source API. @@ -2791,6 +2795,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_fence_acquire_shared_cluster_cluster_reexports_abi_i1042() { + let public: unsafe fn() -> () = super::barrier::fence_acquire_shared_cluster_cluster; + let canonical: unsafe fn() -> () = super::__cuda_oxide_intrinsic_abi_v1::i1042; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_fence_mbarrier_init_release_cluster_reexports_abi_i0313() { let public: unsafe fn() -> () = super::barrier::fence_mbarrier_init_release_cluster; let canonical: unsafe fn() -> () = super::__cuda_oxide_intrinsic_abi_v1::i0313; @@ -3504,6 +3514,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_mbarrier_arrive_cluster_relaxed_reexports_abi_i1040() { + let public: unsafe fn(u64) -> () = super::barrier::mbarrier_arrive_cluster_relaxed; + let canonical: unsafe fn(u64) -> () = super::__cuda_oxide_intrinsic_abi_v1::i1040; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_mbarrier_arrive_expect_tx_reexports_abi_i0306() { let public: unsafe fn(*const u64, u32, u32) -> u64 = super::barrier::mbarrier_arrive_expect_tx; @@ -3568,6 +3584,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_mbarrier_try_wait_parity_cluster_relaxed_reexports_abi_i1041() { + let public: unsafe fn(*const u64, u32) -> bool = + super::barrier::mbarrier_try_wait_parity_cluster_relaxed; + let canonical: unsafe fn(*const u64, u32) -> bool = + super::__cuda_oxide_intrinsic_abi_v1::i1041; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_min_bf16_reexports_abi_i0838() { let public: fn(u16, u16) -> u16 = super::bf16::min_bf16; let canonical: fn(u16, u16) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i0838; @@ -5194,6 +5218,32 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_mma_m16n8k64_mxf4_f32_reexports_abi_i1039() { + let public: unsafe fn( + [f32; 4], + [u32; 4], + [u32; 2], + u32, + u16, + u16, + u32, + u16, + u16, + ) -> [f32; 4] = super::matrix::mma_m16n8k64_mxf4_f32; + let canonical: unsafe fn( + [f32; 4], + [u32; 4], + [u32; 2], + u32, + u16, + u16, + u32, + u16, + u16, + ) -> [f32; 4] = super::__cuda_oxide_intrinsic_abi_v1::i1039; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_mma_m16n8k64_s32_s4_reexports_abi_i0145() { let public: unsafe fn([i32; 4], [u32; 4], [u32; 2]) -> [i32; 4] = super::matrix::mma_m16n8k64_s32_s4; diff --git a/cuda-oxide/crates/cuda-macros/Cargo.toml b/cuda-oxide/crates/cuda-macros/Cargo.toml index 16a25c8459..440290c02a 100644 --- a/cuda-oxide/crates/cuda-macros/Cargo.toml +++ b/cuda-oxide/crates/cuda-macros/Cargo.toml @@ -45,6 +45,7 @@ reserved-oxide-symbols = { workspace = true } oxide-tile-export = { workspace = true } [dev-dependencies] +ctm-rs = { workspace = true } trybuild = "1" # Dev-only cycle (cargo allows it): the launch_requires_unsafe compile-fail # test type-checks a real cuda_launch! expansion, which needs the host-side diff --git a/cuda-oxide/crates/cuda-macros/src/cuda_module/launchers.rs b/cuda-oxide/crates/cuda-macros/src/cuda_module/launchers.rs index 279a15cce4..d8603f5130 100644 --- a/cuda-oxide/crates/cuda-macros/src/cuda_module/launchers.rs +++ b/cuda-oxide/crates/cuda-macros/src/cuda_module/launchers.rs @@ -407,7 +407,7 @@ fn generate_cuda_module_legacy_launch_method(kernel: &CudaModuleKernel) -> Token #vis unsafe fn #fn_name #impl_generics ( &self, #stream: &::cuda_core::CudaStream, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result<(), ::cuda_core::DriverError> #where_clause @@ -493,7 +493,7 @@ fn generate_cuda_module_prepared_launch_method(kernel: &CudaModuleKernel) -> Tok #vis unsafe fn #unchecked_name #impl_generics ( &self, #stream: &::cuda_core::CudaStream, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result<(), ::cuda_core::DriverError> #where_clause @@ -556,7 +556,7 @@ fn generate_cuda_module_legacy_async_launch_method(kernel: &CudaModuleKernel) -> #[allow(clippy::multiple_bound_locations, clippy::too_many_arguments)] #vis unsafe fn #fn_name #impl_generics ( &self, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result<::cuda_host::AsyncKernelLaunch<#async_lifetime>, ::cuda_core::DriverError> #where_clause @@ -681,7 +681,7 @@ fn generate_cuda_module_prepared_async_launch_method(kernel: &CudaModuleKernel) #[allow(clippy::multiple_bound_locations, clippy::too_many_arguments)] #vis unsafe fn #unchecked_name #impl_generics ( &self, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result< ::cuda_host::AsyncKernelLaunch<#async_lifetime>, @@ -784,7 +784,7 @@ fn generate_cuda_module_legacy_owned_async_launch_method( #[allow(clippy::multiple_bound_locations, clippy::too_many_arguments)] #vis unsafe fn #fn_name #impl_generics ( &self, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result<::cuda_host::OwnedAsyncKernelLaunch<#resources_ty>, ::cuda_core::DriverError> #where_clause @@ -955,7 +955,7 @@ fn generate_cuda_module_prepared_owned_async_launch_method( #[allow(clippy::multiple_bound_locations, clippy::too_many_arguments)] #vis unsafe fn #unchecked_name #impl_generics ( &self, - #config: ::cuda_core::simt::LaunchConfig, + #config: ::cuda_core::LaunchConfig, #(#params),* ) -> ::core::result::Result< ::cuda_host::OwnedAsyncKernelLaunch<#resources_ty>, diff --git a/cuda-oxide/crates/cuda-macros/src/cuda_module/mod.rs b/cuda-oxide/crates/cuda-macros/src/cuda_module/mod.rs index 58d0902427..7de03804a9 100644 --- a/cuda-oxide/crates/cuda-macros/src/cuda_module/mod.rs +++ b/cuda-oxide/crates/cuda-macros/src/cuda_module/mod.rs @@ -236,7 +236,7 @@ pub(crate) fn expand_cuda_module_inner( /// specialization and no conflicting entry definition. pub unsafe fn load_async( device_id: usize, - ) -> ::core::result::Result { + ) -> ::core::result::Result { // SAFETY: upheld by this function's caller. unsafe { load_async_named(device_id, env!("CARGO_PKG_NAME")) } } @@ -251,7 +251,7 @@ pub(crate) fn expand_cuda_module_inner( pub unsafe fn load_async_named( device_id: usize, name: &str, - ) -> ::core::result::Result { + ) -> ::core::result::Result { ::cuda_host::load_cuda_module_from_async_context(device_id, |ctx| { // SAFETY: upheld by this function's caller. unsafe { load_named(ctx, name) } @@ -262,14 +262,14 @@ pub(crate) fn expand_cuda_module_inner( quote! { pub fn load_async( device_id: usize, - ) -> ::core::result::Result { + ) -> ::core::result::Result { load_async_named(device_id, env!("CARGO_PKG_NAME")) } pub fn load_async_named( device_id: usize, name: &str, - ) -> ::core::result::Result { + ) -> ::core::result::Result { ::cuda_host::load_cuda_module_from_async_context(device_id, |ctx| load_named(ctx, name)) } } diff --git a/cuda-oxide/crates/cuda-macros/src/launch.rs b/cuda-oxide/crates/cuda-macros/src/launch.rs index cc47b3830f..2ad12e9753 100644 --- a/cuda-oxide/crates/cuda-macros/src/launch.rs +++ b/cuda-oxide/crates/cuda-macros/src/launch.rs @@ -646,18 +646,18 @@ pub(crate) fn expand_cuda_launch_async(input: CudaLaunchAsyncInput) -> TokenStre let len_name = internal_ident(&format!("__cuda_oxide_arg_{i}_len")); quote! { let #tmp_name = &#expr; - #launch_ident.push_arg(Box::new(#tmp_name.cu_deviceptr())); + #launch_ident.push_arg(#tmp_name.cu_deviceptr()); let #len_name = #tmp_name.len() as u64; - #launch_ident.push_arg(Box::new(#len_name)); + #launch_ident.push_arg(#len_name); } } CudaLaunchArg::SliceMutWithLen(expr) => { let len_name = internal_ident(&format!("__cuda_oxide_arg_{i}_len")); quote! { let #tmp_name = &mut #expr; - #launch_ident.push_arg(Box::new(#tmp_name.cu_deviceptr())); + #launch_ident.push_arg(#tmp_name.cu_deviceptr()); let #len_name = #tmp_name.len() as u64; - #launch_ident.push_arg(Box::new(#len_name)); + #launch_ident.push_arg(#len_name); } } CudaLaunchArg::Closure { .. } => { @@ -694,7 +694,7 @@ pub(crate) fn expand_cuda_launch_async(input: CudaLaunchAsyncInput) -> TokenStre #error_ident, ) }); - let mut #launch_ident = cuda_async::simt::launch::AsyncKernelLaunchBuilder::new( + let mut #launch_ident = cuda_async::launch::AsyncKernelLaunchBuilder::new( std::sync::Arc::new(#function_ident), ); #(#arg_code)* @@ -715,7 +715,7 @@ pub(crate) fn expand_cuda_launch_async(input: CudaLaunchAsyncInput) -> TokenStre #error_ident, ) }); - let mut #launch_ident = cuda_async::simt::launch::AsyncKernelLaunchBuilder::new( + let mut #launch_ident = cuda_async::launch::AsyncKernelLaunchBuilder::new( std::sync::Arc::new(#function_ident), ); #(#arg_code)* @@ -735,7 +735,7 @@ pub(crate) fn expand_cuda_launch_async(input: CudaLaunchAsyncInput) -> TokenStre #error_ident, ) }); - let mut #launch_ident = cuda_async::simt::launch::AsyncKernelLaunchBuilder::new( + let mut #launch_ident = cuda_async::launch::AsyncKernelLaunchBuilder::new( std::sync::Arc::new(#function_ident), ); #(#arg_code)* diff --git a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs b/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs index b10f4cd455..1e653c70ef 100644 --- a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs +++ b/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs @@ -381,6 +381,17 @@ impl Parse for UnrollArgs { }); } + if input.peek(Ident) { + let fork = input.fork(); + let ident: Ident = fork.parse()?; + if ident == "never" && fork.is_empty() { + let _: Ident = input.parse()?; + return Ok(UnrollArgs { + factor: ConstU32Expr::literal(u32::MAX - 1), + }); + } + } + let args: Punctuated = Punctuated::parse_terminated(input)?; let mut values = args .into_iter() diff --git a/cuda-oxide/crates/cuda-macros/src/lib.rs b/cuda-oxide/crates/cuda-macros/src/lib.rs index 0a3d738b55..50e4b423d4 100644 --- a/cuda-oxide/crates/cuda-macros/src/lib.rs +++ b/cuda-oxide/crates/cuda-macros/src/lib.rs @@ -1247,7 +1247,7 @@ pub fn cuda_launch(input: TokenStream) -> TokenStream { /// /// ```ignore /// use cuda_host::cuda_launch_async; -/// use cuda_core::simt::LaunchConfig; +/// use cuda_core::LaunchConfig; /// /// // SAFETY: ABI, lifetimes, geometry, and resources match vecadd. /// let op = unsafe { diff --git a/cuda-oxide/crates/cuda-macros/src/tests/cuda_module.rs b/cuda-oxide/crates/cuda-macros/src/tests/cuda_module.rs index be7a69ac79..3502c5a7f2 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/cuda_module.rs +++ b/cuda-oxide/crates/cuda-macros/src/tests/cuda_module.rs @@ -711,7 +711,7 @@ fn launch_contract_generates_prepared_and_unsafe_raw_paths() { expanded.contains("__cuda_oxide_prepared:&::cuda_core::PreparedLaunch<__map_CudaKernel>") ); assert!(expanded.contains("unsafefnmap_unchecked(")); - assert!(expanded.contains("__cuda_oxide_config:::cuda_core::simt::LaunchConfig")); + assert!(expanded.contains("__cuda_oxide_config:::cuda_core::LaunchConfig")); assert!(expanded.contains("min_alignment:128u32")); assert!(expanded.contains("with_min_compute_capability(8u32,0u32)")); assert!(expanded.contains("with_u32_coordinates()")); diff --git a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs b/cuda-oxide/crates/cuda-macros/src/tests/launch.rs index 4cee9d382d..3043c881b9 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs +++ b/cuda-oxide/crates/cuda-macros/src/tests/launch.rs @@ -371,6 +371,13 @@ fn malformed_unroll_attr_records_error() { fn partial_unroll_factor_must_be_at_least_two() { assert!(syn::parse_str::("0").is_err()); assert!(syn::parse_str::("1").is_err()); + assert_eq!( + syn::parse_str::("never") + .unwrap() + .factor + .literal_value, + Some(u32::MAX - 1) + ); assert_eq!( syn::parse_str::("2") .unwrap() diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.rs index a80b8e0c6f..259dcce9d4 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.rs @@ -19,21 +19,19 @@ impl FakeModule { } mod cuda_async { - pub mod simt { - pub mod launch { - use cuda_core::CudaFunction; - use cuda_core::simt::LaunchConfig; - use std::sync::Arc; + pub mod launch { + use cuda_core::CudaFunction; + use cuda_core::LaunchConfig; + use std::sync::Arc; - pub struct AsyncKernelLaunchBuilder; + pub struct AsyncKernelLaunchBuilder; - impl AsyncKernelLaunchBuilder { - pub fn new(_function: Arc) -> Self { - Self - } - - pub unsafe fn finalize_unchecked(self, _config: LaunchConfig) {} + impl AsyncKernelLaunchBuilder { + pub fn new(_function: Arc) -> Self { + Self } + + pub unsafe fn finalize_unchecked(self, _config: LaunchConfig) {} } } } @@ -42,7 +40,7 @@ fn launch_without_unsafe(module: &FakeModule) { let _ = cuda_launch_async! { kernel: raw, module: module, - config: cuda_core::simt::LaunchConfig { + config: cuda_core::LaunchConfig { grid_dim: (1, 2, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.stderr index 98269f9792..e3825fd8a0 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/async_launch_macro_requires_unsafe.stderr @@ -1,14 +1,14 @@ -error[E0133]: call to unsafe function `cuda_async::simt::launch::AsyncKernelLaunchBuilder::finalize_unchecked` is unsafe and requires unsafe block - --> tests/compile_fail/async_launch_macro_requires_unsafe.rs:42:13 +error[E0133]: call to unsafe function `cuda_async::launch::AsyncKernelLaunchBuilder::finalize_unchecked` is unsafe and requires unsafe block + --> tests/compile_fail/async_launch_macro_requires_unsafe.rs:40:13 | -42 | let _ = cuda_launch_async! { +40 | let _ = cuda_launch_async! { | _____________^ -43 | | kernel: raw, -44 | | module: module, -45 | | config: cuda_core::simt::LaunchConfig { +41 | | kernel: raw, +42 | | module: module, +43 | | config: cuda_core::LaunchConfig { ... | -50 | | args: [], -51 | | }; +48 | | args: [], +49 | | }; | |_____^ call to unsafe function | = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.rs index 402f69f667..6c3734e051 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.rs @@ -22,7 +22,7 @@ fn unchecked_without_unsafe( module .contracted_unchecked( stream, - cuda_core::simt::LaunchConfig { + cuda_core::LaunchConfig { grid_dim: (1, 1, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.stderr index 4a90e0a54a..5d059edbc3 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/contract_unchecked_requires_unsafe.stderr @@ -4,7 +4,7 @@ error[E0133]: call to unsafe function `LoadedModule::contracted_unchecked` is un 22 | / module 23 | | .contracted_unchecked( 24 | | stream, -25 | | cuda_core::simt::LaunchConfig { +25 | | cuda_core::LaunchConfig { ... | 30 | | input, 31 | | ) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_include_kernel_boundary.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_include_kernel_boundary.rs index 1293feb7a5..aa8aca52c5 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_include_kernel_boundary.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_include_kernel_boundary.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use cuda_core::CudaStream; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_macros::cuda_module; #[cuda_module] diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_nested_type_mismatch.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_nested_type_mismatch.rs index 6e43cd2287..59575c3e65 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_nested_type_mismatch.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_nested_type_mismatch.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use cuda_core::CudaStream; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_macros::cuda_module; #[cuda_module] diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_pub_super_scope.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_pub_super_scope.rs index 5fedcf6534..b29ea97625 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_pub_super_scope.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/cuda_module_pub_super_scope.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use cuda_core::CudaStream; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_macros::cuda_module; #[cuda_module] diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.rs index a66be86723..57896d07f0 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.rs @@ -32,7 +32,7 @@ fn launch( // The spelling chose the three-word (ptr, len, width) marshalling, but // the resolved device slice is two words. The ABI bound must reject the // call before a mis-shaped packet can exist. - let cfg = cuda_core::simt::LaunchConfig { + let cfg = cuda_core::LaunchConfig { grid_dim: (1, 1, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.stderr index 6a068a9dd8..a4cedec322 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_fakes_row_width.stderr @@ -12,4 +12,5 @@ note: required by a bound in `LoadedModule::alias_fakes_row_width` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::alias_fakes_row_width` ... 22 | pub fn alias_fakes_row_width(mut out: DisjointSlice) { + | --------------------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.rs index fbbcaff005..88445a9f9f 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.rs @@ -32,7 +32,7 @@ fn launch( // The spelling chose the two-word (ptr, len) marshalling, so this raw // launch would push two kernel parameters for a three-parameter kernel. // The ABI bound must reject the call. - let cfg = cuda_core::simt::LaunchConfig { + let cfg = cuda_core::LaunchConfig { grid_dim: (1, 1, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.stderr index 1e2dd6faa7..4bc77dcf68 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_alias_hides_row_width.stderr @@ -12,4 +12,5 @@ note: required by a bound in `LoadedModule::alias_hides_row_width` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::alias_hides_row_width` ... 22 | pub fn alias_hides_row_width(mut out: DisjointSlice) { + | --------------------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_fake_disjoint_slice.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_fake_disjoint_slice.stderr index a0d8471d81..530af8ce95 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_fake_disjoint_slice.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_fake_disjoint_slice.stderr @@ -26,6 +26,7 @@ note: required by a bound in `LoadedModule::prepare_lookalike` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::prepare_lookalike` ... 21 | pub fn lookalike(mut out: DisjointSlice) { + | --------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) error[E0277]: the trait bound `for<'__cuda_oxide_disjoint> DisjointSlice<'__cuda_oxide_disjoint, u32>: cuda_device::disjoint::__LaunchContractDisjointSliceAbi` is not satisfied @@ -56,4 +57,5 @@ note: required by a bound in `LoadedModule::prepare_lookalike` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::prepare_lookalike` ... 21 | pub fn lookalike(mut out: DisjointSlice) { + | --------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_misleading_index_alias.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_misleading_index_alias.stderr index 3208581be6..50bc23cd6f 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_misleading_index_alias.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_misleading_index_alias.stderr @@ -22,4 +22,5 @@ note: required by a bound in `LoadedModule::prepare_misleading_name` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::prepare_misleading_name` ... 13 | pub fn misleading_name(mut out: DisjointSlice) { + | --------------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_reordered_disjoint_alias.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_reordered_disjoint_alias.stderr index 851917094c..6509665d5d 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_reordered_disjoint_alias.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_reordered_disjoint_alias.stderr @@ -22,6 +22,7 @@ note: required by a bound in `LoadedModule::prepare_reordered` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::prepare_reordered` ... 17 | pub fn reordered(mut out: DisjointSlice) { + | --------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) error[E0277]: the trait bound `for<'__cuda_oxide_disjoint> cuda_device::DisjointSlice<'__cuda_oxide_disjoint, u64>: cuda_device::disjoint::__LaunchContractDisjointSliceAbi` is not satisfied @@ -48,4 +49,5 @@ note: required by a bound in `LoadedModule::prepare_reordered` | ^^^^^^^^^^^^^^ required by this bound in `LoadedModule::prepare_reordered` ... 17 | pub fn reordered(mut out: DisjointSlice) { + | --------- required by a bound in this associated function = note: this error originates in the attribute macro `cuda_module` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.rs index 112ee2f037..6edd83f0b3 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.rs @@ -28,7 +28,7 @@ fn launch_without_unsafe( kernel: dummy, stream: stream, module: module, - config: cuda_core::simt::LaunchConfig::for_num_elems(1), + config: cuda_core::LaunchConfig::for_num_elems(1), args: [x] } .unwrap(); diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.stderr index f63438e3eb..fe6ee6e6b3 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_requires_unsafe.stderr @@ -5,7 +5,7 @@ error[E0133]: call to unsafe function `launch_kernel_on_stream` is unsafe and re 28 | | kernel: dummy, 29 | | stream: stream, 30 | | module: module, -31 | | config: cuda_core::simt::LaunchConfig::for_num_elems(1), +31 | | config: cuda_core::LaunchConfig::for_num_elems(1), 32 | | args: [x] 33 | | } | |_____^ call to unsafe function diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_invalid_unroll.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_invalid_unroll.stderr index fe4d9eb24e..d79f0da8ea 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_invalid_unroll.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_invalid_unroll.stderr @@ -1,7 +1,6 @@ error[E0080]: evaluation panicked: partial unroll factor must be in 2..=1024, or 0 for full unrolling --> $WORKSPACE/crates/cuda-device/src/thread.rs | - | pub fn __unroll_config() { | const { validate_unroll_factor(FACTOR) } | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cuda_device::__unroll_config::<1>::{constant#0}` failed inside this call | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_unresolved.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_unresolved.stderr index 970bdf2d2e..6cf8574256 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_unresolved.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_unresolved.stderr @@ -1,14 +1,3 @@ -warning: `feature(generic_const_exprs)` is not supported with the next-generation trait solver - --> tests/compile_fail/policy_config_unresolved.rs:4:12 - | -4 | #![feature(generic_const_exprs)] - | ^^^^^^^^^^^^^^^^^^^ - | - = note: `-Znext-solver=globally` is currently enabled by default for testing - = note: reverted the setting to `-Znext-solver=coherence` for this crate - = note: the currently stable trait solver will be used for this crate - = note: see issues #160895 for more information - error[E0599]: no associated function or constant named `MAX_THREADS` found for type parameter `P` in the current scope --> tests/compile_fail/policy_config_unresolved.rs:10:20 | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_zero_threads.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_zero_threads.stderr index b6d60f72e7..9d71df7e96 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_zero_threads.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_config_zero_threads.stderr @@ -1,7 +1,6 @@ error[E0080]: evaluation panicked: launch_bounds maximum threads must be greater than zero --> $WORKSPACE/crates/cuda-device/src/thread.rs | - | pub fn __launch_bounds_config() { | const { validate_launch_bounds(MAX_THREADS) } | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cuda_device::__launch_bounds_config::<0, 0>::{constant#0}` failed inside this call | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_exact_block_too_large.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_exact_block_too_large.stderr index e65b8597b7..abdd18341c 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_exact_block_too_large.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_exact_block_too_large.stderr @@ -1,14 +1,3 @@ -warning: `feature(generic_const_exprs)` is not supported with the next-generation trait solver - --> tests/compile_fail/policy_contract_exact_block_too_large.rs:4:12 - | -4 | #![feature(generic_const_exprs)] - | ^^^^^^^^^^^^^^^^^^^ - | - = note: `-Znext-solver=globally` is currently enabled by default for testing - = note: reverted the setting to `-Znext-solver=coherence` for this crate - = note: the currently stable trait solver will be used for this crate - = note: see issues #160895 for more information - error[E0080]: evaluation panicked: launch_contract exact block exceeds launch_bounds maximum threads --> tests/compile_fail/policy_contract_exact_block_too_large.rs:20:1 | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_large_3d_block.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_large_3d_block.stderr index 0b2d131856..74e3e4d53a 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_large_3d_block.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_large_3d_block.stderr @@ -1,14 +1,3 @@ -warning: `feature(generic_const_exprs)` is not supported with the next-generation trait solver - --> tests/compile_fail/policy_contract_large_3d_block.rs:4:12 - | -4 | #![feature(generic_const_exprs)] - | ^^^^^^^^^^^^^^^^^^^ - | - = note: `-Znext-solver=globally` is currently enabled by default for testing - = note: reverted the setting to `-Znext-solver=coherence` for this crate - = note: the currently stable trait solver will be used for this crate - = note: see issues #160895 for more information - error[E0080]: evaluation panicked: launch_contract exact block exceeds launch_bounds maximum threads --> tests/compile_fail/policy_contract_large_3d_block.rs:20:1 | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_wrong_brand.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_wrong_brand.stderr index 4769fb9425..b08b165228 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_wrong_brand.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_wrong_brand.stderr @@ -1,14 +1,3 @@ -warning: `feature(generic_const_exprs)` is not supported with the next-generation trait solver - --> tests/compile_fail/policy_contract_wrong_brand.rs:4:12 - | -4 | #![feature(generic_const_exprs)] - | ^^^^^^^^^^^^^^^^^^^ - | - = note: `-Znext-solver=globally` is currently enabled by default for testing - = note: reverted the setting to `-Znext-solver=coherence` for this crate - = note: the currently stable trait solver will be used for this crate - = note: see issues #160895 for more information - error[E0308]: mismatched types --> tests/compile_fail/policy_contract_wrong_brand.rs:41:16 | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_zero_threads.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_zero_threads.stderr index 7b867238bb..a46d8be4a4 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_zero_threads.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/policy_contract_zero_threads.stderr @@ -1,14 +1,3 @@ -warning: `feature(generic_const_exprs)` is not supported with the next-generation trait solver - --> tests/compile_fail/policy_contract_zero_threads.rs:4:12 - | -4 | #![feature(generic_const_exprs)] - | ^^^^^^^^^^^^^^^^^^^ - | - = note: `-Znext-solver=globally` is currently enabled by default for testing - = note: reverted the setting to `-Znext-solver=coherence` for this crate - = note: the currently stable trait solver will be used for this crate - = note: see issues #160895 for more information - error[E0080]: evaluation panicked: launch_bounds maximum threads must be greater than zero --> tests/compile_fail/policy_contract_zero_threads.rs:20:1 | diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.rs index 5a1fa7d265..b2394780a5 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.rs @@ -20,7 +20,7 @@ fn launch_without_unsafe( module .uncontracted( stream, - cuda_core::simt::LaunchConfig { + cuda_core::LaunchConfig { grid_dim: (1, 2, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.stderr index 19a3595186..ea7d2fd9a1 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/uncontracted_launch_requires_unsafe.stderr @@ -4,7 +4,7 @@ error[E0133]: call to unsafe function `LoadedModule::uncontracted` is unsafe and 20 | / module 21 | | .uncontracted( 22 | | stream, -23 | | cuda_core::simt::LaunchConfig { +23 | | cuda_core::LaunchConfig { ... | 28 | | 7, 29 | | ) diff --git a/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock b/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock index 4a78f429b9..1054e0ea68 100644 --- a/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock +++ b/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock @@ -2,11 +2,19 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -20,6 +28,10 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "device-only-kernels" version = "0.0.0" diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/async_launch_macro_in_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/pass/async_launch_macro_in_unsafe.rs index c983dd75f5..b2986f44d7 100644 --- a/cuda-oxide/crates/cuda-macros/tests/pass/async_launch_macro_in_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/pass/async_launch_macro_in_unsafe.rs @@ -19,21 +19,19 @@ impl FakeModule { } mod cuda_async { - pub mod simt { - pub mod launch { - use cuda_core::CudaFunction; - use cuda_core::simt::LaunchConfig; - use std::sync::Arc; + pub mod launch { + use cuda_core::CudaFunction; + use cuda_core::LaunchConfig; + use std::sync::Arc; - pub struct AsyncKernelLaunchBuilder; + pub struct AsyncKernelLaunchBuilder; - impl AsyncKernelLaunchBuilder { - pub fn new(_function: Arc) -> Self { - Self - } - - pub unsafe fn finalize_unchecked(self, _config: LaunchConfig) {} + impl AsyncKernelLaunchBuilder { + pub fn new(_function: Arc) -> Self { + Self } + + pub unsafe fn finalize_unchecked(self, _config: LaunchConfig) {} } } } @@ -43,7 +41,7 @@ fn launch_with_unsafe(module: &FakeModule) { cuda_launch_async! { kernel: raw, module: module, - config: cuda_core::simt::LaunchConfig { + config: cuda_core::LaunchConfig { grid_dim: (1, 2, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/const_generic_hygiene.rs b/cuda-oxide/crates/cuda-macros/tests/pass/const_generic_hygiene.rs index 0ede621af9..85fb0a1bca 100644 --- a/cuda-oxide/crates/cuda-macros/tests/pass/const_generic_hygiene.rs +++ b/cuda-oxide/crates/cuda-macros/tests/pass/const_generic_hygiene.rs @@ -3,7 +3,7 @@ #![allow(dead_code, non_upper_case_globals)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaModule, CudaStream}; use cuda_macros::{cuda_launch, kernel}; use std::sync::Arc; diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/cuda_module_inline_namespaces.rs b/cuda-oxide/crates/cuda-macros/tests/pass/cuda_module_inline_namespaces.rs index c964a060d1..aa475b823e 100644 --- a/cuda-oxide/crates/cuda-macros/tests/pass/cuda_module_inline_namespaces.rs +++ b/cuda-oxide/crates/cuda-macros/tests/pass/cuda_module_inline_namespaces.rs @@ -5,7 +5,7 @@ #![allow(dead_code)] use cuda_core::CudaStream; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_macros::{cuda_module, kernel}; #[derive(Clone, Copy)] diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/grid_constant_arguments.rs b/cuda-oxide/crates/cuda-macros/tests/pass/grid_constant_arguments.rs index 779aafd31c..fe07e8383c 100644 --- a/cuda-oxide/crates/cuda-macros/tests/pass/grid_constant_arguments.rs +++ b/cuda-oxide/crates/cuda-macros/tests/pass/grid_constant_arguments.rs @@ -2,7 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 #![allow(dead_code)] +#![feature(f16)] +use ctm_rs::{memory::F16Tile128, tma::TensorMap}; use cuda_device::{cuda_module, kernel}; #[derive(Clone, Copy)] @@ -13,6 +15,11 @@ pub struct Descriptor([u32; 32]); mod kernels { use super::*; + #[kernel] + pub fn typed(#[grid_constant] descriptor: &TensorMap) { + let _ = descriptor.as_descriptor(); + } + #[kernel] pub fn plain(#[grid_constant] descriptor: &Descriptor) { let _ = descriptor.0[0]; @@ -33,11 +40,13 @@ fn typed_launches( module: &kernels::LoadedModule, stream: &cuda_core::CudaStream, descriptor: Descriptor, + typed_descriptor: TensorMap, ) { - let config = cuda_core::simt::LaunchConfig::for_num_elems(1); + let config = cuda_core::LaunchConfig::for_num_elems(1); // SAFETY: this function only type-checks the by-value host ABI; it is never run. unsafe { module.plain(stream, config, descriptor).unwrap(); + module.typed(stream, config, typed_descriptor).unwrap(); module.generic(stream, config, descriptor, 3u32).unwrap(); } } diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/uncontracted_launch_in_unsafe.rs b/cuda-oxide/crates/cuda-macros/tests/pass/uncontracted_launch_in_unsafe.rs index 5fc62ce57c..0fa24556be 100644 --- a/cuda-oxide/crates/cuda-macros/tests/pass/uncontracted_launch_in_unsafe.rs +++ b/cuda-oxide/crates/cuda-macros/tests/pass/uncontracted_launch_in_unsafe.rs @@ -21,7 +21,7 @@ fn launch_with_unsafe( module .uncontracted( stream, - cuda_core::simt::LaunchConfig { + cuda_core::LaunchConfig { grid_dim: (1, 2, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0, diff --git a/cuda-oxide/crates/cuda-oxide-codegen/Cargo.toml b/cuda-oxide/crates/cuda-oxide-codegen/Cargo.toml index 0b2b091e6e..4d6164ac34 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/Cargo.toml +++ b/cuda-oxide/crates/cuda-oxide-codegen/Cargo.toml @@ -18,6 +18,8 @@ thiserror = { workspace = true } dialect-mir = { workspace = true } dialect-iket = { workspace = true } dialect-nvvm = { workspace = true } +dialect-cute = { workspace = true } +cuda-oxide-mlir-export = { workspace = true } iket-lower = { workspace = true } mir-lower = { workspace = true } mir-transforms = { workspace = true } @@ -26,3 +28,7 @@ libnvvm-sys = { workspace = true } nvvm-transforms = { workspace = true } tempfile = { workspace = true } rustc-demangle = { workspace = true } +serde_json = { workspace = true } +sha2 = { workspace = true } +libloading = { workspace = true } +object = { workspace = true } diff --git a/cuda-oxide/crates/cuda-oxide-codegen/README.md b/cuda-oxide/crates/cuda-oxide-codegen/README.md index 9ac68d1d56..db30c9df38 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/README.md +++ b/cuda-oxide/crates/cuda-oxide-codegen/README.md @@ -8,7 +8,7 @@ lowering, and LLVM tooling as the rustc frontend uses. The only supported public surface is the `experimental` module. The point is the absence of a dependency: this crate has **no rustc linkage**. -It needs neither `rustc_private` nor a nightly toolchain matched to +It needs neither `rustc_private` nor a toolchain pin matched to `rustc_driver`, so a caller that can build the IR itself can reach PTX without the compiler-plugin machinery `rustc-codegen-cuda` requires. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs index 88ffa0fc6d..0b061f09ac 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs @@ -666,6 +666,7 @@ impl Compiler { let scratch = ScratchDirectory::new()?; let ll_path = scratch.path().join("module.ll"); let ptx_path = scratch.path().join("module.ptx"); + let cubin_path = scratch.path().join("module.cubin"); let backend_options = BackendOptions { iket: crate::options::IketInstrumentation::Auto, target_arch: Some(options.target.sm()), @@ -687,6 +688,8 @@ impl Compiler { OutputFiles { llvm_ir: &ll_path, ptx: &ptx_path, + cubin: &cubin_path, + mlir: None, stale_before_export: &[], }, ); @@ -1009,6 +1012,7 @@ impl From for CompileError { Self::LibdeviceUnavailable { message } } PipelineError::Export(message) => Self::Export { message }, + PipelineError::BackendB(message) => Self::Codegen { message }, PipelineError::TargetSelection { target, reason } => { Self::TargetSelection { target, reason } } diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs new file mode 100644 index 0000000000..bee7f1349d --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_backend.rs @@ -0,0 +1,1630 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile MLIR with the pinned CUTLASS library and validate its CUDA image. + +use crate::cutlass_compiler::CutlassCompilerLibrary; +use crate::error::PipelineError; +use object::{Object, ObjectSection, ObjectSymbol}; +use serde_json::json; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; +use std::io::Write; +use std::path::{Path, PathBuf}; + +/// Exact public CUTLASS profile implemented by this checkout. +pub const CUTLASS_FULL_CUTE_MLIR22_PROFILE: &str = "cutlass-full-cute-mlir22"; + +/// Official CUTLASS release providing the compiler-library ABI. +pub const CUTLASS_COMPILER_VERSION: &str = "4.7.0"; + +/// SHA-256 of the official x86_64/CUDA-13 CUTLASS 4.7 compiler library. +pub const CUTLASS_COMPILER_LIBRARY_SHA256: &str = + "57df017e3585a10443c74c8b4cd99bda854242fb2f4c9534cf56d58c2c741628"; + +/// Version-pinned ABI mode for external `gpu.module` device images. Concrete +/// CUTLASS wrapper ABIs require FunctionMetadata and cannot represent this +/// direct CUDA-launch path. +const CUTLASS_EXTERNAL_MODULE_ABI: &str = + "tbd-external-module-passthrough; zero FunctionMetadata; validated CUDA KPARAM layout"; + +/// Exact first-stage pipeline validated with the fingerprint-pinned library. +/// +/// The C API pass manager is already anchored on `builtin.module`; adding a +/// second `builtin.module(...)` wrapper changes the pass nesting and does not +/// lower this input. The leading canonicalization is also part of the pinned +/// contract: CUTLASS 4.7's default first stage does not accept the generic +/// operation form emitted by the profile. +const CUTLASS_PRECOMPILED_PIPELINE: &str = "canonicalize,cute-to-nvvm{check-inline-asm=false cubin-format=bin opt-level=3 use-software-pipeline-pass=true use-fold-static-pass=true use-insert-range-information-pass=true use-loop-invariant-pass=true use-strength-reduction-pass=true use-eliminate-unnecessary-sync-pass=true use-infer-loop-attrs-pass=true enable-cuda-dialect=true cuda-dialect-external-module=true}"; + +#[cfg(test)] +#[path = "tests/cutlass/sm100.rs"] +mod sm100_tests; + +#[cfg(test)] +#[path = "tests/cutlass/tcgen05.rs"] +mod tcgen05_tests; + +#[cfg(test)] +#[path = "tests/cutlass/grid_constant.rs"] +mod grid_constant_tests; + +#[cfg(test)] +#[path = "tests/cutlass/primitives.rs"] +mod primitives_tests; + +/// Compiler library and MLIR profile used by the CUTLASS backend. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutlassBackendConfig { + /// Absolute path to the exact `libCutlassCompiler.so` installed by + /// `cargo oxide toolchain install cutlass`. + pub compiler_library: PathBuf, + /// Pinned MLIR mapping profile. Only reviewed in-tree profiles are accepted. + pub profile: String, + /// Official compiler optimization level (2 or 3), defaulting to 3. + pub optimization_level: u32, +} + +impl CutlassBackendConfig { + /// Construct a backend configuration for the pinned in-tree profile. + pub fn new(compiler_library: PathBuf) -> Self { + Self { + compiler_library, + profile: CUTLASS_FULL_CUTE_MLIR22_PROFILE.to_owned(), + optimization_level: 3, + } + } +} + +fn precompiled_pipeline(config: &CutlassBackendConfig) -> Result { + if !matches!(config.optimization_level, 2 | 3) { + return Err(backend_error("CUTLASS optimization level must be 2 or 3")); + } + Ok(CUTLASS_PRECOMPILED_PIPELINE.replace( + "opt-level=3", + &format!("opt-level={}", config.optimization_level), + )) +} + +/// Device compiler selected after the shared MIR/CuTe preparation stages. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum DeviceBackend { + /// CUDA Oxide MIR/NVVM/LLVM backend. + #[default] + Native, + /// Official CUTLASS compiler-library path, producing a CUDA image. + CutlassMlir(CutlassBackendConfig), +} + +#[derive(Debug)] +pub(crate) struct CutlassRunOutput { + pub(crate) diagnostics: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum KernelParameterKind { + GlobalPointer64, + ScalarBool, + /// Physical four-byte CUDA parameter layout shared by i32 and f32. + Scalar32, + Scalar64, + GridConstant { + size: u16, + alignment: u16, + }, +} + +impl KernelParameterKind { + const fn size(self) -> u16 { + match self { + Self::GlobalPointer64 | Self::Scalar64 => 8, + Self::Scalar32 => 4, + // MLIR i1 lowers to a one-byte CUDA parameter, matching Rust bool. + Self::ScalarBool => 1, + Self::GridConstant { size, .. } => size, + } + } + + const fn alignment(self) -> u16 { + match self { + Self::GlobalPointer64 | Self::Scalar64 => 8, + Self::Scalar32 => 4, + Self::ScalarBool => 1, + Self::GridConstant { alignment, .. } => alignment, + } + } + + const fn cuda_parameter_space(self) -> u32 { + match self { + Self::GlobalPointer64 => 5, + Self::ScalarBool | Self::Scalar32 | Self::Scalar64 | Self::GridConstant { .. } => 0, + } + } +} + +type ExpectedKernelAbis = BTreeMap>; + +/// Compile one rendered profile module and atomically publish a verified CUDA +/// image suitable for `cuModuleLoadData`. +/// +/// CUTLASS 4.7 returns a host relocatable object. Its exact, fingerprint-pinned +/// contract contains one local data symbol named `kernels_binary`; that symbol +/// is either a raw CUDA ELF image or a CUDA fatbinary. A fatbinary must contain +/// exactly one matching, uncompressed cubin, which is extracted so the public +/// `Cubin` artifact kind always carries an actual CUDA ELF image. +pub(crate) fn compile_to_cubin( + config: &CutlassBackendConfig, + rendered_mlir: &str, + target: &str, + expected_kernels: &[String], + cubin_output: &Path, + mlir_output: Option<&Path>, +) -> Result { + if !matches!( + config.profile.as_str(), + CUTLASS_FULL_CUTE_MLIR22_PROFILE | cuda_oxide_mlir_export::CutlassPrimitivesMlir22::NAME + ) { + return Err(backend_error(format!( + "unsupported MLIR profile {:?}; expected {:?} or {:?}", + config.profile, + CUTLASS_FULL_CUTE_MLIR22_PROFILE, + cuda_oxide_mlir_export::CutlassPrimitivesMlir22::NAME + ))); + } + if expected_kernels.is_empty() { + return Err(backend_error( + "the CUTLASS backend requires at least one named kernel entry", + )); + } + parse_sm_target(target)?; + let pipeline = precompiled_pipeline(config)?; + + let compiler = + CutlassCompilerLibrary::load(&config.compiler_library, CUTLASS_COMPILER_LIBRARY_SHA256) + .map_err(|error| { + backend_error(format!("could not load official CUTLASS compiler: {error}")) + })?; + let execution_mlir = add_execution_manifest( + rendered_mlir, + compiler.path(), + &compiler.fingerprint().to_string(), + target, + &config.profile, + &pipeline, + config.optimization_level, + )?; + let expected_abis = expected_kernel_abis(&execution_mlir, expected_kernels) + .map_err(|reason| backend_error(format!("invalid Backend-B kernel ABI: {reason}")))?; + if let Some(path) = mlir_output { + publish_atomically(path, execution_mlir.as_bytes())?; + } + + let object = compiler + .compile_precompiled_mlir_to_object(&execution_mlir, target, &pipeline) + .map_err(|error| backend_error(format!("official CUTLASS compilation failed: {error}")))?; + let image = extract_kernels_binary(&object.elf) + .map_err(|reason| backend_error(format!("invalid CUTLASS ObjectArtifact: {reason}")))?; + let (cubin, source_kind) = + extract_validated_cubin(&image, target, &expected_abis).map_err(|reason| { + backend_error(format!( + "CUTLASS ObjectArtifact contains an invalid CUDA image: {reason}" + )) + })?; + publish_atomically(cubin_output, &cubin)?; + + Ok(CutlassRunOutput { + diagnostics: vec![format!( + "official CUTLASS {CUTLASS_COMPILER_VERSION} emitted a validated direct-ABI cubin extracted from {source_kind} ({} bytes)", + cubin.len() + )], + }) +} + +fn add_execution_manifest( + rendered_mlir: &str, + compiler_library: &Path, + compiler_sha256: &str, + target: &str, + profile: &str, + pipeline: &str, + optimization_level: u32, +) -> Result { + let pipeline_sha256 = format!("{:x}", Sha256::digest(pipeline)); + let manifest = json!({ + "profile": profile, + "target": target, + "cutlass_compiler_version": CUTLASS_COMPILER_VERSION, + "cutlass_compiler_library": compiler_library, + "cutlass_compiler_library_sha256": compiler_sha256, + "cutlass_external_module_abi": CUTLASS_EXTERNAL_MODULE_ABI, + "precompiled_pipeline_sha256": pipeline_sha256, + "optimization_level": optimization_level, + }); + let manifest = serde_json::to_string_pretty(&manifest).map_err(|error| { + backend_error(format!( + "could not render CUTLASS execution manifest: {error}" + )) + })?; + let mut output = String::from("// cuda-oxide CUTLASS Backend-B execution manifest\n"); + for line in manifest.lines() { + output.push_str("// "); + output.push_str(line); + output.push('\n'); + } + output.push_str(rendered_mlir); + Ok(output) +} + +/// Read the expected launch argument sizes and alignments from exported MLIR. +/// +/// A Rust slice uses two arguments: pointer, then `i64` length. This profile also +/// accepts `i1`, `i32`, `i64`, and `f32` scalars in their original representations. +/// A grid-constant reference carries the value's bytes at its Rust alignment; +/// MLIR represents its address as an opaque pointer. +fn expected_kernel_abis( + rendered_mlir: &str, + expected_kernels: &[String], +) -> Result { + const PREFIX: &str = "\"func.func\"() <{"; + let mut abis = BTreeMap::new(); + for kernel in expected_kernels { + if kernel.contains(['"', '\\']) { + return Err(format!("kernel name {kernel:?} cannot be matched safely")); + } + if abis.contains_key(kernel) { + return Err(format!("duplicate expected kernel {kernel:?}")); + } + let name_marker = format!("sym_name = \"{kernel}\""); + let signatures = rendered_mlir + .lines() + .filter_map(|line| { + let properties = line.trim().strip_prefix(PREFIX)?; + let (properties, _) = properties.split_once("}> ({")?; + properties.contains(&name_marker).then_some(properties) + }) + .collect::>(); + if signatures.len() != 1 { + return Err(format!( + "expected exactly one flattened func.func signature for kernel {kernel:?}, found {}", + signatures.len() + )); + } + let properties = parse_abi_dictionary(signatures[0])?; + let quoted_name = format!("\"{kernel}\""); + if properties.get("sym_name") != Some("ed_name.as_str()) { + return Err(format!("kernel {kernel:?} has an ambiguous symbol name")); + } + for property in properties.keys() { + if !matches!(*property, "sym_name" | "function_type" | "arg_attrs") { + return Err(format!( + "kernel {kernel:?} has unsupported function property {property:?}" + )); + } + } + let signature = properties + .get("function_type") + .and_then(|value| value.strip_prefix('(')?.strip_suffix(") -> ()")) + .ok_or_else(|| format!("kernel {kernel:?} has an unsupported function type"))?; + let attributes = properties + .get("arg_attrs") + .map(|value| { + let entries = value + .strip_prefix('[') + .and_then(|value| value.strip_suffix(']')) + .ok_or_else(|| format!("kernel {kernel:?} arg_attrs is not an array"))?; + split_abi_items(entries) + }) + .transpose()?; + let types = split_abi_items(signature)?; + if let Some(attributes) = &attributes + && attributes.len() != types.len() + { + return Err(format!( + "kernel {kernel:?} has {} argument attributes for {} parameters", + attributes.len(), + types.len() + )); + } + let parameters = types + .iter() + .enumerate() + .map(|(ordinal, parameter)| { + let attributes = attributes.as_ref().map(|items| items[ordinal]); + parse_parameter_abi(parameter, attributes) + .map_err(|reason| format!("kernel {kernel:?} parameter {ordinal}: {reason}")) + }) + .collect::, _>>()?; + abis.insert(kernel.clone(), parameters); + } + Ok(abis) +} + +fn parse_parameter_abi( + parameter: &str, + attributes: Option<&str>, +) -> Result { + let attributes = attributes + .map(|value| { + let entries = value + .strip_prefix('{') + .and_then(|value| value.strip_suffix('}')) + .ok_or_else(|| "argument attributes must be a dictionary".to_owned())?; + parse_abi_dictionary(entries) + }) + .transpose()? + .unwrap_or_default(); + if !attributes.is_empty() { + if parameter != "!llvm.ptr" + || attributes.len() != 3 + || attributes.get("nvvm.grid_constant") != Some(&"unit") + { + return Err("unsupported argument attributes; expected an aligned grid-constant by-value pointer".to_owned()); + } + let size = attributes + .get("llvm.byval") + .and_then(|value| value.strip_prefix("!llvm.array<")?.strip_suffix(" x i8>")) + .and_then(|value| value.parse::().ok()) + .filter(|size| (1..=0x3fff).contains(size)) + .ok_or_else(|| { + "llvm.byval must be a nonempty byte array fitting the CUDA parameter size field" + .to_owned() + })?; + let alignment = attributes + .get("llvm.align") + .and_then(|value| value.strip_suffix(" : i64")) + .and_then(|value| value.parse::().ok()) + .filter(|alignment| alignment.is_power_of_two()) + .ok_or_else(|| { + "llvm.align must be a nonzero power-of-two i64 fitting the CUDA parameter block" + .to_owned() + })?; + return Ok(KernelParameterKind::GridConstant { size, alignment }); + } + match parameter { + "!llvm.ptr" => Ok(KernelParameterKind::GlobalPointer64), + "i1" => Ok(KernelParameterKind::ScalarBool), + "i32" | "f32" => Ok(KernelParameterKind::Scalar32), + "i64" => Ok(KernelParameterKind::Scalar64), + other => Err(format!( + "unsupported flattened parameter type {other:?}; only !llvm.ptr, i1, i32, i64, and f32 have a pinned direct-launch ABI" + )), + } +} + +fn parse_abi_dictionary(input: &str) -> Result, String> { + let mut result = BTreeMap::new(); + for item in split_abi_items(input)? { + let (name, value) = if item == "nvvm.grid_constant" { + (item, "unit") + } else { + item.split_once(" = ") + .ok_or_else(|| format!("malformed ABI dictionary entry {item:?}"))? + }; + if result.insert(name, value).is_some() { + return Err(format!("duplicate ABI dictionary entry {name:?}")); + } + } + Ok(result) +} + +/// Split argument lists at commas, keeping nested types and attributes intact. +fn split_abi_items(input: &str) -> Result, String> { + if input.trim().is_empty() { + return Ok(Vec::new()); + } + let mut items = Vec::new(); + let mut brackets = Vec::new(); + let mut quoted = false; + let mut escaped = false; + let mut start = 0; + for (offset, byte) in input.bytes().enumerate() { + if quoted { + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + quoted = false; + } + continue; + } + match byte { + b'"' => quoted = true, + b'(' => brackets.push(b')'), + b'[' => brackets.push(b']'), + b'{' => brackets.push(b'}'), + b'<' => brackets.push(b'>'), + b'>' if offset > 0 && input.as_bytes()[offset - 1] == b'-' => {} + b')' | b']' | b'}' | b'>' => { + if brackets.pop() != Some(byte) { + return Err("unbalanced delimiters in kernel ABI".to_owned()); + } + } + b',' if brackets.is_empty() => { + items.push(input[start..offset].trim()); + start = offset + 1; + } + _ => {} + } + } + if quoted || !brackets.is_empty() { + return Err("unterminated kernel ABI property".to_owned()); + } + items.push(input[start..].trim()); + if items.contains(&"") { + return Err("empty entry in kernel ABI".to_owned()); + } + Ok(items) +} + +fn extract_kernels_binary(object_bytes: &[u8]) -> Result, String> { + let file = object::File::parse(object_bytes).map_err(|error| error.to_string())?; + if file.format() != object::BinaryFormat::Elf + || file.kind() != object::ObjectKind::Relocatable + || file.architecture() != object::Architecture::X86_64 + { + return Err(format!( + "expected an x86-64 ELF relocatable, got {:?}/{:?}/{:?}", + file.format(), + file.kind(), + file.architecture() + )); + } + + let mut symbols = file + .symbols() + .filter(|symbol| symbol.name() == Ok("kernels_binary")); + let symbol = symbols + .next() + .ok_or_else(|| "missing exact `kernels_binary` data symbol".to_owned())?; + if symbols.next().is_some() { + return Err("multiple `kernels_binary` symbols are ambiguous".to_owned()); + } + if !symbol.is_definition() || symbol.kind() != object::SymbolKind::Data || symbol.size() == 0 { + return Err("`kernels_binary` is not a nonempty defined data symbol".to_owned()); + } + let section_index = symbol + .section_index() + .ok_or_else(|| "`kernels_binary` does not name an object section".to_owned())?; + let section = file + .section_by_index(section_index) + .map_err(|error| error.to_string())?; + let section_data = section.data().map_err(|error| error.to_string())?; + let relative = symbol + .address() + .checked_sub(section.address()) + .ok_or_else(|| "`kernels_binary` address precedes its section".to_owned())?; + let start = usize::try_from(relative) + .map_err(|_| "`kernels_binary` offset does not fit in memory".to_owned())?; + let size = usize::try_from(symbol.size()) + .map_err(|_| "`kernels_binary` size does not fit in memory".to_owned())?; + let end = start + .checked_add(size) + .ok_or_else(|| "`kernels_binary` range overflows".to_owned())?; + section_data + .get(start..end) + .map(<[u8]>::to_vec) + .ok_or_else(|| "`kernels_binary` range is outside its section".to_owned()) +} + +fn extract_validated_cubin( + bytes: &[u8], + target: &str, + expected_abis: &ExpectedKernelAbis, +) -> Result<(Vec, &'static str), String> { + const FATBIN_MAGIC: u32 = 0xba55_ed50; + if bytes.starts_with(b"\x7fELF") { + validate_cuda_cubin(bytes, target, expected_abis)?; + return Ok((bytes.to_vec(), "a raw cubin")); + } + if read_u32(bytes, 0)? != FATBIN_MAGIC { + return Err("image is neither a CUDA ELF cubin nor a CUDA fatbinary".to_owned()); + } + let cubin = extract_cubin_from_fatbin(bytes, target, expected_abis)?; + Ok((cubin, "a CUDA fatbinary")) +} + +fn extract_cubin_from_fatbin( + bytes: &[u8], + target: &str, + expected_abis: &ExpectedKernelAbis, +) -> Result, String> { + const HEADER_SIZE: usize = 16; + const COMMON_SIZE: usize = 64; + const FATBIN_VERSION: u16 = 1; + const FATBIN_KIND_ELF: u16 = 2; + const COMPRESSED_FLAGS: u64 = 0xf000; + + if bytes.len() < HEADER_SIZE { + return Err("truncated CUDA fatbinary header".to_owned()); + } + if read_u16(bytes, 4)? != FATBIN_VERSION { + return Err("unsupported CUDA fatbinary version".to_owned()); + } + let header_size = usize::from(read_u16(bytes, 6)?); + if header_size < HEADER_SIZE || header_size % 8 != 0 { + return Err("invalid CUDA fatbinary header size".to_owned()); + } + let payload_size = usize::try_from(read_u64(bytes, 8)?) + .map_err(|_| "CUDA fatbinary size does not fit in memory".to_owned())?; + let total_size = header_size + .checked_add(payload_size) + .ok_or_else(|| "CUDA fatbinary size overflows".to_owned())?; + if total_size != bytes.len() || payload_size % 8 != 0 { + return Err("CUDA fatbinary length does not match its header".to_owned()); + } + + let expected_sm = parse_sm_target_value(target)?; + let mut cursor = header_size; + let mut cubin = None; + while cursor < bytes.len() { + if bytes.len() - cursor < COMMON_SIZE { + return Err("truncated CUDA fatbinary code header".to_owned()); + } + let kind = read_u16(bytes, cursor)?; + let version = read_u16(bytes, cursor + 2)?; + let code_offset = usize::try_from(read_u32(bytes, cursor + 4)?) + .map_err(|_| "fatbinary code offset does not fit in memory".to_owned())?; + let code_size = usize::try_from(read_u64(bytes, cursor + 8)?) + .map_err(|_| "fatbinary code size does not fit in memory".to_owned())?; + let compressed_size = read_u32(bytes, cursor + 16)?; + let arch = read_u32(bytes, cursor + 28)?; + let flags = read_u64(bytes, cursor + 40)?; + if kind != FATBIN_KIND_ELF || version != 0x0101 { + return Err(format!( + "unsupported fatbinary entry kind/version {kind:#x}/{version:#x}" + )); + } + if compressed_size != 0 || flags & COMPRESSED_FLAGS != 0 { + return Err( + "compressed CUDA fatbinary entries are not supported by this pinned contract" + .to_owned(), + ); + } + if arch != expected_sm { + return Err(format!( + "fatbinary entry targets sm_{arch}, expected {target}" + )); + } + if code_offset < COMMON_SIZE || code_offset % 8 != 0 || code_size == 0 { + return Err("invalid CUDA fatbinary code range".to_owned()); + } + let code_start = cursor + .checked_add(code_offset) + .ok_or_else(|| "fatbinary code offset overflows".to_owned())?; + let next = code_start + .checked_add(code_size) + .ok_or_else(|| "fatbinary code size overflows".to_owned())?; + let code = bytes + .get(code_start..next) + .ok_or_else(|| "fatbinary code range is out of bounds".to_owned())?; + validate_cuda_cubin(code, target, expected_abis)?; + if cubin.replace(code.to_vec()).is_some() { + return Err( + "CUDA fatbinary contains multiple cubin entries for one exact-target compilation" + .to_owned(), + ); + } + cursor = next; + } + cubin.ok_or_else(|| "CUDA fatbinary contains no cubin entries".to_owned()) +} + +fn validate_cuda_cubin( + bytes: &[u8], + target: &str, + expected_abis: &ExpectedKernelAbis, +) -> Result<(), String> { + const ELF_HEADER_SIZE: usize = 64; + const SECTION_HEADER_SIZE: usize = 64; + const EM_CUDA: u16 = 190; + + if bytes.len() < ELF_HEADER_SIZE { + return Err(format!("cubin is only {} bytes", bytes.len())); + } + if &bytes[..4] != b"\x7fELF" || bytes[4] != 2 || bytes[5] != 1 || bytes[6] != 1 { + return Err("expected a little-endian ELF64 cubin".to_owned()); + } + if read_u16(bytes, 16)? != 2 || read_u16(bytes, 18)? != EM_CUDA { + return Err("ELF image is not an executable NVIDIA CUDA cubin".to_owned()); + } + let expected_sm = parse_sm_target_value(target)?; + let actual_sm = (read_u32(bytes, 48)? >> 8) & 0xff; + if actual_sm != expected_sm { + return Err(format!( + "cubin ELF flags target sm_{actual_sm}, expected {target}" + )); + } + let section_offset = usize::try_from(read_u64(bytes, 40)?) + .map_err(|_| "section table offset does not fit in memory".to_owned())?; + let section_entry_size = usize::from(read_u16(bytes, 58)?); + let section_count = usize::from(read_u16(bytes, 60)?); + let names_index = usize::from(read_u16(bytes, 62)?); + if section_entry_size < SECTION_HEADER_SIZE + || section_count == 0 + || names_index >= section_count + { + return Err("invalid or unsupported cubin section table".to_owned()); + } + let section_table_end = section_entry_size + .checked_mul(section_count) + .and_then(|size| section_offset.checked_add(size)) + .ok_or_else(|| "cubin section table overflows".to_owned())?; + if section_table_end > bytes.len() { + return Err("cubin section table is out of bounds".to_owned()); + } + let names_header = section_offset + names_index * section_entry_size; + let names_offset = usize::try_from(read_u64(bytes, names_header + 24)?) + .map_err(|_| "section-name table offset does not fit in memory".to_owned())?; + let names_size = usize::try_from(read_u64(bytes, names_header + 32)?) + .map_err(|_| "section-name table size does not fit in memory".to_owned())?; + let names_end = names_offset + .checked_add(names_size) + .ok_or_else(|| "section-name table overflows".to_owned())?; + let names = bytes + .get(names_offset..names_end) + .ok_or_else(|| "cubin section-name table is out of bounds".to_owned())?; + + let mut text_sections = BTreeMap::::new(); + let mut parameter_sections = BTreeMap::::new(); + let mut target_note = None; + for index in 0..section_count { + let header = section_offset + index * section_entry_size; + let name_offset = usize::try_from(read_u32(bytes, header)?) + .map_err(|_| "section name offset does not fit in memory".to_owned())?; + let name = nul_terminated(names, name_offset)?; + if let Some(kernel) = name.strip_prefix(".text.") { + let body_offset = usize::try_from(read_u64(bytes, header + 24)?) + .map_err(|_| format!("section {name:?} offset does not fit in memory"))?; + let size = usize::try_from(read_u64(bytes, header + 32)?) + .map_err(|_| format!("section {name:?} size does not fit in memory"))?; + let body_end = body_offset + .checked_add(size) + .ok_or_else(|| format!("section {name:?} overflows"))?; + if body_end > bytes.len() { + return Err(format!("section {name:?} is out of bounds")); + } + text_sections.insert(kernel.to_owned(), size); + } else if let Some(kernel) = name.strip_prefix(".nv.info.") { + if !expected_abis.contains_key(kernel) { + continue; + } + if read_u32(bytes, header + 4)? != 0x7000_0000 { + return Err(format!("section {name:?} is not CUDA_INFO")); + } + let body_offset = usize::try_from(read_u64(bytes, header + 24)?) + .map_err(|_| format!("section {name:?} offset does not fit in memory"))?; + let size = usize::try_from(read_u64(bytes, header + 32)?) + .map_err(|_| format!("section {name:?} size does not fit in memory"))?; + let body_end = body_offset + .checked_add(size) + .ok_or_else(|| format!("section {name:?} overflows"))?; + let body = bytes + .get(body_offset..body_end) + .ok_or_else(|| format!("section {name:?} is out of bounds"))?; + if parameter_sections.insert(kernel.to_owned(), body).is_some() { + return Err(format!("cubin contains multiple {name:?} sections")); + } + } else if name == ".note.nv.tkinfo" { + if target_note.is_some() { + return Err("cubin contains multiple .note.nv.tkinfo sections".to_owned()); + } + if read_u32(bytes, header + 4)? != 7 { + return Err(".note.nv.tkinfo is not an ELF note section".to_owned()); + } + if read_u64(bytes, header + 48)? != 4 { + return Err(".note.nv.tkinfo does not have 4-byte alignment".to_owned()); + } + let body_offset = usize::try_from(read_u64(bytes, header + 24)?) + .map_err(|_| ".note.nv.tkinfo offset does not fit in memory".to_owned())?; + let size = usize::try_from(read_u64(bytes, header + 32)?) + .map_err(|_| ".note.nv.tkinfo size does not fit in memory".to_owned())?; + let body_end = body_offset + .checked_add(size) + .ok_or_else(|| ".note.nv.tkinfo range overflows".to_owned())?; + target_note = Some( + bytes + .get(body_offset..body_end) + .ok_or_else(|| ".note.nv.tkinfo is out of bounds".to_owned())?, + ); + } + } + validate_target_note( + target_note.ok_or_else(|| "cubin is missing .note.nv.tkinfo".to_owned())?, + target, + )?; + for (kernel, expected_parameters) in expected_abis { + match text_sections.get(kernel) { + Some(size) if *size != 0 => {} + Some(_) => return Err(format!("kernel section .text.{kernel} is empty")), + None => return Err(format!("missing kernel code section .text.{kernel}")), + } + let parameter_section = parameter_sections + .get(kernel) + .ok_or_else(|| format!("missing kernel parameter section .nv.info.{kernel}"))?; + validate_kernel_parameter_abi(kernel, parameter_section, expected_parameters)?; + } + Ok(()) +} + +/// Check each cubin's launch argument layout against the exported MLIR. +/// +/// CUDA records parameter sizes and offsets in each kernel's `.nv.info` section. +/// We support the record format emitted by the pinned CUTLASS 4.7 / CUDA 13.x +/// toolchain and reject unknown or malformed records. +fn validate_kernel_parameter_abi( + kernel: &str, + section: &[u8], + expected: &[KernelParameterKind], +) -> Result<(), String> { + const EIFMT_NVAL: u8 = 1; + const EIFMT_BVAL: u8 = 2; + const EIFMT_HVAL: u8 = 3; + const EIFMT_SVAL: u8 = 4; + const EIATTR_KPARAM_INFO: u8 = 0x17; + const EIATTR_CBANK_PARAM_SIZE: u8 = 0x19; + const KPARAM_PAYLOAD_SIZE: u16 = 12; + + #[derive(Clone, Copy, Debug)] + struct ActualParameter { + offset: u16, + size: u16, + space: u32, + } + + let mut parameters = BTreeMap::::new(); + let mut parameter_block_size = None; + let mut cursor = 0; + while cursor < section.len() { + let header = section + .get(cursor..cursor + 4) + .ok_or_else(|| format!(".nv.info.{kernel} has a truncated attribute header"))?; + let format = header[0]; + let attribute = header[1]; + let immediate = u16::from_le_bytes([header[2], header[3]]); + let payload_size = match format { + EIFMT_NVAL | EIFMT_BVAL | EIFMT_HVAL => 0, + EIFMT_SVAL => usize::from(immediate), + _ => { + return Err(format!( + ".nv.info.{kernel} uses unsupported attribute format {format:#x}" + )); + } + }; + let payload_start = cursor + 4; + let next = payload_start + .checked_add(payload_size) + .ok_or_else(|| format!(".nv.info.{kernel} attribute range overflows"))?; + let payload = section + .get(payload_start..next) + .ok_or_else(|| format!(".nv.info.{kernel} has a truncated attribute payload"))?; + + match attribute { + EIATTR_KPARAM_INFO => { + if format != EIFMT_SVAL || immediate != KPARAM_PAYLOAD_SIZE { + return Err(format!( + ".nv.info.{kernel} has a malformed EIATTR_KPARAM_INFO record" + )); + } + let index = read_u32(payload, 0)?; + let ordinal = read_u16(payload, 4)?; + let offset = read_u16(payload, 6)?; + let flags = read_u32(payload, 8)?; + if index != 0 { + return Err(format!( + "kernel {kernel:?} parameter ordinal {ordinal} has unsupported index {index}" + )); + } + let size = u16::try_from((flags >> 18) & 0x3fff) + .expect("14-bit CUDA parameter size fits u16"); + let space = (flags >> 8) & 0xf; + if parameters + .insert( + ordinal, + ActualParameter { + offset, + size, + space, + }, + ) + .is_some() + { + return Err(format!( + "kernel {kernel:?} has duplicate parameter ordinal {ordinal}" + )); + } + } + EIATTR_CBANK_PARAM_SIZE => { + if format != EIFMT_HVAL { + return Err(format!( + ".nv.info.{kernel} has a malformed EIATTR_CBANK_PARAM_SIZE record" + )); + } + if parameter_block_size.replace(immediate).is_some() { + return Err(format!( + "kernel {kernel:?} has multiple parameter-block size records" + )); + } + } + _ => {} + } + cursor = next; + } + + if parameters.len() != expected.len() { + return Err(format!( + "kernel {kernel:?} has {} CUDA parameters, expected {} from its flattened MLIR signature", + parameters.len(), + expected.len() + )); + } + let mut expected_offset = 0_u16; + for (index, expected_kind) in expected.iter().copied().enumerate() { + let ordinal = u16::try_from(index) + .map_err(|_| format!("kernel {kernel:?} has too many parameters"))?; + let alignment_mask = expected_kind.alignment() - 1; + expected_offset = expected_offset + .checked_add(alignment_mask) + .map(|offset| offset & !alignment_mask) + .ok_or_else(|| format!("kernel {kernel:?} parameter alignment overflows"))?; + let actual = parameters + .get(&ordinal) + .ok_or_else(|| format!("kernel {kernel:?} is missing parameter ordinal {ordinal}"))?; + if actual.offset != expected_offset + || actual.size != expected_kind.size() + || actual.space != expected_kind.cuda_parameter_space() + { + return Err(format!( + "kernel {kernel:?} parameter {ordinal} is offset/size/space {}/{}/{}, expected {}/{}/{}", + actual.offset, + actual.size, + actual.space, + expected_offset, + expected_kind.size(), + expected_kind.cuda_parameter_space() + )); + } + expected_offset = expected_offset + .checked_add(expected_kind.size()) + .ok_or_else(|| format!("kernel {kernel:?} parameter block overflows"))?; + } + match parameter_block_size { + Some(size) if size == expected_offset => Ok(()), + // ptxas omits parameter-bank metadata for a zero-argument entry. + // All records were still parsed and the parameter count checked above. + None if expected.is_empty() => Ok(()), + Some(size) => Err(format!( + "kernel {kernel:?} parameter block is {size} bytes, expected {expected_offset}" + )), + None => Err(format!( + "kernel {kernel:?} is missing EIATTR_CBANK_PARAM_SIZE" + )), + } +} + +fn validate_target_note(section: &[u8], target: &str) -> Result<(), String> { + const NVIDIA_NOTE_TYPE: u32 = 0x7d0; + const NVIDIA_OWNER: &[u8] = b"NVIDIA Corp\0"; + + if section.len() < 12 { + return Err(".note.nv.tkinfo is truncated".to_owned()); + } + let name_size = usize::try_from(read_u32(section, 0)?) + .map_err(|_| "NVIDIA note owner size does not fit in memory".to_owned())?; + let descriptor_size = usize::try_from(read_u32(section, 4)?) + .map_err(|_| "NVIDIA note descriptor size does not fit in memory".to_owned())?; + if read_u32(section, 8)? != NVIDIA_NOTE_TYPE { + return Err(".note.nv.tkinfo has an unexpected NVIDIA note type".to_owned()); + } + let name_start = 12_usize; + let name_end = name_start + .checked_add(name_size) + .ok_or_else(|| "NVIDIA note owner range overflows".to_owned())?; + let descriptor_start = align_up_4(name_end)?; + let descriptor_end = descriptor_start + .checked_add(descriptor_size) + .ok_or_else(|| "NVIDIA note descriptor range overflows".to_owned())?; + let note_end = align_up_4(descriptor_end)?; + if note_end != section.len() { + return Err(".note.nv.tkinfo must contain exactly one bounded NVIDIA note".to_owned()); + } + let owner = section + .get(name_start..name_end) + .ok_or_else(|| "NVIDIA note owner is out of bounds".to_owned())?; + if owner != NVIDIA_OWNER { + return Err(".note.nv.tkinfo has an unexpected owner".to_owned()); + } + let descriptor = section + .get(descriptor_start..descriptor_end) + .ok_or_else(|| "NVIDIA note descriptor is out of bounds".to_owned())?; + let tokens = descriptor + .split(|byte| *byte == 0 || byte.is_ascii_whitespace()) + .filter(|token| !token.is_empty()) + .collect::>(); + let mut arch_values = tokens + .windows(2) + .filter(|window| window[0] == b"-arch") + .map(|window| window[1]); + let arch = arch_values + .next() + .ok_or_else(|| ".note.nv.tkinfo has no bounded -arch argument".to_owned())?; + if arch_values.next().is_some() { + return Err(".note.nv.tkinfo has multiple -arch arguments".to_owned()); + } + if arch != target.as_bytes() { + return Err(format!( + ".note.nv.tkinfo targets {}, expected {target}", + String::from_utf8_lossy(arch) + )); + } + Ok(()) +} + +fn align_up_4(value: usize) -> Result { + value + .checked_add(3) + .map(|value| value & !3) + .ok_or_else(|| "ELF note alignment overflows".to_owned()) +} + +fn parse_sm_target(target: &str) -> Result { + parse_sm_target_value(target).map_err(backend_error) +} + +fn parse_sm_target_value(target: &str) -> Result { + let digits = target + .strip_prefix("sm_") + .and_then(|value| value.strip_suffix('a')) + .ok_or_else(|| format!("unsupported CUTLASS device target {target:?}"))?; + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(format!("unsupported CUTLASS device target {target:?}")); + } + digits + .parse::() + .map_err(|_| format!("CUTLASS device target {target:?} is out of range")) +} + +fn publish_atomically(path: &Path, bytes: &[u8]) -> Result<(), PipelineError> { + let parent = path + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or(Path::new(".")); + let mut temporary = tempfile::NamedTempFile::new_in(parent).map_err(|error| { + backend_error(format!( + "could not create a temporary artifact next to {}: {error}", + path.display() + )) + })?; + temporary.write_all(bytes).map_err(|error| { + backend_error(format!( + "could not stage artifact {}: {error}", + path.display() + )) + })?; + temporary.as_file().sync_all().map_err(|error| { + backend_error(format!( + "could not sync artifact {}: {error}", + path.display() + )) + })?; + temporary.persist(path).map_err(|error| { + backend_error(format!( + "could not publish artifact {}: {}", + path.display(), + error.error + )) + })?; + Ok(()) +} + +fn read_u16(bytes: &[u8], offset: usize) -> Result { + let value: [u8; 2] = bytes + .get(offset..offset + 2) + .ok_or_else(|| "truncated binary field".to_owned())? + .try_into() + .expect("slice length checked"); + Ok(u16::from_le_bytes(value)) +} + +fn read_u32(bytes: &[u8], offset: usize) -> Result { + let value: [u8; 4] = bytes + .get(offset..offset + 4) + .ok_or_else(|| "truncated binary field".to_owned())? + .try_into() + .expect("slice length checked"); + Ok(u32::from_le_bytes(value)) +} + +fn read_u64(bytes: &[u8], offset: usize) -> Result { + let value: [u8; 8] = bytes + .get(offset..offset + 8) + .ok_or_else(|| "truncated binary field".to_owned())? + .try_into() + .expect("slice length checked"); + Ok(u64::from_le_bytes(value)) +} + +fn nul_terminated(bytes: &[u8], offset: usize) -> Result<&str, String> { + let tail = bytes + .get(offset..) + .ok_or_else(|| "section name offset is out of bounds".to_owned())?; + let end = tail + .iter() + .position(|byte| *byte == 0) + .ok_or_else(|| "unterminated ELF section name".to_owned())?; + std::str::from_utf8(&tail[..end]).map_err(|_| "ELF section name is not UTF-8".to_owned()) +} + +fn backend_error(message: impl Into) -> PipelineError { + PipelineError::BackendB(message.into()) +} + +#[cfg(test)] +mod tests { + use super::*; + use object::write::{Object as WriteObject, Symbol, SymbolSection}; + use object::{ + Architecture, BinaryFormat, Endianness, SectionKind, SymbolFlags, SymbolKind, SymbolScope, + }; + + #[test] + fn optimization_level_is_validated_and_changes_pipeline_provenance() { + let mut config = CutlassBackendConfig::new(PathBuf::from("/compiler.so")); + let default = precompiled_pipeline(&config).unwrap(); + assert!(default.contains("opt-level=3")); + config.optimization_level = 2; + let matched = precompiled_pipeline(&config).unwrap(); + assert!(matched.contains("opt-level=2")); + assert_ne!(Sha256::digest(default), Sha256::digest(matched)); + config.optimization_level = 4; + assert!(precompiled_pipeline(&config).is_err()); + } + + fn pair_abi() -> Vec { + vec![ + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + ] + } + + fn abi_map(kernel: &str, parameters: Vec) -> ExpectedKernelAbis { + BTreeMap::from([(kernel.to_owned(), parameters)]) + } + + fn parameter_info(parameters: &[KernelParameterKind]) -> Vec { + let mut info = Vec::new(); + let mut offset = 0_u16; + let mut offsets = Vec::new(); + for parameter in parameters { + let mask = parameter.alignment() - 1; + offset = (offset + mask) & !mask; + offsets.push(offset); + offset += parameter.size(); + } + for (ordinal, parameter) in parameters.iter().copied().enumerate().rev() { + info.extend_from_slice(&[4, 0x17]); + info.extend_from_slice(&12_u16.to_le_bytes()); + info.extend_from_slice(&0_u32.to_le_bytes()); + info.extend_from_slice(&(ordinal as u16).to_le_bytes()); + info.extend_from_slice(&offsets[ordinal].to_le_bytes()); + let flags = (u32::from(parameter.size()) << 18) + | (0x1f << 12) + | (parameter.cuda_parameter_space() << 8); + info.extend_from_slice(&flags.to_le_bytes()); + } + info.extend_from_slice(&[3, 0x19]); + info.extend_from_slice(&offset.to_le_bytes()); + info + } + + fn minimal_cubin_with_arches( + kernel: &str, + sm: u32, + arches: &[&str], + parameters: &[KernelParameterKind], + ) -> Vec { + let names = format!("\0.shstrtab\0.text.{kernel}\0.note.nv.tkinfo\0.nv.info.{kernel}\0") + .into_bytes(); + let text_name = format!(".text.{kernel}"); + let info_name = format!(".nv.info.{kernel}"); + let text_name_offset = names + .windows(text_name.len()) + .position(|window| window == text_name.as_bytes()) + .unwrap(); + let note_name_offset = names + .windows(b".note.nv.tkinfo".len()) + .position(|window| window == b".note.nv.tkinfo") + .unwrap(); + let info_name_offset = names + .windows(info_name.len()) + .position(|window| window == info_name.as_bytes()) + .unwrap(); + let mut descriptor = Vec::new(); + descriptor.extend_from_slice(b"-O\0"); + for arch in arches { + descriptor.extend_from_slice(b"-arch\0"); + descriptor.extend_from_slice(arch.as_bytes()); + descriptor.push(0); + } + let mut note = Vec::new(); + note.extend_from_slice(&12_u32.to_le_bytes()); + note.extend_from_slice(&(descriptor.len() as u32).to_le_bytes()); + note.extend_from_slice(&0x7d0_u32.to_le_bytes()); + note.extend_from_slice(b"NVIDIA Corp\0"); + note.extend_from_slice(&descriptor); + while note.len() % 4 != 0 { + note.push(0); + } + let info = parameter_info(parameters); + + let names_offset = 64_usize; + let code_offset = (names_offset + names.len() + 7) & !7; + let note_offset = (code_offset + 8 + 3) & !3; + let info_offset = (note_offset + note.len() + 3) & !3; + let section_offset = (info_offset + info.len() + 7) & !7; + let mut bytes = vec![0_u8; section_offset + 5 * 64]; + bytes[..4].copy_from_slice(b"\x7fELF"); + bytes[4] = 2; + bytes[5] = 1; + bytes[6] = 1; + bytes[16..18].copy_from_slice(&2_u16.to_le_bytes()); + bytes[18..20].copy_from_slice(&190_u16.to_le_bytes()); + bytes[40..48].copy_from_slice(&(section_offset as u64).to_le_bytes()); + bytes[48..52].copy_from_slice(&(sm << 8).to_le_bytes()); + bytes[52..54].copy_from_slice(&64_u16.to_le_bytes()); + bytes[58..60].copy_from_slice(&64_u16.to_le_bytes()); + bytes[60..62].copy_from_slice(&5_u16.to_le_bytes()); + bytes[62..64].copy_from_slice(&1_u16.to_le_bytes()); + bytes[names_offset..names_offset + names.len()].copy_from_slice(&names); + bytes[code_offset..code_offset + 8].copy_from_slice(&[1; 8]); + bytes[note_offset..note_offset + note.len()].copy_from_slice(¬e); + bytes[info_offset..info_offset + info.len()].copy_from_slice(&info); + + let names_header = section_offset + 64; + bytes[names_header..names_header + 4].copy_from_slice(&1_u32.to_le_bytes()); + bytes[names_header + 4..names_header + 8].copy_from_slice(&3_u32.to_le_bytes()); + bytes[names_header + 24..names_header + 32] + .copy_from_slice(&(names_offset as u64).to_le_bytes()); + bytes[names_header + 32..names_header + 40] + .copy_from_slice(&(names.len() as u64).to_le_bytes()); + + let text_header = section_offset + 128; + bytes[text_header..text_header + 4] + .copy_from_slice(&(text_name_offset as u32).to_le_bytes()); + bytes[text_header + 4..text_header + 8].copy_from_slice(&1_u32.to_le_bytes()); + bytes[text_header + 24..text_header + 32] + .copy_from_slice(&(code_offset as u64).to_le_bytes()); + bytes[text_header + 32..text_header + 40].copy_from_slice(&8_u64.to_le_bytes()); + + let note_header = section_offset + 192; + bytes[note_header..note_header + 4] + .copy_from_slice(&(note_name_offset as u32).to_le_bytes()); + bytes[note_header + 4..note_header + 8].copy_from_slice(&7_u32.to_le_bytes()); + bytes[note_header + 24..note_header + 32] + .copy_from_slice(&(note_offset as u64).to_le_bytes()); + bytes[note_header + 32..note_header + 40] + .copy_from_slice(&(note.len() as u64).to_le_bytes()); + bytes[note_header + 48..note_header + 56].copy_from_slice(&4_u64.to_le_bytes()); + + let info_header = section_offset + 256; + bytes[info_header..info_header + 4] + .copy_from_slice(&(info_name_offset as u32).to_le_bytes()); + bytes[info_header + 4..info_header + 8].copy_from_slice(&0x7000_0000_u32.to_le_bytes()); + bytes[info_header + 24..info_header + 32] + .copy_from_slice(&(info_offset as u64).to_le_bytes()); + bytes[info_header + 32..info_header + 40] + .copy_from_slice(&(info.len() as u64).to_le_bytes()); + bytes[info_header + 48..info_header + 56].copy_from_slice(&4_u64.to_le_bytes()); + bytes + } + + fn minimal_cubin(kernel: &str, sm: u32) -> Vec { + minimal_cubin_with_arches(kernel, sm, &["sm_120a"], &pair_abi()) + } + + fn fatbin(cubin: &[u8], sm: u32) -> Vec { + const HEADER: usize = 16; + const CODE_OFFSET: usize = 64; + let payload_size = CODE_OFFSET + cubin.len(); + let mut bytes = vec![0_u8; HEADER + payload_size]; + bytes[..4].copy_from_slice(&0xba55_ed50_u32.to_le_bytes()); + bytes[4..6].copy_from_slice(&1_u16.to_le_bytes()); + bytes[6..8].copy_from_slice(&(HEADER as u16).to_le_bytes()); + bytes[8..16].copy_from_slice(&(payload_size as u64).to_le_bytes()); + bytes[HEADER..HEADER + 2].copy_from_slice(&2_u16.to_le_bytes()); + bytes[HEADER + 2..HEADER + 4].copy_from_slice(&0x0101_u16.to_le_bytes()); + bytes[HEADER + 4..HEADER + 8].copy_from_slice(&(CODE_OFFSET as u32).to_le_bytes()); + bytes[HEADER + 8..HEADER + 16].copy_from_slice(&(cubin.len() as u64).to_le_bytes()); + bytes[HEADER + 28..HEADER + 32].copy_from_slice(&sm.to_le_bytes()); + bytes[HEADER + CODE_OFFSET..].copy_from_slice(cubin); + bytes + } + + fn host_object(image: &[u8]) -> Vec { + let mut object = + WriteObject::new(BinaryFormat::Elf, Architecture::X86_64, Endianness::Little); + let section = + object.add_section(Vec::new(), b".lrodata".to_vec(), SectionKind::ReadOnlyData); + object.section_mut(section).set_data(image.to_vec(), 8); + object.add_symbol(Symbol { + name: b"kernels_binary".to_vec(), + value: 0, + size: image.len() as u64, + kind: SymbolKind::Data, + scope: SymbolScope::Compilation, + weak: false, + section: SymbolSection::Section(section), + flags: SymbolFlags::None, + }); + object.write().unwrap() + } + + #[test] + fn object_symbol_and_fatbinary_are_validated_end_to_end() { + let cubin = minimal_cubin("add_f32", 120); + let image = fatbin(&cubin, 120); + let object = host_object(&image); + let extracted = extract_kernels_binary(&object).unwrap(); + assert_eq!(extracted, image); + let (normalized, source_kind) = + extract_validated_cubin(&extracted, "sm_120a", &abi_map("add_f32", pair_abi())) + .unwrap(); + assert_eq!(normalized, cubin); + assert_eq!(source_kind, "a CUDA fatbinary"); + } + + #[test] + fn validation_rejects_wrong_target_missing_kernel_and_compression() { + let cubin = minimal_cubin("add_f32", 120); + let mut image = fatbin(&cubin, 120); + assert!( + extract_validated_cubin(&image, "sm_100a", &abi_map("add_f32", pair_abi())).is_err() + ); + assert!( + extract_validated_cubin(&image, "sm_120a", &abi_map("add_f16", pair_abi())).is_err() + ); + image[16 + 40..16 + 48].copy_from_slice(&0x1000_u64.to_le_bytes()); + assert!( + extract_validated_cubin(&image, "sm_120a", &abi_map("add_f32", pair_abi())).is_err() + ); + } + + #[test] + fn validation_requires_the_exact_blackwell_target_variant_note() { + let wrong_variant = minimal_cubin_with_arches("add_f32", 120, &["sm_120f"], &pair_abi()); + let error = validate_cuda_cubin(&wrong_variant, "sm_120a", &abi_map("add_f32", pair_abi())) + .unwrap_err(); + assert!(error.contains("sm_120f"), "{error}"); + + let duplicate = + minimal_cubin_with_arches("add_f32", 120, &["sm_120a", "sm_120a"], &pair_abi()); + let error = validate_cuda_cubin(&duplicate, "sm_120a", &abi_map("add_f32", pair_abi())) + .unwrap_err(); + assert!(error.contains("multiple -arch"), "{error}"); + + let mut missing = minimal_cubin("add_f32", 120); + let note_name = missing + .windows(b".note.nv.tkinfo".len()) + .position(|window| window == b".note.nv.tkinfo") + .unwrap(); + missing[note_name] = b'_'; + let error = + validate_cuda_cubin(&missing, "sm_120a", &abi_map("add_f32", pair_abi())).unwrap_err(); + assert!(error.contains("missing .note.nv.tkinfo"), "{error}"); + } + + #[test] + fn flattened_elementwise_and_gemv_signatures_have_fixed_u64_slots() { + let elementwise = r#" + "func.func"() <{function_type = (!llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i64) -> (), sym_name = "add_f32"}> ({ + "func.func"() <{function_type = (!llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i64) -> (), sym_name = "add_f16"}> ({ + "#; + let elementwise_abis = + expected_kernel_abis(elementwise, &["add_f32".to_owned(), "add_f16".to_owned()]) + .unwrap(); + let elementwise_parameters = vec![ + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + ]; + assert_eq!(elementwise_abis["add_f32"], elementwise_parameters); + assert_eq!(elementwise_abis["add_f16"], elementwise_parameters); + + let gemv = r#" + "func.func"() <{function_type = (!llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i64, i64, i64) -> (), sym_name = "nvfp4_gemv"}> ({ + "#; + let gemv_abis = expected_kernel_abis(gemv, &["nvfp4_gemv".to_owned()]).unwrap(); + assert_eq!(gemv_abis["nvfp4_gemv"].len(), 12); + assert_eq!( + gemv_abis["nvfp4_gemv"], + vec![ + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::Scalar64, + KernelParameterKind::Scalar64, + ] + ); + } + + #[test] + fn cubin_parameter_metadata_matches_elementwise_and_gemv_abis() { + let elementwise = vec![ + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + ]; + let cubin = minimal_cubin_with_arches("add_f32", 120, &["sm_120a"], &elementwise); + validate_cuda_cubin(&cubin, "sm_120a", &abi_map("add_f32", elementwise.clone())).unwrap(); + + let gemv = vec![ + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + KernelParameterKind::Scalar64, + KernelParameterKind::Scalar64, + ]; + let cubin = minimal_cubin_with_arches("nvfp4_gemv", 120, &["sm_120a"], &gemv); + validate_cuda_cubin(&cubin, "sm_120a", &abi_map("nvfp4_gemv", gemv.clone())).unwrap(); + + let mut malformed = parameter_info(&elementwise); + // The first record describes ordinal 5 at offset 40. Make that slot + // overlap ordinal 0; the post-object ABI validator must fail closed. + malformed[10..12].copy_from_slice(&0_u16.to_le_bytes()); + let error = validate_kernel_parameter_abi("add_f32", &malformed, &elementwise).unwrap_err(); + assert!(error.contains("offset/size/space"), "{error}"); + } + + const MIXED_SCALAR_SIGNATURE: &str = r#""func.func"() <{function_type = (!llvm.ptr, i1, i32, i1, i64, i1) -> (), sym_name = "mixed_abi"}> ({"#; + + const FLOAT_SCALAR_SIGNATURE: &str = r#""func.func"() <{function_type = (!llvm.ptr, i1, f32, i1, i64, f32) -> (), sym_name = "float_abi"}> ({"#; + + #[test] + fn flattened_f32_uses_four_byte_scalar_slots_and_preserves_padding_checks() { + let abis = expected_kernel_abis(FLOAT_SCALAR_SIGNATURE, &["float_abi".into()]).unwrap(); + let parameters = &abis["float_abi"]; + assert_eq!(parameters[2].size() as usize, std::mem::size_of::()); + assert_eq!( + parameters[2].alignment() as usize, + std::mem::align_of::() + ); + assert_eq!(parameters[2].cuda_parameter_space(), 0); + assert_eq!(parameters[5], parameters[2]); + + // Explicit CUDA KPARAM positions are independent of the production + // alignment calculation. The pinned compiler fixture below exercises + // this same mixed-width signature through object generation. + let layout: [(u16, u16, u16, u32); 6] = [ + (0, 0, 8, 5), + (1, 8, 1, 0), + (2, 12, 4, 0), + (3, 16, 1, 0), + (4, 24, 8, 0), + (5, 32, 4, 0), + ]; + let mut records = vec![]; + for (ordinal, offset, size, space) in layout { + records.extend_from_slice(&[4, 0x17, 12, 0]); + records.extend_from_slice(&0_u32.to_le_bytes()); + records.extend_from_slice(&ordinal.to_le_bytes()); + records.extend_from_slice(&offset.to_le_bytes()); + records.extend_from_slice( + &((u32::from(size) << 18) | (0x1f << 12) | (space << 8)).to_le_bytes(), + ); + } + records.extend_from_slice(&[3, 0x19, 36, 0]); + validate_kernel_parameter_abi("float_abi", &records, parameters).unwrap(); + for (offset, bytes) in [ + (42, 9_u16.to_le_bytes().to_vec()), + (44, ((8_u32 << 18) | (0x1f << 12)).to_le_bytes().to_vec()), + ( + 44, + ((4_u32 << 18) | (0x1f << 12) | (5 << 8)) + .to_le_bytes() + .to_vec(), + ), + (98, 40_u16.to_le_bytes().to_vec()), + ] { + let mut malformed = records.clone(); + malformed[offset..offset + bytes.len()].copy_from_slice(&bytes); + assert!(validate_kernel_parameter_abi("float_abi", &malformed, parameters).is_err()); + } + } + + #[test] + fn float_argument_support_does_not_admit_unreviewed_types_or_attributes() { + for ty in ["f16", "bf16", "f64", "vector<2xf32>", "!llvm.struct<(f32)>"] { + let signature = FLOAT_SCALAR_SIGNATURE.replace("f32", ty); + assert!( + expected_kernel_abis(&signature, &["float_abi".into()]).is_err(), + "{ty}" + ); + } + assert!(parse_parameter_abi("f32", Some("{llvm.align = 4 : i64}")).is_err()); + } + + #[test] + #[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] + fn cutlass_f32_direct_launch_abi_compiles() { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mlir = format!( + r#"module attributes {{gpu.container_module}} {{ + gpu.module @kernels {{ + {FLOAT_SCALAR_SIGNATURE} + ^bb0(%out: !llvm.ptr, %b0: i1, %x: f32, %b1: i1, %wide: i64, %epsilon: f32): + %w = arith.uitofp %wide : i64 to f32 + %sum = arith.addf %x, %epsilon : f32 + %y = arith.select %b0, %sum, %x : f32 + %z = arith.select %b1, %y, %w : f32 + llvm.store %z, %out : f32, !llvm.ptr + func.return + }}) {{cute.kernel, gpu.kernel}} : () -> () + }} +}}"# + ); + let temp = tempfile::tempdir().unwrap(); + let output = temp.path().join("float_abi.cubin"); + compile_to_cubin( + &CutlassBackendConfig::new(compiler_path.into()), + &mlir, + "sm_120a", + &["float_abi".into()], + &output, + None, + ) + .unwrap(); + assert!(std::fs::metadata(output).unwrap().len() > 0); + } + + #[test] + fn flattened_bool_and_i32_match_pinned_cuda_parameter_records() { + let abis = expected_kernel_abis(MIXED_SCALAR_SIGNATURE, &["mixed_abi".into()]).unwrap(); + let parameters = &abis["mixed_abi"]; + // CUTLASS 4.7 / CUDA 13 emits i1 as one byte, matching the original + // Rust bool carrier passed to push_kernel_scalar by typed launchers. + assert_eq!( + parameters + .iter() + .map(|kind| kind.size()) + .collect::>(), + [8, 1, 4, 1, 8, 1] + ); + assert_eq!(parameters[1].size() as usize, std::mem::size_of::()); + assert_eq!( + parameters[1].alignment() as usize, + std::mem::align_of::() + ); + assert_eq!(parameters[2].size() as usize, std::mem::size_of::()); + assert_eq!( + parameters[2].alignment() as usize, + std::mem::align_of::() + ); + + fn records(layout: &[(u16, u16, u16, u32)], total: u16) -> Vec { + let mut info = vec![]; + for &(ordinal, offset, size, space) in layout.iter().rev() { + info.extend_from_slice(&[4, 0x17, 12, 0]); + info.extend_from_slice(&0_u32.to_le_bytes()); + info.extend_from_slice(&ordinal.to_le_bytes()); + info.extend_from_slice(&offset.to_le_bytes()); + info.extend_from_slice( + &((u32::from(size) << 18) | (0x1f << 12) | (space << 8)).to_le_bytes(), + ); + } + info.extend_from_slice(&[3, 0x19]); + info.extend_from_slice(&total.to_le_bytes()); + info + } + // Explicit offsets come from the pinned compiler's .nv.info output, + // independently of KernelParameterKind's alignment/size calculation. + let layout = [ + (0, 0, 8, 5), + (1, 8, 1, 0), + (2, 12, 4, 0), + (3, 16, 1, 0), + (4, 24, 8, 0), + (5, 32, 1, 0), + ]; + validate_kernel_parameter_abi("mixed_abi", &records(&layout, 33), parameters).unwrap(); + for (slot, invalid) in [ + (1, (1, 8, 4, 0)), + (2, (2, 9, 4, 0)), + (2, (2, 12, 4, 5)), + (4, (4, 17, 8, 0)), + ] { + let mut malformed = layout; + malformed[slot] = invalid; + assert!( + validate_kernel_parameter_abi("mixed_abi", &records(&malformed, 33), parameters) + .unwrap_err() + .contains("offset/size/space") + ); + } + assert!( + validate_kernel_parameter_abi("mixed_abi", &records(&layout, 40), parameters) + .unwrap_err() + .contains("parameter block") + ); + } + + #[test] + #[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] + fn cutlass_mixed_bool_i32_direct_launch_abi_compiles() { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mlir = format!( + r#"module attributes {{gpu.container_module}} {{ + gpu.module @kernels {{ + {MIXED_SCALAR_SIGNATURE} + ^bb0(%out: !llvm.ptr, %b0: i1, %n: i32, %b1: i1, %wide: i64, %b2: i1): + %w = arith.trunci %wide : i64 to i32 + %sum = arith.addi %n, %w : i32 + %x = arith.select %b0, %sum, %n : i32 + %y = arith.select %b1, %x, %w : i32 + %z = arith.select %b2, %y, %sum : i32 + llvm.store %z, %out : i32, !llvm.ptr + func.return + }}) {{cute.kernel, gpu.kernel}} : () -> () + }} +}}"# + ); + let temp = tempfile::tempdir().unwrap(); + let output = temp.path().join("mixed_abi.cubin"); + compile_to_cubin( + &CutlassBackendConfig::new(compiler_path.into()), + &mlir, + "sm_120a", + &["mixed_abi".into()], + &output, + None, + ) + .unwrap(); + assert!(std::fs::metadata(output).unwrap().len() > 0); + } + + #[test] + fn flattened_signature_parser_rejects_unsupported_or_ambiguous_abis() { + let unsupported = + r#""func.func"() <{function_type = (!llvm.ptr, i16) -> (), sym_name = "kernel"}> ({"#; + let error = expected_kernel_abis(unsupported, &["kernel".to_owned()]).unwrap_err(); + assert!(error.contains("unsupported flattened parameter"), "{error}"); + + let duplicate = format!("{unsupported}\n{unsupported}").replace("i16", "i64"); + let error = expected_kernel_abis(&duplicate, &["kernel".to_owned()]).unwrap_err(); + assert!(error.contains("found 2"), "{error}"); + } + + #[test] + fn profile_and_target_are_fail_closed() { + assert_eq!( + CutlassBackendConfig::new("/tmp/compiler.so".into()).profile, + CUTLASS_FULL_CUTE_MLIR22_PROFILE + ); + assert_eq!(parse_sm_target_value("sm_120a").unwrap(), 120); + assert!(parse_sm_target_value("sm_120").is_err()); + assert!(parse_sm_target_value("compute_120").is_err()); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_compiler.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_compiler.rs new file mode 100644 index 0000000000..fdc0ee02f6 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/cutlass_compiler.rs @@ -0,0 +1,895 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Runtime-loaded bindings for the CUTLASS 4.7 compiler C API. +//! +//! Compile textual `PreCompiledMlir` through a pinned CuTe pipeline to the +//! C API's `ObjectArtifact`. The result is a host ELF object containing the +//! CUDA image; the backend extracts that image before the driver can load it. +//! These wrappers release compiler and artifact handles with the matching +//! C API destructors. + +use libloading::{Library, Symbol}; +use sha2::{Digest, Sha256}; +use std::ffi::{c_char, c_int, c_void}; +use std::fmt; +use std::fs::{self, File}; +use std::io::Read; +#[cfg(target_os = "linux")] +use std::os::fd::AsRawFd; +use std::path::{Path, PathBuf}; +use std::ptr; +use std::slice; +use std::str::FromStr; +use thiserror::Error; + +const CUTLASS_ARTIFACT_PRE_COMPILED_MLIR: c_int = 1; +const CUTLASS_ARTIFACT_OBJECT: c_int = 4; +// CUTLASS 4.7's two concrete wrapper ABIs require FunctionMetadata, which the +// `gpu.module` external-module path deliberately does not carry. The +// fingerprint-pinned 4.7 library accepts TBD as a no-wrapper passthrough for +// this path. We assert below that the resulting ObjectArtifact still carries +// zero FunctionMetadata entries so a future change cannot silently introduce +// a host wrapper or reinterpret the direct CUDA kernel ABI. +const CUTLASS_ABI_TBD_EXTERNAL_MODULE_PASSTHROUGH: c_int = 0; + +type CompilerRef = *mut c_void; +type ArtifactsRef = *mut c_void; + +type CompilerCreateFn = unsafe extern "C" fn() -> CompilerRef; +type CompilerDestroyFn = unsafe extern "C" fn(CompilerRef); +type CompilerSetDeviceTargetFn = + unsafe extern "C" fn(CompilerRef, *const c_char, usize, *mut *mut c_char, *mut usize) -> c_int; +type CompilerSetAbiFn = + unsafe extern "C" fn(CompilerRef, c_int, *mut *mut c_char, *mut usize) -> c_int; +type CompilerSetPipelineFn = + unsafe extern "C" fn(CompilerRef, c_int, *const c_char, usize) -> c_int; +type CompilerCompileToFn = unsafe extern "C" fn( + CompilerRef, + ArtifactsRef, + c_int, + *mut *mut c_char, + *mut usize, +) -> ArtifactsRef; +type ArtifactsFromTextFn = unsafe extern "C" fn(c_int, *const c_char, usize) -> ArtifactsRef; +type ArtifactsGetTypeFn = unsafe extern "C" fn(ArtifactsRef) -> c_int; +type ArtifactsGetDataFn = unsafe extern "C" fn(ArtifactsRef, *mut *mut u8, *mut usize) -> c_int; +type ArtifactsFunctionCountFn = unsafe extern "C" fn(ArtifactsRef, *mut usize) -> c_int; +type ArtifactsDestroyFn = unsafe extern "C" fn(ArtifactsRef); + +unsafe extern "C" { + fn free(ptr: *mut c_void); +} + +/// A validated SHA-256 digest used to pin the proprietary compiler library. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct Sha256Fingerprint([u8; 32]); + +impl Sha256Fingerprint { + /// Return the raw 32-byte digest. + #[cfg(test)] + pub fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } +} + +impl fmt::Display for Sha256Fingerprint { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for byte in self.0 { + write!(f, "{byte:02x}")?; + } + Ok(()) + } +} + +impl FromStr for Sha256Fingerprint { + type Err = FingerprintParseError; + + fn from_str(value: &str) -> Result { + if value.len() != 64 { + return Err(FingerprintParseError::Length { + actual: value.len(), + }); + } + + let mut digest = [0_u8; 32]; + let encoded = value.as_bytes(); + for (index, byte) in digest.iter_mut().enumerate() { + let offset = index * 2; + let high = + hex_nibble(encoded[offset]).ok_or(FingerprintParseError::NonHex { offset })?; + let low = hex_nibble(encoded[offset + 1]) + .ok_or(FingerprintParseError::NonHex { offset: offset + 1 })?; + *byte = (high << 4) | low; + } + Ok(Self(digest)) + } +} + +fn hex_nibble(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } +} + +/// A malformed SHA-256 fingerprint supplied for compiler provenance. +#[derive(Clone, Debug, PartialEq, Eq, Error)] +pub enum FingerprintParseError { + /// A SHA-256 hexadecimal digest must contain exactly 64 characters. + #[error("SHA-256 fingerprint must contain 64 hexadecimal characters, got {actual}")] + Length { + /// Number of characters supplied by the caller. + actual: usize, + }, + /// The digest contains a character outside the hexadecimal alphabet. + #[error("SHA-256 fingerprint contains a non-hexadecimal byte at offset {offset}")] + NonHex { + /// Byte offset of the invalid two-character group. + offset: usize, + }, +} + +/// Errors produced while loading or calling the CUTLASS compiler C API. +#[derive(Debug, Error)] +pub enum CutlassCompilerError { + /// The expected compiler-library fingerprint is malformed. + #[error(transparent)] + InvalidFingerprint(#[from] FingerprintParseError), + + /// Runtime loading only accepts an explicit absolute path. + #[error("CUTLASS compiler library path must be absolute: `{path}`")] + RelativeLibraryPath { + /// Rejected path. + path: PathBuf, + }, + + /// The path does not name a regular file (symlinks are rejected). + #[error("CUTLASS compiler library path is not a regular file: `{path}`")] + LibraryNotRegular { + /// Rejected path. + path: PathBuf, + }, + + /// Filesystem inspection of the library failed. + #[error("could not inspect CUTLASS compiler library `{path}`: {source}")] + LibraryMetadata { + /// Inspected path. + path: PathBuf, + /// Underlying filesystem error. + #[source] + source: std::io::Error, + }, + + /// Opening the validated library file failed. + #[error("could not open CUTLASS compiler library `{path}`: {source}")] + LibraryOpen { + /// Opened path. + path: PathBuf, + /// Underlying filesystem error. + #[source] + source: std::io::Error, + }, + + /// Reading the library to calculate its digest failed. + #[error("could not fingerprint CUTLASS compiler library `{path}`: {source}")] + LibraryRead { + /// Read path. + path: PathBuf, + /// Underlying filesystem error. + #[source] + source: std::io::Error, + }, + + /// The library's actual digest did not match the pinned provenance. + #[error( + "CUTLASS compiler library fingerprint mismatch for `{path}`: expected {expected}, got {actual}" + )] + FingerprintMismatch { + /// Validated path. + path: PathBuf, + /// Pinned digest. + expected: Sha256Fingerprint, + /// Digest calculated from the opened file. + actual: Sha256Fingerprint, + }, + + /// `dlopen` rejected the validated library. + #[error("could not load CUTLASS compiler library `{path}`: {source}")] + LibraryLoad { + /// Loaded path. + path: PathBuf, + /// Dynamic-loader error. + #[source] + source: libloading::Error, + }, + + /// A required 4.7 C API symbol is absent. + #[error("CUTLASS compiler library is missing required symbol `{symbol}`: {source}")] + SymbolNotFound { + /// Required symbol name. + symbol: &'static str, + /// Dynamic-loader error. + #[source] + source: libloading::Error, + }, + + /// A C API constructor unexpectedly returned a null handle. + #[error("CUTLASS C API `{operation}` returned a null handle")] + NullHandle { + /// C function that returned null. + operation: &'static str, + }, + + /// Local input validation rejected the compile request. + #[error("invalid CUTLASS compile input: {reason}")] + InvalidInput { + /// Human-readable validation failure. + reason: &'static str, + }, + + /// Compiler configuration failed. + #[error("CUTLASS C API `{operation}` failed with status {status}: {message}")] + Configure { + /// C function that failed. + operation: &'static str, + /// Raw C status. + status: c_int, + /// C-allocated diagnostic, when available. + message: String, + }, + + /// The C API rejected a source-stage pass pipeline. + #[error( + "CUTLASS C API `cutlass_compiler_set_pipeline` rejected artifact type {artifact_type} with status {status}" + )] + PipelineConfiguration { + /// Source artifact type associated with the pipeline. + artifact_type: c_int, + /// Raw C status. + status: c_int, + }, + + /// The C API could not parse the supplied textual MLIR. + #[error("CUTLASS could not create a {artifact_kind} artifact from textual MLIR")] + MlirParse { + /// Human-readable C API artifact kind. + artifact_kind: &'static str, + }, + + /// The multi-stage compile operation failed. + #[error("CUTLASS C API `cutlass_compiler_compile_to` failed: {message}")] + Compile { + /// C-allocated diagnostic, when available. + message: String, + }, + + /// An artifact query failed or violated its documented pointer contract. + #[error("CUTLASS C API `{operation}` failed: {reason}")] + ArtifactQuery { + /// C function that failed. + operation: &'static str, + /// Failure detail. + reason: String, + }, + + /// The returned `ObjectArtifact` was not an ELF object as documented. + #[error("CUTLASS ObjectArtifact is not an ELF object")] + NotElfObject, + + /// The pinned external-module passthrough unexpectedly gained wrapper + /// metadata, so its direct CUDA-kernel ABI can no longer be assumed. + #[error( + "CUTLASS 4.7 external-module passthrough returned {count} FunctionMetadata entries; expected zero" + )] + UnexpectedFunctionMetadata { + /// Number of entries reported by the post-object C API. + count: usize, + }, +} + +/// Output of the official CompiledMlir-to-Object compilation path. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutlassObject { + /// Raw ELF bytes returned by `cutlass_artifacts_get_data`. + pub elf: Vec, +} + +/// A loaded and SHA-pinned CUTLASS 4.7 compiler library. +pub struct CutlassCompilerLibrary { + api: Api, +} + +impl CutlassCompilerLibrary { + /// Validate and load an official compiler library from an absolute path. + /// + /// The path's final component must be a regular file rather than a + /// symlink. Its SHA-256 digest must exactly match `expected_sha256`. On + /// Linux, `dlopen` is performed through the already-open file descriptor, + /// so replacing the pathname after validation cannot redirect the load to + /// a different inode. + pub fn load( + path: impl AsRef, + expected_sha256: &str, + ) -> Result { + let expected = expected_sha256.parse()?; + let validated = validate_library_file(path.as_ref(), expected)?; + let api = Api::load(validated)?; + Ok(Self { api }) + } + + /// Return the absolute path used for provenance validation. + pub fn path(&self) -> &Path { + &self.api.path + } + + /// Return the SHA-256 digest calculated before the library was loaded. + pub fn fingerprint(&self) -> Sha256Fingerprint { + self.api.fingerprint + } + + /// Compile textual `PreCompiledMlir` to a host ELF object. + /// + /// `precompiled_pipeline` replaces the library's first-stage CuTe pass + /// pipeline. The caller must pin it together with the compiler-library + /// digest: it is part of the executable compiler contract, not a user + /// supplied optimization knob. The returned object contains the CUDA + /// image selected by that pipeline; extracting and validating that image + /// is deliberately handled by the higher-level backend boundary. + pub fn compile_precompiled_mlir_to_object( + &self, + mlir: &str, + device_target: &str, + precompiled_pipeline: &str, + ) -> Result { + validate_compile_input(mlir, device_target)?; + if precompiled_pipeline.is_empty() { + return Err(CutlassCompilerError::InvalidInput { + reason: "PreCompiledMlir pipeline must not be empty", + }); + } + if precompiled_pipeline.as_bytes().contains(&0) { + return Err(CutlassCompilerError::InvalidInput { + reason: "PreCompiledMlir pipeline must not contain NUL bytes", + }); + } + + let compiler = CompilerHandle::create_cute(&self.api)?; + compiler.set_device_target(device_target)?; + compiler.set_external_module_passthrough_abi()?; + compiler.set_pipeline(CUTLASS_ARTIFACT_PRE_COMPILED_MLIR, precompiled_pipeline)?; + + let mut input = ArtifactsHandle::from_textual_precompiled_mlir(&self.api, mlir)?; + compile_object(compiler, &mut input) + } +} + +fn validate_compile_input(mlir: &str, device_target: &str) -> Result<(), CutlassCompilerError> { + if mlir.is_empty() { + return Err(CutlassCompilerError::InvalidInput { + reason: "textual MLIR must not be empty", + }); + } + if mlir.as_bytes().contains(&0) { + return Err(CutlassCompilerError::InvalidInput { + reason: "textual MLIR must not contain NUL bytes", + }); + } + if device_target.is_empty() { + return Err(CutlassCompilerError::InvalidInput { + reason: "device target must not be empty", + }); + } + if device_target.as_bytes().contains(&0) { + return Err(CutlassCompilerError::InvalidInput { + reason: "device target must not contain NUL bytes", + }); + } + Ok(()) +} + +fn compile_object<'api>( + compiler: CompilerHandle<'api>, + input: &mut ArtifactsHandle<'api>, +) -> Result { + let object = compiler.compile_to_object(input)?; + if object.artifact_type() != CUTLASS_ARTIFACT_OBJECT { + return Err(CutlassCompilerError::ArtifactQuery { + operation: "cutlass_artifacts_get_type", + reason: "compile returned an artifact whose type is not Object".to_string(), + }); + } + let function_metadata_count = object.function_count()?; + if function_metadata_count != 0 { + return Err(CutlassCompilerError::UnexpectedFunctionMetadata { + count: function_metadata_count, + }); + } + + let elf = object.data()?; + if !elf.starts_with(b"\x7fELF") { + return Err(CutlassCompilerError::NotElfObject); + } + Ok(CutlassObject { elf }) +} + +struct ValidatedLibraryFile { + path: PathBuf, + file: File, + fingerprint: Sha256Fingerprint, +} + +fn validate_library_file( + path: &Path, + expected: Sha256Fingerprint, +) -> Result { + if !path.is_absolute() { + return Err(CutlassCompilerError::RelativeLibraryPath { + path: path.to_path_buf(), + }); + } + + let metadata = + fs::symlink_metadata(path).map_err(|source| CutlassCompilerError::LibraryMetadata { + path: path.to_path_buf(), + source, + })?; + if !metadata.file_type().is_file() { + return Err(CutlassCompilerError::LibraryNotRegular { + path: path.to_path_buf(), + }); + } + + let mut file = File::open(path).map_err(|source| CutlassCompilerError::LibraryOpen { + path: path.to_path_buf(), + source, + })?; + let opened_metadata = + file.metadata() + .map_err(|source| CutlassCompilerError::LibraryMetadata { + path: path.to_path_buf(), + source, + })?; + if !opened_metadata.is_file() { + return Err(CutlassCompilerError::LibraryNotRegular { + path: path.to_path_buf(), + }); + } + + let actual = + fingerprint_reader(&mut file).map_err(|source| CutlassCompilerError::LibraryRead { + path: path.to_path_buf(), + source, + })?; + if actual != expected { + return Err(CutlassCompilerError::FingerprintMismatch { + path: path.to_path_buf(), + expected, + actual, + }); + } + + Ok(ValidatedLibraryFile { + path: path.to_path_buf(), + file, + fingerprint: actual, + }) +} + +fn fingerprint_reader(reader: &mut impl Read) -> Result { + let mut hasher = Sha256::new(); + let mut buffer = [0_u8; 64 * 1024]; + loop { + let count = reader.read(&mut buffer)?; + if count == 0 { + break; + } + hasher.update(&buffer[..count]); + } + Ok(Sha256Fingerprint(hasher.finalize().into())) +} + +struct Api { + // Keep the dynamic-library handle alive for every resolved function + // pointer. On Linux, keep the opened descriptor alive too so provenance + // continues to name the loaded inode even if its pathname is replaced. + _library: Library, + _loaded_file: File, + path: PathBuf, + fingerprint: Sha256Fingerprint, + cute_compiler_create: CompilerCreateFn, + compiler_destroy: CompilerDestroyFn, + compiler_set_device_target: CompilerSetDeviceTargetFn, + compiler_set_abi: CompilerSetAbiFn, + compiler_set_pipeline: CompilerSetPipelineFn, + compiler_compile_to: CompilerCompileToFn, + artifacts_from_textual_form: ArtifactsFromTextFn, + artifacts_get_type: ArtifactsGetTypeFn, + artifacts_get_data: ArtifactsGetDataFn, + artifacts_function_count: ArtifactsFunctionCountFn, + artifacts_destroy: ArtifactsDestroyFn, +} + +impl Api { + fn load(validated: ValidatedLibraryFile) -> Result { + #[cfg(target_os = "linux")] + let load_path = PathBuf::from(format!("/proc/self/fd/{}", validated.file.as_raw_fd())); + #[cfg(not(target_os = "linux"))] + let load_path = validated.path.clone(); + + let library = unsafe { Library::new(&load_path) }.map_err(|source| { + CutlassCompilerError::LibraryLoad { + path: validated.path.clone(), + source, + } + })?; + + unsafe { + Ok(Self { + cute_compiler_create: resolve(&library, "cutlass_cute_compiler_create")?, + compiler_destroy: resolve(&library, "cutlass_compiler_destroy")?, + compiler_set_device_target: resolve( + &library, + "cutlass_compiler_set_device_target", + )?, + compiler_set_abi: resolve(&library, "cutlass_compiler_set_abi")?, + compiler_set_pipeline: resolve(&library, "cutlass_compiler_set_pipeline")?, + compiler_compile_to: resolve(&library, "cutlass_compiler_compile_to")?, + artifacts_from_textual_form: resolve( + &library, + "cutlass_artifacts_from_textual_form", + )?, + artifacts_get_type: resolve(&library, "cutlass_artifacts_get_type")?, + artifacts_get_data: resolve(&library, "cutlass_artifacts_get_data")?, + artifacts_function_count: resolve(&library, "cutlass_artifacts_function_count")?, + artifacts_destroy: resolve(&library, "cutlass_artifacts_destroy")?, + path: validated.path, + fingerprint: validated.fingerprint, + _loaded_file: validated.file, + _library: library, + }) + } + } +} + +unsafe fn resolve( + library: &Library, + name: &'static str, +) -> Result { + let symbol: Symbol = unsafe { library.get(name.as_bytes()) }.map_err(|source| { + CutlassCompilerError::SymbolNotFound { + symbol: name, + source, + } + })?; + Ok(unsafe { *symbol.into_raw() }) +} + +struct CompilerHandle<'api> { + api: &'api Api, + raw: CompilerRef, +} + +impl<'api> CompilerHandle<'api> { + fn create_cute(api: &'api Api) -> Result { + let raw = unsafe { (api.cute_compiler_create)() }; + if raw.is_null() { + return Err(CutlassCompilerError::NullHandle { + operation: "cutlass_cute_compiler_create", + }); + } + Ok(Self { api, raw }) + } + + fn set_device_target(&self, target: &str) -> Result<(), CutlassCompilerError> { + let mut error_ptr = ptr::null_mut(); + let mut error_len = 0; + let status = unsafe { + (self.api.compiler_set_device_target)( + self.raw, + target.as_ptr().cast(), + target.len(), + &mut error_ptr, + &mut error_len, + ) + }; + let message = unsafe { take_error_message(error_ptr, error_len) }; + if status == 0 { + Ok(()) + } else { + Err(CutlassCompilerError::Configure { + operation: "cutlass_compiler_set_device_target", + status, + message, + }) + } + } + + fn set_external_module_passthrough_abi(&self) -> Result<(), CutlassCompilerError> { + let mut error_ptr = ptr::null_mut(); + let mut error_len = 0; + let status = unsafe { + (self.api.compiler_set_abi)( + self.raw, + CUTLASS_ABI_TBD_EXTERNAL_MODULE_PASSTHROUGH, + &mut error_ptr, + &mut error_len, + ) + }; + let message = unsafe { take_error_message(error_ptr, error_len) }; + if status == 0 { + Ok(()) + } else { + Err(CutlassCompilerError::Configure { + operation: "cutlass_compiler_set_abi", + status, + message, + }) + } + } + + fn set_pipeline( + &self, + artifact_type: c_int, + pipeline: &str, + ) -> Result<(), CutlassCompilerError> { + let status = unsafe { + (self.api.compiler_set_pipeline)( + self.raw, + artifact_type, + pipeline.as_ptr().cast(), + pipeline.len(), + ) + }; + if status == 0 { + Ok(()) + } else { + Err(CutlassCompilerError::PipelineConfiguration { + artifact_type, + status, + }) + } + } + + fn compile_to_object( + &self, + input: &mut ArtifactsHandle<'api>, + ) -> Result, CutlassCompilerError> { + let mut error_ptr = ptr::null_mut(); + let mut error_len = 0; + let raw = unsafe { + (self.api.compiler_compile_to)( + self.raw, + input.raw, + CUTLASS_ARTIFACT_OBJECT, + &mut error_ptr, + &mut error_len, + ) + }; + let message = unsafe { take_error_message(error_ptr, error_len) }; + if raw.is_null() { + return Err(CutlassCompilerError::Compile { message }); + } + Ok(ArtifactsHandle { api: self.api, raw }) + } +} + +impl Drop for CompilerHandle<'_> { + fn drop(&mut self) { + unsafe { (self.api.compiler_destroy)(self.raw) }; + } +} + +struct ArtifactsHandle<'api> { + api: &'api Api, + raw: ArtifactsRef, +} + +impl<'api> ArtifactsHandle<'api> { + fn from_textual_precompiled_mlir( + api: &'api Api, + mlir: &str, + ) -> Result { + Self::from_textual_mlir( + api, + mlir, + CUTLASS_ARTIFACT_PRE_COMPILED_MLIR, + "PreCompiledMlir", + ) + } + + fn from_textual_mlir( + api: &'api Api, + mlir: &str, + artifact_type: c_int, + artifact_kind: &'static str, + ) -> Result { + let raw = unsafe { + (api.artifacts_from_textual_form)(artifact_type, mlir.as_ptr().cast(), mlir.len()) + }; + if raw.is_null() { + return Err(CutlassCompilerError::MlirParse { artifact_kind }); + } + Ok(Self { api, raw }) + } + + fn artifact_type(&self) -> c_int { + unsafe { (self.api.artifacts_get_type)(self.raw) } + } + + fn data(&self) -> Result, CutlassCompilerError> { + let mut data_ptr = ptr::null_mut(); + let mut data_len = 0; + let status = + unsafe { (self.api.artifacts_get_data)(self.raw, &mut data_ptr, &mut data_len) }; + if status != 0 { + unsafe { free_if_nonnull(data_ptr) }; + return Err(CutlassCompilerError::ArtifactQuery { + operation: "cutlass_artifacts_get_data", + reason: format!("returned status {status}"), + }); + } + unsafe { copy_malloc_bytes(data_ptr, data_len, "cutlass_artifacts_get_data") } + } + + fn function_count(&self) -> Result { + let mut count = 0; + let status = unsafe { (self.api.artifacts_function_count)(self.raw, &mut count) }; + if status == 0 { + Ok(count) + } else { + Err(CutlassCompilerError::ArtifactQuery { + operation: "cutlass_artifacts_function_count", + reason: format!("returned status {status}"), + }) + } + } +} + +impl Drop for ArtifactsHandle<'_> { + fn drop(&mut self) { + unsafe { (self.api.artifacts_destroy)(self.raw) }; + } +} + +unsafe fn take_error_message(ptr: *mut c_char, len: usize) -> String { + if ptr.is_null() { + return "no diagnostic returned".to_string(); + } + if len > isize::MAX as usize { + unsafe { free(ptr.cast()) }; + return format!("diagnostic length {len} exceeds Rust's addressable slice limit"); + } + let bytes = unsafe { slice::from_raw_parts(ptr.cast::(), len) }; + let message = String::from_utf8_lossy(bytes).into_owned(); + unsafe { free(ptr.cast()) }; + message +} + +unsafe fn copy_malloc_bytes( + ptr: *mut u8, + len: usize, + operation: &'static str, +) -> Result, CutlassCompilerError> { + if ptr.is_null() { + if len == 0 { + return Ok(Vec::new()); + } + return Err(CutlassCompilerError::ArtifactQuery { + operation, + reason: format!("returned a null pointer with non-zero length {len}"), + }); + } + if len > isize::MAX as usize { + unsafe { free(ptr.cast()) }; + return Err(CutlassCompilerError::ArtifactQuery { + operation, + reason: format!("returned length {len} exceeds Rust's addressable slice limit"), + }); + } + let bytes = unsafe { slice::from_raw_parts(ptr, len) }.to_vec(); + unsafe { free(ptr.cast()) }; + Ok(bytes) +} + +unsafe fn free_if_nonnull(ptr: *mut u8) { + if !ptr.is_null() { + unsafe { free(ptr.cast()) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + fn fingerprint(data: &[u8]) -> Sha256Fingerprint { + fingerprint_reader(&mut Cursor::new(data)).expect("in-memory digest cannot fail") + } + + #[test] + fn fingerprint_parser_is_exact_and_round_trips() { + let expected = fingerprint(b"CUTLASS 4.7 compiler"); + let lower = expected.to_string(); + let upper = lower.to_ascii_uppercase(); + + assert_eq!(lower.parse::().unwrap(), expected); + assert_eq!(upper.parse::().unwrap(), expected); + assert_eq!(expected.as_bytes().len(), 32); + assert!(matches!( + "00".parse::(), + Err(FingerprintParseError::Length { actual: 2 }) + )); + assert!(matches!( + format!("{}zz", &lower[..62]).parse::(), + Err(FingerprintParseError::NonHex { offset: 62 }) + )); + let unicode = format!("{}é", "0".repeat(62)); + assert_eq!(unicode.len(), 64); + assert!(matches!( + unicode.parse::(), + Err(FingerprintParseError::NonHex { offset: 62 }) + )); + } + + #[test] + fn library_validation_requires_an_absolute_regular_file() { + let digest = fingerprint(b"unused"); + let relative = Path::new("libCutlassCompiler.so"); + assert!(matches!( + validate_library_file(relative, digest), + Err(CutlassCompilerError::RelativeLibraryPath { .. }) + )); + + let directory = tempfile::tempdir().unwrap(); + assert!(matches!( + validate_library_file(directory.path(), digest), + Err(CutlassCompilerError::LibraryNotRegular { .. }) + )); + } + + #[test] + fn library_validation_checks_the_opened_file_digest() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("libCutlassCompiler.so"); + let contents = b"not a real library, sufficient for validation"; + fs::write(&path, contents).unwrap(); + + let expected = fingerprint(contents); + let validated = validate_library_file(&path, expected).unwrap(); + assert_eq!(validated.path, path); + assert_eq!(validated.fingerprint, expected); + + let wrong = fingerprint(b"different library"); + assert!(matches!( + validate_library_file(&path, wrong), + Err(CutlassCompilerError::FingerprintMismatch { + expected: rejected, + actual, + .. + }) if rejected == wrong && actual == expected + )); + } + + #[cfg(unix)] + #[test] + fn library_validation_rejects_a_symlink() { + use std::os::unix::fs::symlink; + + let directory = tempfile::tempdir().unwrap(); + let target = directory.path().join("real.so"); + let link = directory.path().join("libCutlassCompiler.so"); + let contents = b"compiler"; + fs::write(&target, contents).unwrap(); + symlink(&target, &link).unwrap(); + + assert!(matches!( + validate_library_file(&link, fingerprint(contents)), + Err(CutlassCompilerError::LibraryNotRegular { .. }) + )); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/error.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/error.rs index d47931b2dd..9febb1d74c 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/error.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/error.rs @@ -32,6 +32,8 @@ pub enum PipelineError { LibdeviceUnavailable { message: String }, /// LLVM IR export failed. Export(String), + /// The selected official CUTLASS MLIR backend failed. + BackendB(String), /// The requested CUDA target could not be parsed, or cannot lower a /// feature the module needs. Distinct from `PtxGeneration`: this is a /// problem with the requested target itself, decided before `llc` runs. @@ -81,6 +83,7 @@ impl std::fmt::Display for PipelineError { write!(f, "libdevice linking is unavailable: {message}") } Self::Export(msg) => write!(f, "Export failed: {}", msg), + Self::BackendB(msg) => write!(f, "CUTLASS MLIR backend failed: {msg}"), Self::TargetSelection { reason, .. } => write!(f, "{reason}"), Self::PtxGeneration(msg) => write!(f, "PTX generation failed: {}", msg), Self::Optimization(msg) => write!(f, "LLVM optimization failed: {msg}"), diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs index b98af30924..c668285b8b 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs index 31ad15d2aa..568cc0701c 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `clc` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs index 03e9b01832..6856e03f01 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs index 8b71148cd8..2f9470f016 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs index a70f45bc9a..a0300a9cc4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs index b9dd3af0a8..145e58b9cb 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs index e18d1602b7..5532f7f335 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs index e3a82698e6..99b91d3a36 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `elect` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs index 5e62065edf..84ef1ddd66 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs index 1e7852cd3e..f7ae666f94 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs index 06f2b97853..d1233ea0bc 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs index afb45c1cd6..dc9679d517 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs index 92cc6bef9f..21ac163b92 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs index cbd62399d2..eb92247be1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `mbarrier_extended` intrinsics. @@ -522,4 +522,134 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ result_range: None, }), }, + GeneratedIntrinsicTarget { + marker: "v1:i1040", + id: "mbarrier_arrive_cluster_relaxed", + abi_id: "i1040", + dialect_op: "nvvm.mbarrier_arrive_cluster_relaxed", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm( + 90, + )]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ], + selections: &[], + llvm: None, + }, + GeneratedIntrinsicTarget { + marker: "v1:i1041", + id: "mbarrier_try_wait_parity_cluster_relaxed", + abi_id: "i1041", + dialect_op: "nvvm.mbarrier_try_wait_parity_cluster_relaxed", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm( + 90, + )]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ], + selections: &[], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "IntrArgMemOnly", + "IntrConvergent", + "IntrNoCallback", + "IntrReadMem", + "NoCapture", + ], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1042", + id: "fence_acquire_shared_cluster_cluster", + abi_id: "i1042", + dialect_op: "nvvm.fence_acquire_shared_cluster_cluster", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm( + 90, + )]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ], + selections: &[GeneratedSelectionAlternative { + source_record: "INT_FENCE_ACQUIRE_SYNC_RESTRICT_CLUSTER_CLUSTER", + asm: "fence.acquire.sync_restrict::shared::cluster.cluster;", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCallback"], + result_no_undef: false, + result_range: None, + }), + }, ]; diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs index fae12bae4d..14c3e2e588 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated target requirements and separately imported LLVM/selection facts. @@ -162,6 +162,7 @@ pub enum GeneratedRegisterMmaKind { Standard, F8f6f4, Mxf8f6f4, + Mxf4, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum GeneratedRegisterMmaAccumulator { @@ -867,6 +868,9 @@ pub fn generated_intrinsic_operation_matches( GeneratedRegisterMmaKind::Mxf8f6f4 => { op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf8f6f4 } + GeneratedRegisterMmaKind::Mxf4 => { + op.kind_or_inferred(ctx) == RegisterMmaKindAttr::Mxf4 + } }; let accumulator_matches = match accumulator { GeneratedRegisterMmaAccumulator::F16 => { diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs index 082f67415a..1fb52887a1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs index 03d04d4ca7..b93de0cf82 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs index b16813cd15..c4bd304fb1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs index 5d2d7956f6..21c3d5da36 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs index 218f67468c..692a73beb5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `prmt` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs index c7cc2c8688..f67f400c30 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `redux` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs index d39dec27bb..fdf784cbf3 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `register_mma` intrinsics. @@ -9712,4 +9712,50 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ result_range: None, }), }, + GeneratedIntrinsicTarget { + marker: "v1:i1039", + id: "mma_m16n8k64_mxf4_f32", + abi_id: "i1039", + dialect_op: "nvvm.register_mma", + variant: GeneratedIntrinsicVariant::RegisterMma { + shape: GeneratedRegisterMmaShape::M16n8k64, + operation: GeneratedRegisterMmaOperation::Multiply, + kind: GeneratedRegisterMmaKind::Mxf4, + accumulator: GeneratedRegisterMmaAccumulator::F32, + a_element: GeneratedRegisterMmaElement::E2m1, + b_element: GeneratedRegisterMmaElement::E2m1, + a_layout: GeneratedRegisterMmaLayout::Row, + b_layout: GeneratedRegisterMmaLayout::Col, + overflow: GeneratedRegisterMmaOverflow::NotApplicable, + }, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(87), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(120), + ]), + }, + backend_requirements: &[GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(87), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(120), + ]), + }, + }], + selections: &[GeneratedSelectionAlternative { + source_record: "anonymous_19058", + asm: "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 \n\t\t{{$rd0, $rd1, $rd2, $rd3}},\n\t\t{{$ra0, $ra1, $ra2, $ra3}},\n\t\t{{$rb0, $rb1}},\n\t\t{{$rc0, $rc1, $rc2, $rc3}},\n\t\t$scale_a, {{$byte_id_a, $thread_id_a}},\n\t\t$scale_b, {{$byte_id_b, $thread_id_b}};", + predicates: &["Subtarget->hasMMABlockScale()"], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCallback", "IntrNoMem"], + result_no_undef: false, + result_range: None, + }), + }, ]; diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs index 60a21dbe3f..4879de292a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs index a41ae115e0..99b72cda00 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs index 74fc6f9745..7643253888 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs index 5d932855a2..c5cc33b4c7 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs index 425aff37d8..47c9a30ead 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `sreg` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs index cc5ca29161..cde7d74fb3 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs index 80010ce4bb..9f91e3e4e0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `sync` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs index f85ca20ab1..af657ccb8a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_cp` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs index 8fe001d6b5..c76ca85005 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_ld` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs index be40bde5d2..c2e6f05061 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_mma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs index 1e621efbe7..0d52b8ca98 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_mma_ws.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_mma_ws` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs index d3fd5dbdff..8b9850a3c9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_other` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs index 82f51f8e93..3eb4ac2e77 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tcgen05_st` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs index 26ca71b44d..bd470b5d33 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs index 7c98e49459..08d48ecaff 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `clc` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs index 734046a818..0021713ea8 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs index b6912253dc..af9eca62ae 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs index 1de82cf936..3781371112 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs index 9166dcc661..6ec580813f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs index 9f0a3fdeb1..1ac767695a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs index fb0ec54a76..716a6dd5ec 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `elect` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs index 26b58a196c..0d5862ac32 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs index 6366219601..326b18ec16 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs index ccb9d78ade..d49447c7f1 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs index 8443460449..4695f42f57 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs index 84a9a99d29..6c15db9167 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs index 489f4002ad..782b9ab218 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `mbarrier_extended` intrinsics. @@ -604,3 +604,158 @@ fn nanosleep_target_matches_the_catalog() { &["IntrConvergent", "IntrHasSideEffects", "IntrNoMem"] as &[&str] ); } + +#[test] +fn mbarrier_arrive_cluster_relaxed_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1040").unwrap(); + assert_eq!(target.id, "mbarrier_arrive_cluster_relaxed"); + assert_eq!(target.abi_id, "i1040"); + assert_eq!(target.dialect_op, "nvvm.mbarrier_arrive_cluster_relaxed"); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 86); + assert_eq!( + target.requirement.hardware, + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(90)]) } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert!(target.llvm.is_none()); +} + +#[test] +fn mbarrier_try_wait_parity_cluster_relaxed_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1041").unwrap(); + assert_eq!(target.id, "mbarrier_try_wait_parity_cluster_relaxed"); + assert_eq!(target.abi_id, "i1041"); + assert_eq!( + target.dialect_op, + "nvvm.mbarrier_try_wait_parity_cluster_relaxed" + ); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 86); + assert_eq!( + target.requirement.hardware, + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(90)]) } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "IntrArgMemOnly", + "IntrConvergent", + "IntrNoCallback", + "IntrReadMem", + "NoCapture" + ] as &[&str] + ); +} + +#[test] +fn fence_acquire_shared_cluster_cluster_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1042").unwrap(); + assert_eq!(target.id, "fence_acquire_shared_cluster_cluster"); + assert_eq!(target.abi_id, "i1042"); + assert_eq!( + target.dialect_op, + "nvvm.fence_acquire_shared_cluster_cluster" + ); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 86); + assert_eq!( + target.requirement.hardware, + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(90)]) } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(86), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::MinimumSm(90), + ]), + }, + }, + ] + } + ); + assert_eq!( + target.selections, + &[GeneratedSelectionAlternative { + source_record: "INT_FENCE_ACQUIRE_SYNC_RESTRICT_CLUSTER_CLUSTER", + asm: "fence.acquire.sync_restrict::shared::cluster.cluster;", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCallback"] as &[&str] + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs index 3563eb4391..0915796fd0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests, grouped per family shard. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs index 0b11ce5341..98db4aaf4e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs index 6bd00d7b5f..e0833bcf44 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs index 662abc5144..52a9eec446 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs index 7946d3a243..394e138410 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs index 68b0189d30..4f8f92df86 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `prmt` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs index 02715a5695..5e9be59f08 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `redux` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs index 85b894a235..3445c4167e 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_float.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `register_mma_float` intrinsics. @@ -7860,3 +7860,60 @@ fn mma_m16n8k32_mxf8f6f4_f32_e5m2_e5m2_target_matches_the_catalog() { &["IntrNoCallback", "IntrNoMem"] as &[&str] ); } + +#[test] +fn mma_m16n8k64_mxf4_f32_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1039").unwrap(); + assert_eq!(target.id, "mma_m16n8k64_mxf4_f32"); + assert_eq!(target.abi_id, "i1039"); + assert_eq!(target.dialect_op, "nvvm.register_mma"); + assert_eq!( + target.variant, + GeneratedIntrinsicVariant::RegisterMma { + shape: GeneratedRegisterMmaShape::M16n8k64, + operation: GeneratedRegisterMmaOperation::Multiply, + kind: GeneratedRegisterMmaKind::Mxf4, + accumulator: GeneratedRegisterMmaAccumulator::F32, + a_element: GeneratedRegisterMmaElement::E2m1, + b_element: GeneratedRegisterMmaElement::E2m1, + a_layout: GeneratedRegisterMmaLayout::Row, + b_layout: GeneratedRegisterMmaLayout::Col, + overflow: GeneratedRegisterMmaOverflow::NotApplicable + } + ); + assert_eq!(target.requirement.minimum_ptx.encoded(), 87); + assert_eq!( + target.requirement.hardware, + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::ExactArchitecture(120)]) } + ); + assert_eq!( + target.backend_requirements, + const { + &[GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(87), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(120), + ]), + }, + }] + } + ); + assert_eq!( + target.selections, + &[GeneratedSelectionAlternative { + source_record: "anonymous_19058", + asm: "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 \n\t\t{{$rd0, $rd1, $rd2, $rd3}},\n\t\t{{$ra0, $ra1, $ra2, $ra3}},\n\t\t{{$rb0, $rb1}},\n\t\t{{$rc0, $rc1, $rc2, $rc3}},\n\t\t$scale_a, {{$byte_id_a, $thread_id_a}},\n\t\t$scale_b, {{$byte_id_b, $thread_id_b}};", + predicates: &["Subtarget->hasMMABlockScale()"], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCallback", "IntrNoMem"] as &[&str] + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs index 039355ee52..2840232222 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/register_mma_s32.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `register_mma_s32` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs index 6b65cae93b..9ef28b4104 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs index a7c82b59bb..dc839e76a9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs index dda7dcaba5..b7c4b6ec20 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs index cd986ff572..acd75e6af9 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs index 0f049292aa..3079cc1921 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `sreg` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs index 04de6265b2..4ec333adc8 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs index 56e87ad452..a0d9049a20 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `sync` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs index 4386115e63..827fc47376 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_cp` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs index 1017a9f455..719b7d4f2f 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_ld` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs index a077eb0e5b..e32e1b4ebb 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_mma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs index a6b1cbc212..bb30d706d4 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_mma_ws.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_mma_ws` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs index 4fe601b16c..f58b3845c6 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_other` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs index aaa9c828d9..3dae5d3a7a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tcgen05_st` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs index 2072748ea7..09d565f877 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs index 72b99c50d2..70f661f183 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `vote` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs index 901e09e656..6f453a5af0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs index 87fb6f18d8..e54059dcaf 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs index 69e7fd10fe..71d40627c0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs index 1fe226c429..25b0fb411d 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tests/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs index a1858f9e89..578073fbd3 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tma` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs index 44f829e5df..5ab177077a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `vote` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs index bc7aad6911..0af0a26bd0 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs index 12f9894c87..9d5d580445 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs index 0b1891da9d..e4460f77fe 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs index 3a205d31cc..ee4dc522d5 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/generated_intrinsic_targets/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs index 5269d0ba64..1d4ac4ac87 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs @@ -728,6 +728,8 @@ mod tests { let files = OutputFiles { llvm_ir: std::path::Path::new("unused.ll"), ptx: std::path::Path::new("unused.ptx"), + cubin: std::path::Path::new("unused.cubin"), + mlir: None, stale_before_export: &[], }; let request = ModulePipelineRequest::for_rust_pipeline( @@ -735,6 +737,8 @@ mod tests { nvvm, &backend, debug_kind, + None, + &[], files, PipelineTrace::default(), ); @@ -767,6 +771,7 @@ mod tests { let root = tempfile::tempdir().unwrap(); let llvm_ir = root.path().join("module.ll"); let ptx = root.path().join("module.ptx"); + let cubin = root.path().join("module.cubin"); let backend = BackendOptions { device_arch_hint: Some(DeviceArchHint::parse("sm_120".to_string())), verbose: true, @@ -777,9 +782,13 @@ mod tests { nvvm, &backend, llvm_export::export::DebugKind::Off, + None, + &[], OutputFiles { llvm_ir: &llvm_ir, ptx: &ptx, + cubin: &cubin, + mlir: None, stale_before_export: &[], }, PipelineTrace::default(), diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/lib.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/lib.rs index ad0bcca57a..ba090dc822 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/lib.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/lib.rs @@ -11,11 +11,13 @@ //! the rustc frontend. //! //! This crate intentionally has no rustc linkage. It does not require -//! `rustc_private` or a nightly toolchain matched to `rustc_driver`. +//! `rustc_private` or a toolchain pin matched to `rustc_driver`. #![warn(missing_docs)] mod api; +mod cutlass_backend; +mod cutlass_compiler; mod error; mod export; mod generated; @@ -112,6 +114,8 @@ pub mod experimental { /// This is not part of the standalone frontend contract. #[doc(hidden)] pub mod __private { + #[doc(hidden)] + pub use crate::cutlass_backend::{CutlassBackendConfig, DeviceBackend}; #[doc(hidden)] pub use crate::error::PipelineError; #[doc(hidden)] diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/pipeline.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/pipeline.rs index 115e0d154c..d9b9c863c3 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/pipeline.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/pipeline.rs @@ -9,6 +9,9 @@ //! every destructive compiler stage after that boundary so the rustc frontend //! and the standalone frontend cannot silently diverge. +use crate::cutlass_backend::{ + CUTLASS_FULL_CUTE_MLIR22_PROFILE, CutlassBackendConfig, compile_to_cubin, +}; use crate::error::PipelineError; use crate::export::{ DebugExport, DeviceExternDecl, export_llvm_ir, function_local_static_placement_for_llc, @@ -30,11 +33,18 @@ use crate::ptx::{ }; use crate::target::detect_features_in_llvm_text; use crate::verify::verify_operation; +use cuda_oxide_mlir_export::render_module as render_mlir_module; +use cuda_oxide_mlir_export::{ + CutlassFullCuteMlir22, CutlassPrimitivesMlir22, ExportManifest, MlirConsumerProfile, +}; use llvm_export::export::{DebugKind, FunctionLocalStaticPlacement, NvvmIrDialect}; +use pliron::builtin::ops::ModuleOp; use pliron::context::{Context, Ptr}; use pliron::linked_list::ContainsLinkedList; use pliron::operation::Operation; use pliron::printable::Printable; +use std::ffi::OsString; +use std::io; use std::path::{Path, PathBuf}; /// Files used to publish one compiler result. @@ -45,6 +55,10 @@ pub struct OutputFiles<'a> { pub llvm_ir: &'a Path, /// PTX output. This remains absent for an NVVM IR result. pub ptx: &'a Path, + /// Cubin output. This is populated only by the selected CUTLASS backend. + pub cubin: &'a Path, + /// Optional textual MLIR observation of the prepared semantic module. + pub mlir: Option<&'a Path>, /// Older competing artifacts to remove immediately before export. pub stale_before_export: &'a [PathBuf], } @@ -138,16 +152,21 @@ pub struct ModulePipelineRequest<'a> { files: OutputFiles<'a>, trace: PipelineTrace, generated_marker_policy: GeneratedMarkerPolicy, + cutlass_backend: Option<&'a CutlassBackendConfig>, + expected_kernels: &'a [String], } impl<'a> ModulePipelineRequest<'a> { /// Build the link-capable request used by the rustc frontend. #[doc(hidden)] + #[allow(clippy::too_many_arguments)] pub fn for_rust_pipeline( device_externs: &'a [DeviceExternDecl], request_nvvm_ir: bool, backend: &'a BackendOptions, debug_kind: DebugKind, + cutlass_backend: Option<&'a CutlassBackendConfig>, + expected_kernels: &'a [String], files: OutputFiles<'a>, trace: PipelineTrace, ) -> Self { @@ -160,6 +179,8 @@ impl<'a> ModulePipelineRequest<'a> { files, trace, generated_marker_policy: GeneratedMarkerPolicy::Required, + cutlass_backend, + expected_kernels, } } @@ -179,6 +200,8 @@ impl<'a> ModulePipelineRequest<'a> { files, trace: PipelineTrace::default(), generated_marker_policy: GeneratedMarkerPolicy::Optional, + cutlass_backend: None, + expected_kernels: &[], } } } @@ -191,6 +214,8 @@ pub enum ModuleArtifactKind { Ptx, /// NVVM-compatible LLVM IR for a later libNVVM/nvJitLink step. NvvmIr, + /// Final cubin generated by the official CUTLASS compiler library. + Cubin, } /// Successful shared-pipeline result. @@ -205,7 +230,7 @@ pub struct ModulePipelineOutput { pub diagnostics: Vec, } -/// Compile an already-translated `dialect-mir` module to PTX or NVVM IR. +/// Compile a `dialect-mir` module to PTX, NVVM IR, or a CUTLASS cubin. #[doc(hidden)] pub fn compile_translated_module( ctx: &mut Context, @@ -255,6 +280,14 @@ pub fn compile_translated_module( } let backend = validated.options(); + // Primitives must reject semantic source objects before preparation can + // fold away their operations or erase their pointee types. + if let Some(config) = request.cutlass_backend { + validate_cutlass_profile_name(&config.profile)?; + if config.profile == CutlassPrimitivesMlir22::NAME { + verify_source_for_backend(ctx, module, Some(config))?; + } + } let promote_and_unroll = !request.debug_kind.variables_enabled(); if !promote_and_unroll && has_iket_operations(ctx, module) { return Err(PipelineError::Lowering( @@ -294,11 +327,98 @@ pub fn compile_translated_module( .emit(format!("{}", module.deref(ctx).disp(ctx))); } + // Check tensor origins, pipeline state, and operation ordering before + // either backend lowers the CuTe operations. + verify_source_for_backend(ctx, module, request.cutlass_backend)?; + // After the pinning above, `backend.target_arch` is exactly the target // the rest of the pipeline will compile for (or `None`, in which case // materialization fails with its explicit-target requirement). materialize_iket(ctx, module, backend.target_arch.as_deref(), &backend.iket)?; + // Export or compile the prepared module while CuTe operations are intact. + // Skip MLIR translation unless a dump or the CUTLASS backend was requested. + let cutlass_rendered = if request.files.mlir.is_some() || request.cutlass_backend.is_some() { + Some(render_live_cutlass_mlir( + ctx, + module, + &validated, + request + .cutlass_backend + .map_or(CUTLASS_FULL_CUTE_MLIR22_PROFILE, |c| c.profile.as_str()), + )?) + } else { + None + }; + if let Some(path) = request.files.mlir { + validate_cutlass_mlir_output_path(path, &request.files)?; + } + if let Some(config) = request.cutlass_backend { + validate_cutlass_compile_policy(request.backend, request.debug_kind)?; + if matches!( + request.output_policy, + OutputPolicy::ExternalLinkAllowed { + request_nvvm_ir: true + } + ) { + return Err(PipelineError::BackendB( + "CUDA_OXIDE_EMIT_NVVM_IR cannot be combined with CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir; the CUTLASS backend already produces a final cubin" + .to_owned(), + )); + } + let rendered = cutlass_rendered + .as_ref() + .expect("the selected CUTLASS backend always renders its input"); + remove_stale_files(request.files.stale_before_export)?; + if request.trace.verbose { + request.trace.emit(format!( + "\n=== Running official CUTLASS compiler → {} ===", + request.files.cubin.display() + )); + } + let generated = compile_to_cubin( + config, + &rendered.text, + &rendered.target, + request.expected_kernels, + request.files.cubin, + request.files.mlir, + )?; + for diagnostic in &generated.diagnostics { + request.trace.emit(diagnostic); + } + if request.trace.verbose { + if let Some(path) = request.files.mlir { + request + .trace + .emit(format!("CUTLASS MLIR written to {}", path.display())); + } + request.trace.emit(format!( + "✓ cubin written to {} (target: {})", + request.files.cubin.display(), + rendered.target + )); + } + return Ok(ModulePipelineOutput { + artifact_kind: ModuleArtifactKind::Cubin, + target: rendered.target.clone(), + diagnostics: generated.diagnostics, + }); + } + if let (Some(path), Some(rendered)) = (request.files.mlir, cutlass_rendered.as_ref()) { + std::fs::write(path, &rendered.text).map_err(|error| { + PipelineError::Export(format!( + "could not write requested CUTLASS MLIR to {}: {error}", + path.display() + )) + })?; + if request.trace.verbose { + request + .trace + .emit(format!("CUTLASS MLIR written to {}", path.display())); + } + } + // Calls need structured extern declarations before lowering so pointer // address spaces are preserved by the call converter. if !request.device_externs.is_empty() { @@ -626,6 +746,257 @@ pub fn compile_translated_module( }) } +fn validate_cutlass_compile_policy( + backend: &BackendOptions, + debug_kind: DebugKind, +) -> Result<(), PipelineError> { + if backend.no_fma { + return Err(PipelineError::BackendB( + "CUDA_OXIDE_NO_FMA/--no-fmad is not implemented by the pinned CUTLASS backend; refusing to compile with different floating-point contraction semantics" + .to_owned(), + )); + } + if backend.no_opt { + return Err(PipelineError::BackendB( + "unoptimized device code is not implemented by the pinned CUTLASS backend; refusing to silently use its fixed optimized pipeline" + .to_owned(), + )); + } + if debug_kind != DebugKind::Off { + return Err(PipelineError::BackendB( + "device debug information is not implemented by the pinned CUTLASS backend; use the native backend for line tables or variable debug information" + .to_owned(), + )); + } + Ok(()) +} + +fn validate_cutlass_mlir_output_path( + output: &Path, + files: &OutputFiles<'_>, +) -> Result<(), PipelineError> { + let normalized_output = normalize_artifact_path(output).map_err(|error| { + PipelineError::Export(format!( + "could not safely normalize requested CUTLASS MLIR path {}: {error}", + output.display() + )) + })?; + for native in std::iter::once(files.llvm_ir) + .chain(std::iter::once(files.ptx)) + .chain(std::iter::once(files.cubin)) + .chain(files.stale_before_export.iter().map(PathBuf::as_path)) + { + let normalized_native = normalize_artifact_path(native).map_err(|error| { + PipelineError::Export(format!( + "could not safely normalize ordinary compiler artifact path {} while checking CUTLASS MLIR output {}: {error}", + native.display(), + output.display() + )) + })?; + if normalized_output == normalized_native { + return Err(PipelineError::Export(format!( + "requested CUTLASS MLIR path {} collides with ordinary compiler artifact {}", + output.display(), + native.display() + ))); + } + } + Ok(()) +} + +/// Resolve every existing path component, then normalize the remaining tail. +/// +/// `canonicalize` alone cannot handle a not-yet-created output file. Walking +/// upward to the first resolvable ancestor preserves symlink resolution for +/// every existing directory while still giving new sibling artifacts a stable +/// absolute identity. A broken symlink is an error rather than an ordinary +/// missing component: treating it as text could let the later write follow a +/// path the collision check never inspected. +fn normalize_artifact_path(path: &Path) -> io::Result { + let mut probe = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir()?.join(path) + }; + let mut unresolved = Vec::::new(); + + let mut resolved = loop { + match std::fs::canonicalize(&probe) { + Ok(resolved) => break resolved, + Err(error) if error.kind() == io::ErrorKind::NotFound => { + match std::fs::symlink_metadata(&probe) { + Ok(_) => return Err(error), + Err(metadata_error) if metadata_error.kind() == io::ErrorKind::NotFound => {} + Err(metadata_error) => return Err(metadata_error), + } + + let component = probe.components().next_back().ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("path {} has no resolvable ancestor", path.display()), + ) + })?; + unresolved.push(component.as_os_str().to_os_string()); + if !probe.pop() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("path {} has no resolvable ancestor", path.display()), + )); + } + } + Err(error) => return Err(error), + } + }; + + for component in unresolved.iter().rev() { + resolved.push(component); + } + normalize_absolute_lexically(&resolved) +} + +fn normalize_absolute_lexically(path: &Path) -> io::Result { + use std::path::Component; + + let mut normalized = PathBuf::new(); + for component in path.components() { + match component { + Component::Prefix(prefix) => normalized.push(prefix.as_os_str()), + Component::RootDir => normalized.push(component.as_os_str()), + Component::CurDir => {} + Component::ParentDir => { + if !normalized.pop() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("path {} escapes its filesystem root", path.display()), + )); + } + } + Component::Normal(name) => normalized.push(name), + } + } + if !normalized.is_absolute() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "normalized artifact path is not absolute: {}", + path.display() + ), + )); + } + Ok(normalized) +} + +struct RenderedCutlassMlir { + text: String, + target: String, +} + +fn validate_cutlass_profile_name(name: &str) -> Result<(), PipelineError> { + match name { + CUTLASS_FULL_CUTE_MLIR22_PROFILE | CutlassPrimitivesMlir22::NAME => Ok(()), + _ => Err(PipelineError::Export(format!( + "unsupported CUTLASS MLIR profile `{name}`" + ))), + } +} + +fn verify_source_for_backend( + ctx: &Context, + module: Ptr, + cutlass: Option<&CutlassBackendConfig>, +) -> Result<(), PipelineError> { + if let Some(config) = cutlass { + validate_cutlass_profile_name(&config.profile)?; + if config.profile == CutlassPrimitivesMlir22::NAME { + let source = Operation::get_op::(module, ctx).ok_or_else(|| { + PipelineError::Export("primitives source must be builtin.module".into()) + })?; + return cuda_oxide_mlir_export::verify_primitives_source(ctx, &source).map_err(|e| { + PipelineError::Verification { + name: "primitives source".into(), + message: e.to_string(), + operation: None, + } + }); + } + } + // Preserve native and full-CuTe validation, including malformed ghost-only + // modules not recognized by the high-level operation promotion list. + dialect_cute::verify::verify_cute_semantics(ctx, module).map_err(|error| { + PipelineError::Lowering(format!( + "dialect-cute semantic verification failed: {error}" + )) + }) +} + +fn render_live_cutlass_mlir( + ctx: &Context, + module: Ptr, + backend: &ValidatedBackendOptions<'_>, + profile_name: &str, +) -> Result { + validate_cutlass_profile_name(profile_name)?; + let gpu_arch = backend + .options() + .target_arch + .clone() + .or_else(|| backend.device_arch().map(|arch| arch.as_ref().sm())) + .ok_or_else(|| { + PipelineError::Export( + "CUTLASS MLIR was requested, but no GPU architecture is available; set \ + CUDA_OXIDE_TARGET to a profile-supported target" + .to_owned(), + ) + })?; + let profile: Result, _> = match profile_name { + CutlassPrimitivesMlir22::NAME => CutlassPrimitivesMlir22::new(gpu_arch.clone()) + .map(|p| Box::new(p) as Box), + _ => CutlassFullCuteMlir22::new(gpu_arch.clone()) + .map(|p| Box::new(p) as Box), + }; + let profile = profile.map_err(|error| { + PipelineError::Export(format!( + "CUTLASS MLIR profile selection failed for `{gpu_arch}`: {error}" + )) + })?; + let source = Operation::get_op::(module, ctx).ok_or_else(|| { + PipelineError::Export( + "CUTLASS MLIR export expected the live source operation to be builtin.module" + .to_owned(), + ) + })?; + let translated = profile.translate_module(ctx, &source).map_err(|error| { + PipelineError::Export(format!("CUTLASS MLIR translation failed: {error}")) + })?; + let text = render_cutlass_mlir(&profile.manifest(), &translated)?; + Ok(RenderedCutlassMlir { + text, + target: gpu_arch.to_owned(), + }) +} + +fn render_cutlass_mlir( + manifest: &ExportManifest, + module: &cuda_oxide_mlir_export::MlirModule, +) -> Result { + let manifest = manifest.to_pretty_json().map_err(|error| { + PipelineError::Export(format!( + "could not render the CUTLASS profile manifest: {error}" + )) + })?; + let mut output = String::from("// cuda-oxide CUTLASS export manifest\n"); + for line in manifest.lines() { + output.push_str("// "); + output.push_str(line); + output.push('\n'); + } + let rendered = render_mlir_module(module).map_err(|error| { + PipelineError::Export(format!("CUTLASS MLIR rendering failed: {error}")) + })?; + output.push_str(&rendered); + Ok(output) +} + /// Backend decision made from the typed module before MIR lowering starts. #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct PreLoweringBackendSelection { @@ -838,53 +1209,335 @@ mod tests { use super::*; #[test] - fn invalid_device_hints_fail_at_the_shared_input_boundary() { - for raw in [ - "compute_120", - "compute_80", - "sm_", - "foo", - "sm_05", - "sm_120f", - "not-an-arch", + fn primitives_render_uses_the_selected_manifest() { + use pliron::op::Op; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "primitive".try_into().unwrap()); + let options = BackendOptions { + target_arch: Some("sm_100a".into()), + ..Default::default() + }; + let validated = ValidatedBackendOptions::new(&options).unwrap(); + let rendered = render_live_cutlass_mlir( + &ctx, + module.get_operation(), + &validated, + "cutlass-primitives-mlir22", + ) + .unwrap(); + assert!( + rendered + .text + .contains("\"profile\": \"cutlass-primitives-mlir22\"") + ); + assert!(!rendered.text.contains("cutlass-full-cute-mlir22")); + assert!( + render_live_cutlass_mlir(&ctx, module.get_operation(), &validated, "unknown").is_err() + ); + } + + #[test] + fn primitives_source_policy_rejects_cute_before_preparation() { + use pliron::builtin::attributes::TypeAttr; + use pliron::op::Op; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "ghost".try_into().unwrap()); + let cute = dialect_cute::types::CuteWorkTileType::get( + &mut ctx, + dialect_cute::attributes::CuteTileGridAttr::new(1, 1, 1), + ); + let ptr = dialect_mir::types::MirPtrType::get_generic(&mut ctx, cute.into(), false); + module + .get_operation() + .deref_mut(&ctx) + .attributes + .set("ghost".try_into().unwrap(), TypeAttr::new(ptr.into())); + let mut config = CutlassBackendConfig::new("/unopened/compiler.so".into()); + config.profile = "cutlass-primitives-mlir22".into(); + let error = + verify_source_for_backend(&ctx, module.get_operation(), Some(&config)).unwrap_err(); + assert!(error.to_string().contains("cute.work_tile"), "{error}"); + assert!(verify_source_for_backend(&ctx, module.get_operation(), None).is_err()); + config.profile = "cutlass-full-cute-mlir22".into(); + assert!(verify_source_for_backend(&ctx, module.get_operation(), Some(&config)).is_err()); + } + + #[test] + fn cutlass_policy_rejects_unimplemented_fma_optimization_and_debug_modes() { + let mut backend = BackendOptions::default(); + assert!(validate_cutlass_compile_policy(&backend, DebugKind::Off).is_ok()); + + backend.no_fma = true; + let error = validate_cutlass_compile_policy(&backend, DebugKind::Off).unwrap_err(); + assert!(error.to_string().contains("--no-fmad"), "{error}"); + + backend.no_fma = false; + backend.no_opt = true; + let error = validate_cutlass_compile_policy(&backend, DebugKind::Off).unwrap_err(); + assert!(error.to_string().contains("unoptimized"), "{error}"); + + backend.no_opt = false; + let error = validate_cutlass_compile_policy(&backend, DebugKind::LineTables).unwrap_err(); + assert!(error.to_string().contains("debug information"), "{error}"); + } + + #[test] + fn cutlass_export_consumes_the_validated_device_hint_and_preserves_explicit_precedence() { + use pliron::op::Op; + + for (explicit, hint, expected) in [ + (None, Some("sm_120a"), "sm_120a"), + (Some("sm_100a"), Some("sm_120a"), "sm_100a"), + (Some("sm_120a"), None, "sm_120a"), ] { - for nvvm in [false, true] { - for explicit in [None, Some("sm_120"), Some("foo")] { - let mut ctx = Context::new(); - let module = typed_mir_test_module(&mut ctx, &[]); - let backend = BackendOptions { - device_arch_hint: Some(crate::options::DeviceArchHint::parse( - raw.to_string(), - )), - target_arch: explicit.map(str::to_string), - ..BackendOptions::default() - }; - let request = ModulePipelineRequest::for_rust_pipeline( - &[], - nvvm, - &backend, - DebugKind::Off, - OutputFiles { - llvm_ir: Path::new("unused.ll"), - ptx: Path::new("unused.ptx"), - stale_before_export: &[], - }, - PipelineTrace::default(), - ); - let Err(error) = compile_translated_module(&mut ctx, module, &request) else { - panic!("invalid hint must fail at input"); - }; - assert!(matches!(error, PipelineError::TargetSelection { .. })); - assert_eq!( - error.to_string(), - format!( - "invalid CUDA_OXIDE_DEVICE_ARCH `{raw}`: expected sm_ with an optional `a` suffix" - ) - ); - assert_eq!( - crate::api::CompileError::from(error).stage(), - crate::api::CompilationStage::Input - ); + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let backend = BackendOptions { + target_arch: explicit.map(str::to_owned), + device_arch_hint: hint + .map(|raw| crate::options::DeviceArchHint::parse(raw.to_owned())), + ..BackendOptions::default() + }; + let validated = ValidatedBackendOptions::new(&backend).unwrap(); + assert!(validated.options().device_arch_hint.is_none()); + let rendered = render_live_cutlass_mlir( + &ctx, + module.get_operation(), + &validated, + CUTLASS_FULL_CUTE_MLIR22_PROFILE, + ) + .expect("the exporter must retain the validated device architecture"); + assert_eq!(rendered.target, expected); + assert!( + rendered + .text + .contains(&format!("\"gpu_arch\": \"{expected}\"")) + ); + } + } + + #[test] + fn cutlass_output_starts_with_a_comment_only_pinned_manifest() { + let manifest = ExportManifest { + profile: "test-profile".to_owned(), + exporter_version: "0.1.0".to_owned(), + pliron_revision: "pliron-rev".to_owned(), + interop_revision: cuda_oxide_mlir_export::INTEROP_REVISION.to_owned(), + consumer_revision: "consumer-rev".to_owned(), + llvm_major: 22, + gpu_arch: "sm_120a".to_owned(), + pointer_model: "opaque".to_owned(), + index_bits: 64, + }; + let module = cuda_oxide_mlir_export::MlirModule { + root: cuda_oxide_mlir_export::MlirOperation::new("builtin.module").unwrap(), + profile: "test-profile".to_owned(), + }; + + let text = render_cutlass_mlir(&manifest, &module).unwrap(); + let module_start = text.find("\"builtin.module\"").unwrap(); + assert!(text.starts_with("// cuda-oxide CUTLASS export manifest\n// {\n")); + assert!(text.contains("// \"consumer_revision\": \"consumer-rev\",\n")); + assert!( + text[..module_start] + .lines() + .all(|line| line.starts_with("//")) + ); + } + + /// Invalid target plans must produce an export diagnostic instead of a formatting panic. + #[test] + fn cutlass_rendering_failure_is_an_export_error() { + use cuda_oxide_mlir_export::{ + MlirModule, MlirOperation, MlirType, MlirValueId, MlirValueUse, + }; + + let mut root = MlirOperation::new("builtin.module").unwrap(); + root.operands.push(MlirValueUse { + id: MlirValueId(99), + ty: MlirType::Integer(32), + }); + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = MlirModule { + root, + profile: profile.name().into(), + }; + let error = render_cutlass_mlir(&profile.manifest(), &target).unwrap_err(); + assert!(matches!(error, PipelineError::Export(message) + if message.contains("CUTLASS MLIR rendering failed") + && message.contains("has no definition"))); + } + + #[test] + fn cutlass_output_must_not_be_overwritten_by_native_artifacts() { + let ll = Path::new("kernel.ll"); + let ptx = Path::new("kernel.ptx"); + let stale = [PathBuf::from("kernel.cubin")]; + let files = OutputFiles { + llvm_ir: ll, + ptx, + cubin: &stale[0], + mlir: Some(Path::new("kernel.mlir")), + stale_before_export: &stale, + }; + + assert!(validate_cutlass_mlir_output_path(Path::new("kernel.mlir"), &files).is_ok()); + for collision in [ll, ptx, stale[0].as_path()] { + let error = validate_cutlass_mlir_output_path(collision, &files).unwrap_err(); + assert!(error.to_string().contains("collides")); + } + } + + #[test] + fn cutlass_relative_output_cannot_alias_an_absolute_native_artifact() { + let current = std::env::current_dir().unwrap(); + let relative = Path::new("cuda-oxide-mlir-collision-check.ll"); + let ll = current.join(relative); + let ptx = current.join("cuda-oxide-mlir-collision-check.ptx"); + let cubin = current.join("cuda-oxide-mlir-collision-check.cubin"); + let files = OutputFiles { + llvm_ir: &ll, + ptx: &ptx, + cubin: &cubin, + mlir: Some(relative), + stale_before_export: &[], + }; + + let error = validate_cutlass_mlir_output_path(relative, &files).unwrap_err(); + assert!(error.to_string().contains("collides")); + } + + #[test] + fn cutlass_output_cannot_alias_a_native_artifact_through_parent_components() { + let root = tempfile::tempdir().unwrap(); + let child = root.path().join("child"); + std::fs::create_dir(&child).unwrap(); + let ll = root.path().join("kernel.ll"); + let ptx = root.path().join("kernel.ptx"); + let cubin = root.path().join("kernel.cubin"); + let alias = child.join("..").join("kernel.ll"); + let files = OutputFiles { + llvm_ir: &ll, + ptx: &ptx, + cubin: &cubin, + mlir: Some(&alias), + stale_before_export: &[], + }; + + let error = validate_cutlass_mlir_output_path(&alias, &files).unwrap_err(); + assert!(error.to_string().contains("collides")); + } + + #[cfg(unix)] + #[test] + fn cutlass_output_resolves_symlinked_parents_before_comparison() { + use std::os::unix::fs::symlink; + + let root = tempfile::tempdir().unwrap(); + let real = root.path().join("real"); + let alias_dir = root.path().join("alias"); + std::fs::create_dir(&real).unwrap(); + symlink(&real, &alias_dir).unwrap(); + let ll = real.join("kernel.ll"); + let ptx = real.join("kernel.ptx"); + let cubin = real.join("kernel.cubin"); + let alias = alias_dir.join("kernel.ll"); + let files = OutputFiles { + llvm_ir: &ll, + ptx: &ptx, + cubin: &cubin, + mlir: Some(&alias), + stale_before_export: &[], + }; + + let error = validate_cutlass_mlir_output_path(&alias, &files).unwrap_err(); + assert!(error.to_string().contains("collides")); + } + + #[cfg(unix)] + #[test] + fn cutlass_output_fails_closed_on_a_broken_symlink() { + use std::os::unix::fs::symlink; + + let root = tempfile::tempdir().unwrap(); + let broken = root.path().join("backend.mlir"); + symlink(root.path().join("missing-target"), &broken).unwrap(); + let ll = root.path().join("kernel.ll"); + let ptx = root.path().join("kernel.ptx"); + let cubin = root.path().join("kernel.cubin"); + let files = OutputFiles { + llvm_ir: &ll, + ptx: &ptx, + cubin: &cubin, + mlir: Some(&broken), + stale_before_export: &[], + }; + + let error = validate_cutlass_mlir_output_path(&broken, &files).unwrap_err(); + assert!(error.to_string().contains("could not safely normalize")); + } + + #[test] + fn invalid_device_hints_fail_at_the_shared_input_boundary() { + let cutlass_config = + CutlassBackendConfig::new(PathBuf::from("/unused/libCutlassCompiler.so")); + for cutlass_backend in [None, Some(&cutlass_config)] { + for raw in [ + "compute_120", + "compute_80", + "sm_", + "foo", + "sm_05", + "sm_120f", + "not-an-arch", + ] { + for nvvm in [false, true] { + for explicit in [None, Some("sm_120"), Some("foo")] { + let mut ctx = Context::new(); + let module = typed_mir_test_module(&mut ctx, &[]); + let backend = BackendOptions { + device_arch_hint: Some(crate::options::DeviceArchHint::parse( + raw.to_string(), + )), + target_arch: explicit.map(str::to_string), + ..BackendOptions::default() + }; + let request = ModulePipelineRequest::for_rust_pipeline( + &[], + nvvm, + &backend, + DebugKind::Off, + cutlass_backend, + &[], + OutputFiles { + llvm_ir: Path::new("unused.ll"), + ptx: Path::new("unused.ptx"), + cubin: Path::new("unused.cubin"), + mlir: None, + stale_before_export: &[], + }, + PipelineTrace::default(), + ); + let Err(error) = compile_translated_module(&mut ctx, module, &request) + else { + panic!("invalid hint must fail at input"); + }; + assert!(matches!(error, PipelineError::TargetSelection { .. })); + assert_eq!( + error.to_string(), + format!( + "invalid CUDA_OXIDE_DEVICE_ARCH `{raw}`: expected sm_ with an optional `a` suffix" + ) + ); + assert_eq!( + crate::api::CompileError::from(error).stage(), + crate::api::CompilationStage::Input + ); + } } } } diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs index 86a8302a07..5fa8e1d10d 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/prep.rs @@ -6,9 +6,49 @@ use crate::error::PipelineError; use crate::mir_pass_registry::{MirPassStage, SelectedMirPasses}; use crate::verify::verify_operation; +use dialect_cute::epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, + CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, + CuteTmaStoreTailOp, +}; +use dialect_cute::gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}; +use dialect_cute::gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, +}; +use dialect_cute::pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, +}; +use dialect_cute::scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, CuteSchedulerNew1dOp, + CuteWorkTileCoordinatesOp, +}; +use dialect_cute::sm100_ops::*; +use dialect_cute::smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, +}; +use dialect_cute::tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, +}; +use dialect_mir::attributes::MirCastKindAttr; +use dialect_mir::ops::{MirAllocaOp, MirCastOp}; +use dialect_mir::types::MirPtrType; use pliron::context::{Context, Ptr}; +use pliron::irbuild::{ + listener::Recorder, + rewriter::{IRRewriter, Rewriter}, +}; +use pliron::linked_list::ContainsLinkedList; +use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; +use pliron::r#type::Typed; /// Controls the reusable dialect-mir preparation stage. #[doc(hidden)] @@ -36,6 +76,10 @@ pub fn prepare_mir_module( preparation: MirPreparation<'_>, ) -> Result<(), PipelineError> { verify_operation(ctx, module, "module")?; + // Capture this before any optional preparation pass mutates the module. + // The extra promotion pass belongs only to the high-level tensor slice; + // ordinary modules keep their existing preparation pipeline. + let has_high_level_cute_operations = has_high_level_cute_operations(ctx, module); let has_pass_pipeline = preparation .mir_pass_pipeline .is_some_and(|pipeline| !pipeline.trim().is_empty()); @@ -45,6 +89,12 @@ pub fn prepare_mir_module( "optional MIR passes are unavailable with full variable debug info".to_string(), )); } + if has_high_level_cute_operations { + return Err(PipelineError::Lowering( + "high-level CuTe operations require MIR variable promotion; full variable debug information is not supported for this kernel" + .to_string(), + )); + } return Ok(()); } @@ -145,6 +195,29 @@ pub fn prepare_mir_module( "module post-borrowed-pointer-read-canonicalization", )?; + if has_high_level_cute_operations { + // Borrowing a local CuTe value (`&mut tensor`) can leave a `mir.cast` + // on its stack-slot pointer. The cast changes only the pointer kind, + // but prevents mem2reg from removing the slot. CuTe ghost values must + // stay in SSA, so remove these casts when every use is a load or store + // address. The next mem2reg pass can then remove the CuTe slot. + fold_ghost_slot_kind_only_casts(ctx, module); + verify_operation(ctx, module, "module post-ghost-slot-kind-cast-folding")?; + + // The first mem2reg pass can remove a temporary pointer slot that was + // the only reason another tensor slot looked as though its address + // escaped. That newly exposed inner slot is promotable only after the + // pointer cleanup above, so give this tensor pipeline one more pass. + pliron::opts::mem2reg::mem2reg(module, ctx, &mut analyses).map_err(|error| { + PipelineError::Verification { + name: "second mem2reg".to_string(), + message: error.disp(ctx).to_string(), + operation: None, + } + })?; + verify_operation(ctx, module, "module post-second-mem2reg")?; + } + mir_transforms::unroll::unroll_annotated_loops(module, ctx, &mut analyses).map_err( |error| PipelineError::Verification { name: "loop-unroll".to_string(), @@ -163,6 +236,183 @@ pub fn prepare_mir_module( ) } +fn has_high_level_cute_operations(ctx: &Context, root: Ptr) -> bool { + let mut pending = vec![root]; + while let Some(operation) = pending.pop() { + if Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::( + operation, ctx, + ) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::( + operation, ctx, + ) + || Operation::is_op::( + operation, ctx, + ) + || Operation::is_op::( + operation, ctx, + ) + || Operation::is_op::( + operation, ctx, + ) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + || Operation::is_op::(operation, ctx) + { + return true; + } + for region in operation.deref(ctx).regions() { + for block in region.deref(ctx).iter(ctx) { + pending.extend(block.deref(ctx).iter(ctx)); + } + } + } + false +} + +/// Fold pointer-kind-only `mir.cast`s on alloca slot pointers whose pointee +/// contains a CuTe ghost type, replacing the cast result with the slot +/// pointer. Only the pointer kind (and its implied authority metadata) is +/// erased; the pointee and address space must match exactly. This runs before +/// mem2reg so a `&mut tensor` reborrow does not keep the ghost slot alive. +fn fold_ghost_slot_kind_only_casts(ctx: &mut Context, root: Ptr) { + let mut foldable = Vec::new(); + let mut pending = vec![root]; + while let Some(operation) = pending.pop() { + for region in operation.deref(ctx).regions() { + for block in region.deref(ctx).iter(ctx) { + for child in block.deref(ctx).iter(ctx) { + pending.push(child); + if let Some(cast) = Operation::get_op::(child, ctx) + && is_ghost_slot_kind_only_cast(ctx, &cast) + { + foldable.push(child); + } + } + } + } + } + let mut rewriter = IRRewriter::::default(); + for cast in foldable { + let (base, result) = { + let operation = cast.deref(ctx); + (operation.get_operand(0), operation.get_result(0)) + }; + result.replace_all_uses_with(ctx, &base); + rewriter.erase_operation(ctx, cast); + } +} + +fn is_ghost_slot_kind_only_cast(ctx: &Context, cast: &MirCastOp) -> bool { + if cast + .get_attr_cast_kind(ctx) + .is_none_or(|kind| *kind != MirCastKindAttr::PtrToPtr) + { + return false; + } + let operation = cast.get_operation().deref(ctx); + if operation.get_num_operands() != 1 || operation.get_num_results() != 1 { + return false; + } + let base = operation.get_operand(0); + let Some(defining_op) = base.defining_op() else { + return false; + }; + if Operation::get_op::(defining_op, ctx).is_none() { + return false; + } + let base_ty = base.get_type(ctx); + let result_ty = operation.get_result(0).get_type(ctx); + let base_ty_ref = base_ty.deref(ctx); + let result_ty_ref = result_ty.deref(ctx); + let (Some(base_ptr), Some(result_ptr)) = ( + base_ty_ref.downcast_ref::(), + result_ty_ref.downcast_ref::(), + ) else { + return false; + }; + if base_ptr.pointee != result_ptr.pointee + || base_ptr.address_space != result_ptr.address_space + || !dialect_cute::verify::type_contains_cute_ghost(ctx, base_ptr.pointee) + { + return false; + } + // Fail closed: folding retypes the value at each use, which is only + // meaning-preserving where the pointer is consumed as a memory address + // (operand 0 of a load or store). A use that carries the pointer as data + // (stored value, call argument, branch operand, ...) must keep the cast. + operation.get_result(0).uses(ctx).iter().all(|r#use| { + let user = r#use.user_op(); + r#use.find_index(ctx) == 0 + && (Operation::get_op::(user, ctx).is_some() + || Operation::get_op::(user, ctx).is_some()) + }) +} + fn select_optional_mir_passes(spec: Option<&str>) -> Result { crate::mir_pass_registry::registry() .select(spec.unwrap_or_default()) @@ -206,8 +456,52 @@ fn run_optional_mir_passes( #[cfg(test)] mod tests { use super::*; + use dialect_cute::{ + attributes::{ + CutePipelineStateAttr, CuteTensorAccessAttr, CuteTensorRoleAttr, CuteTileGridAttr, + CuteTiledMmaPlanAttr, CuteTmaStorePipelineAttr, + }, + epilogue_ops::CuteTmaStoreAcquireOp, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + gemv_ops::CuteTensorMake2DOp, + layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}, + pipeline_ops::CutePipelineStateNewOp, + scheduler_ops::CuteSchedulerNew1dOp, + smem_mma_ops::CuteTiledMmaSliceOp, + tensor_ops::CuteTensorMakeOp, + }; + use dialect_mir::ops::{ + MirAllocaOp, MirFuncOp, MirLoadOp, MirReturnOp, MirStoreOp, MirUndefOp, + }; + use dialect_mir::types::{MirPtrType, address_space}; + use pliron::basic_block::BasicBlock; + use pliron::builtin::attributes::TypeAttr; + use pliron::builtin::op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}; use pliron::builtin::ops::ModuleOp; + use pliron::builtin::types::{FP32Type, FunctionType, IntegerType, Signedness}; use pliron::op::Op; + use pliron::r#type::TypeHandle; + + fn assert_high_level_cute_op_is_detected() { + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + Operation::new( + &mut ctx, + O::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } #[test] fn debug_mode_rejects_requested_mir_passes() { @@ -242,4 +536,359 @@ mod tests { .unwrap_err(); assert!(matches!(error, PipelineError::InvalidMirPassPipeline(_))); } + + #[test] + fn second_mem2reg_promotes_slot_exposed_by_pointer_slot_promotion() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let element: TypeHandle = FP32Type::get(&ctx).into(); + let inner_pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, element, true).into(); + let outer_pointer: TypeHandle = + MirPtrType::get_generic(&mut ctx, inner_pointer, true).into(); + let global_pointer: TypeHandle = + MirPtrType::get(&mut ctx, element, false, address_space::GLOBAL).into(); + let length_type: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let function_type = FunctionType::get(&ctx, vec![element], vec![element]); + let function = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function_op = MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(&mut ctx, "kernel".try_into().unwrap()); + module.append_operation(&mut ctx, function, 0); + + let region = function.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![element]); + entry.insert_at_back(region, &ctx); + let input = entry.deref(&ctx).get_argument(0); + + // This harmless view makes the fixture take the tensor-only second + // promotion path. The nested slots below reproduce the writable + // `TensorMut::store_at(&mut self)` shape that exposed the + // ordering bug. + let data = MirUndefOp::new(&mut ctx, global_pointer).get_operation(); + data.insert_at_back(entry, &ctx); + let data = data.deref(&ctx).get_result(0); + let length = MirUndefOp::new(&mut ctx, length_type).get_operation(); + length.insert_at_back(entry, &ctx); + let length = length.deref(&ctx).get_result(0); + CuteTensorMakeOp::new( + &mut ctx, + data, + length, + element, + element, + CuteTensorAccessAttr::ReadOnly, + 4, + ) + .get_operation() + .insert_at_back(entry, &ctx); + + let inner = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![inner_pointer], + vec![], + vec![], + 0, + ); + inner.insert_at_back(entry, &ctx); + let inner_slot = inner.deref(&ctx).get_result(0); + + let outer = Operation::new( + &mut ctx, + MirAllocaOp::get_concrete_op_info(), + vec![outer_pointer], + vec![], + vec![], + 0, + ); + outer.insert_at_back(entry, &ctx); + let outer_slot = outer.deref(&ctx).get_result(0); + + for (address, value) in [(outer_slot, inner_slot), (inner_slot, input)] { + Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![address, value], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + } + + let load_pointer = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![inner_pointer], + vec![outer_slot], + vec![], + 0, + ); + load_pointer.insert_at_back(entry, &ctx); + let loaded_pointer = load_pointer.deref(&ctx).get_result(0); + let load_value = Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![element], + vec![loaded_pointer], + vec![], + 0, + ); + load_value.insert_at_back(entry, &ctx); + let loaded_value = load_value.deref(&ctx).get_result(0); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![loaded_value], + vec![], + 0, + ) + .insert_at_back(entry, &ctx); + + prepare_mir_module( + &mut ctx, + module.get_operation(), + MirPreparation { + promote_and_unroll: true, + verbose: false, + mir_pass_pipeline: None, + }, + ) + .unwrap(); + + assert!( + entry + .deref(&ctx) + .iter(&ctx) + .all(|operation| !Operation::is_op::(operation, &ctx)), + "both the outer pointer slot and the inner value slot must be promoted" + ); + } + + #[test] + fn full_variable_debug_rejects_high_level_tensor_operations_early() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + let element: TypeHandle = FP32Type::get(&ctx).into(); + let pointer: TypeHandle = + MirPtrType::get(&mut ctx, element, false, address_space::GLOBAL).into(); + let length: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + + let pointer_value = MirUndefOp::new(&mut ctx, pointer).get_operation(); + pointer_value.insert_at_back(block, &ctx); + let pointer_value = pointer_value.deref(&ctx).get_result(0); + let length_value = MirUndefOp::new(&mut ctx, length).get_operation(); + length_value.insert_at_back(block, &ctx); + let length_value = length_value.deref(&ctx).get_result(0); + CuteTensorMakeOp::new( + &mut ctx, + pointer_value, + length_value, + element, + element, + CuteTensorAccessAttr::ReadOnly, + 4, + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = prepare_mir_module( + &mut ctx, + module_operation, + MirPreparation { + promote_and_unroll: false, + verbose: false, + mir_pass_pipeline: None, + }, + ) + .unwrap_err(); + let PipelineError::Lowering(message) = error else { + panic!("expected a full-debug lowering rejection"); + }; + assert!(message.contains("CuTe operations require MIR variable promotion")); + } + + #[test] + fn gemv_view_operations_enable_the_tensor_preparation_gate() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + let storage: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointer: TypeHandle = + MirPtrType::get(&mut ctx, storage, false, address_space::GLOBAL).into(); + let usize_type: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + + let pointer_value = MirUndefOp::new(&mut ctx, pointer).get_operation(); + pointer_value.insert_at_back(block, &ctx); + let pointer_value = pointer_value.deref(&ctx).get_result(0); + let size_value = MirUndefOp::new(&mut ctx, usize_type).get_operation(); + size_value.insert_at_back(block, &ctx); + let size_value = size_value.deref(&ctx).get_result(0); + CuteTensorMake2DOp::new_e2m1( + &mut ctx, + pointer_value, + size_value, + size_value, + size_value, + CuteTensorRoleAttr::Mkl, + 16, + ) + .get_operation() + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } + + #[test] + fn scheduler_operations_enable_the_cute_preparation_gate() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + CuteSchedulerNew1dOp::new(&mut ctx, CuteTileGridAttr::new(16, 16, 1)) + .get_operation() + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } + + #[test] + fn pipeline_operations_enable_the_cute_preparation_gate() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + CutePipelineStateNewOp::new(&mut ctx, CutePipelineStateAttr::producer(3)) + .get_operation() + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } + + #[test] + fn cache_aware_global_store_enables_the_cute_preparation_gate() { + assert_high_level_cute_op_is_detected::( + ); + } + + #[test] + fn sm100_attention_operations_enable_the_cute_preparation_gate() { + assert_high_level_cute_op_is_detected::( + ); + assert_high_level_cute_op_is_detected::( + ); + assert_high_level_cute_op_is_detected::< + dialect_cute::sm100_attention_ops::CuteSm100TmaLoadOp, + >(); + assert_high_level_cute_op_is_detected::< + dialect_cute::sm100_attention_ops::CuteSm100TmaStoreOp, + >(); + } + + #[test] + fn sm100_operations_enable_the_cute_preparation_gate() { + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + } + + #[test] + fn tma_operations_enable_the_cute_preparation_gate() { + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + assert_high_level_cute_op_is_detected::(); + } + + #[test] + fn shared_mma_operations_enable_the_cute_preparation_gate() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + let lane_type: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let lane = MirUndefOp::new(&mut ctx, lane_type).get_operation(); + lane.insert_at_back(block, &ctx); + let lane = lane.deref(&ctx).get_result(0); + + let inner: Layout = "(128,32):(32,1)".parse().unwrap(); + let placement = + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, inner, OffsetUnit::Elements).unwrap(); + CuteTiledMmaSliceOp::new( + &mut ctx, + lane, + CuteTiledMmaPlanAttr::mxf4_128x128x128(placement), + ) + .get_operation() + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } + + #[test] + fn epilogue_operations_enable_the_cute_preparation_gate() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let module_operation = module.get_operation(); + let region = module_operation.deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + CuteTmaStoreAcquireOp::new(&mut ctx, CuteTmaStorePipelineAttr::new(1)) + .get_operation() + .insert_at_back(block, &ctx); + + assert!(has_high_level_cute_operations(&ctx, module_operation)); + } } diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs index 8494b51982..34adae9b92 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs @@ -333,6 +333,7 @@ fn base_llc_args(target: &str) -> Vec { vec![ "-march=nvptx64".to_string(), format!("-mcpu={target}"), + "--nvptx-short-ptr".to_string(), DISABLE_BRANCH_FOLD.to_string(), DISABLE_BLOCK_PLACEMENT.to_string(), ] @@ -1222,6 +1223,7 @@ mod tests { [ "-march=nvptx64", "-mcpu=sm_90", + "--nvptx-short-ptr", "-disable-branch-fold", "-disable-block-placement", ] diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs new file mode 100644 index 0000000000..bdd9a05939 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/grid_constant.rs @@ -0,0 +1,190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Direct CUDA launch ABI regressions; the pinned compiler fixture is never launched. + +use super::*; + +const DESCRIPTOR_ATTRIBUTES: &str = + "{llvm.align = 64 : i64, llvm.byval = !llvm.array<128 x i8>, nvvm.grid_constant}"; + +fn descriptor_kind() -> KernelParameterKind { + KernelParameterKind::GridConstant { + size: 128, + alignment: 64, + } +} + +#[test] +fn parameterless_kernel_may_omit_parameter_size_but_not_claim_parameters() { + validate_kernel_parameter_abi("empty", &[], &[]).unwrap(); + validate_kernel_parameter_abi("empty", &[3, 0x19, 0, 0], &[]).unwrap(); + assert!(validate_kernel_parameter_abi("empty", &[3, 0x19, 8, 0], &[]).is_err()); + assert!( + validate_kernel_parameter_abi("nonempty", &[], &[KernelParameterKind::Scalar64]).is_err() + ); + assert!( + validate_kernel_parameter_abi("empty", ¶meter_records(&[(0, 0, 8, 0)], 8), &[]) + .is_err() + ); + assert!(validate_kernel_parameter_abi("empty", &[4, 0x17, 12, 0], &[]).is_err()); +} + +fn signature(attributes: &str) -> String { + format!( + r#""func.func"() <{{arg_attrs = [{{}}, {attributes}, {{}}], function_type = (i64, !llvm.ptr, !llvm.ptr) -> (), sym_name = "probe"}}> ({{"# + ) +} + +#[test] +fn grid_constant_signature_preserves_descriptor_bytes_and_alignment() { + let input = signature(DESCRIPTOR_ATTRIBUTES); + let abis = expected_kernel_abis(&input, &["probe".into()]).unwrap(); + assert_eq!( + abis["probe"], + vec![ + KernelParameterKind::Scalar64, + descriptor_kind(), + KernelParameterKind::GlobalPointer64, + ] + ); + // The renderer may spell a unit attribute explicitly or canonically elide it. + let explicit = input.replace("nvvm.grid_constant}", "nvvm.grid_constant = unit}"); + assert_eq!( + expected_kernel_abis(&explicit, &["probe".into()]).unwrap(), + abis + ); +} + +#[test] +fn grid_constant_signature_rejects_incomplete_or_malformed_attributes() { + for (original, replacement) in [ + (", nvvm.grid_constant", ""), + ("nvvm.grid_constant", "nvvm.grid_constant = false"), + ("64 : i64", "0 : i64"), + ("64 : i64", "3 : i64"), + ("64 : i64", "64 : i32"), + ("128 x i8", "0 x i8"), + ("128 x i8", "16384 x i8"), + ("128 x i8", "128 x i16"), + ( + "llvm.align = 64 : i64", + "llvm.align = 64 : i64, llvm.align = 64 : i64", + ), + ("nvvm.grid_constant", "nvvm.grid_constant, extra = unit"), + ("!llvm.ptr, !llvm.ptr", "i64, !llvm.ptr"), + ("[{}, {llvm.align", "[{llvm.align"), + ] { + let malformed = signature(DESCRIPTOR_ATTRIBUTES).replace(original, replacement); + assert!( + expected_kernel_abis(&malformed, &["probe".into()]).is_err(), + "accepted malformed ABI: {malformed}" + ); + } +} + +/// Explicit layouts avoid deriving fixture offsets from the code under test. +fn parameter_records(records: &[(u16, u16, u16, u32)], total: u16) -> Vec { + let mut info = Vec::new(); + for &(ordinal, offset, size, space) in records.iter().rev() { + info.extend_from_slice(&[4, 0x17, 12, 0]); + info.extend_from_slice(&0_u32.to_le_bytes()); + info.extend_from_slice(&ordinal.to_le_bytes()); + info.extend_from_slice(&offset.to_le_bytes()); + let flags = (u32::from(size) << 18) | (0x1f << 12) | (space << 8); + info.extend_from_slice(&flags.to_le_bytes()); + } + info.extend_from_slice(&[3, 0x19]); + info.extend_from_slice(&total.to_le_bytes()); + info +} + +#[test] +fn cubin_grid_constant_layout_checks_padding_size_space_and_total() { + let expected = [ + KernelParameterKind::Scalar64, + descriptor_kind(), + descriptor_kind(), + KernelParameterKind::GlobalPointer64, + KernelParameterKind::Scalar64, + ]; + let records = [ + (0, 0, 8, 0), + (1, 64, 128, 0), + (2, 192, 128, 0), + (3, 320, 8, 5), + (4, 328, 8, 0), + ]; + validate_kernel_parameter_abi("probe", ¶meter_records(&records, 336), &expected).unwrap(); + for replacement in [(1, 8, 128, 0), (1, 64, 8, 0), (1, 64, 128, 5)] { + let mut malformed = records; + malformed[1] = replacement; + let error = + validate_kernel_parameter_abi("probe", ¶meter_records(&malformed, 336), &expected) + .unwrap_err(); + assert!(error.contains("offset/size/space"), "{error}"); + } + assert!( + validate_kernel_parameter_abi("probe", ¶meter_records(&records, 280), &expected,) + .unwrap_err() + .contains("parameter block") + ); +} + +/// Run after installing the fingerprint-pinned compiler; no GPU is required: +/// `CUDA_OXIDE_TEST_CUTLASS_COMPILER=/absolute/path/libCutlassCompiler.so +/// cargo test -p cuda-oxide-codegen cutlass_grid_constant_tma_compiles -- --ignored` +#[test] +#[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] +fn cutlass_grid_constant_tma_compiles() { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mlir = format!( + r#"module attributes {{gpu.container_module}} {{ + gpu.module @kernels {{ + {} + ^bb0(%prefix: i64, %desc: !llvm.ptr, %out: !llvm.ptr): + llvm.inline_asm has_side_effects "prefetch.tensormap [$0];", "l" %desc : (!llvm.ptr) -> () + llvm.store %prefix, %out : i64, !llvm.ptr + func.return + }}) {{cute.kernel, gpu.kernel}} : () -> () + }} +}}"#, + signature(DESCRIPTOR_ATTRIBUTES), + ); + let temp = tempfile::tempdir().unwrap(); + let cubin_path = temp.path().join("probe.cubin"); + compile_to_cubin( + &CutlassBackendConfig::new(compiler_path.into()), + &mlir, + "sm_100a", + &["probe".into()], + &cubin_path, + None, + ) + .unwrap_or_else(|error| panic!("grid-constant compilation failed: {error}\n{mlir}")); + let cubin = std::fs::read(cubin_path).unwrap(); + // Keeping the address live in a TMA instruction causes plain `byval` to + // make a 128-byte local copy. With grid_constant the frame must stay empty. + let file = object::File::parse(cubin.as_slice()).unwrap(); + let info = file.section_by_name(".nv.info").unwrap().data().unwrap(); + let mut frame_sizes = Vec::new(); + let mut offset = 0; + while offset < info.len() { + let format = info[offset]; + let attribute = info[offset + 1]; + let size = usize::from(read_u16(info, offset + 2).unwrap()); + if attribute == 0x11 { + assert_eq!((format, size), (4, 8)); + frame_sizes.push(read_u32(info, offset + 8).unwrap()); + } + offset += 4 + if format == 4 { size } else { 0 }; + } + assert_eq!( + frame_sizes, + [0], + "grid-constant descriptor escaped into local memory" + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/primitives.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/primitives.rs new file mode 100644 index 0000000000..0591cb3ba3 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/primitives.rs @@ -0,0 +1,125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile actual Rust-imported primitives IR through production image validation. +use super::*; +use object::ObjectSymbol; + +#[test] +#[ignore = "requires source-built primitives evidence and pinned CUTLASS 4.7; no GPU"] +fn cutlass_primitives_grid_constant_abi() { + let evidence = std::env::var_os("CUDA_OXIDE_PRIMITIVES_EVIDENCE") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/tmp/primitives-api-evidence")); + let mlir = std::fs::read_to_string(evidence.join("primitives.mlir")) + .expect("run scripts/check-primitives-api-compile.py first"); + let names: Vec = [ + "primitives_grid_constant", + "primitives_lifecycle", + "primitives_tma", + "primitives_sync", + "primitives_mma", + ] + .into_iter() + .map(str::to_owned) + .collect(); + let abis = expected_kernel_abis(&mlir, &names).unwrap(); + assert_eq!( + abis[&names[0]][0], + KernelParameterKind::GridConstant { + size: 128, + alignment: 64 + } + ); + assert_eq!( + &abis["primitives_tma"][..4], + &[KernelParameterKind::GridConstant { + size: 128, + alignment: 64 + }; 4] + ); + assert!(abis["primitives_lifecycle"].is_empty()); + assert!(!mlir.contains("\"cute.") && !mlir.contains("!cute.")); + let mut config = CutlassBackendConfig::new( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set pinned compiler library") + .into(), + ); + config.profile = cuda_oxide_mlir_export::CutlassPrimitivesMlir22::NAME.into(); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("primitives.cubin"); + compile_to_cubin(&config, &mlir, "sm_100a", &names, &path, None).unwrap(); + let cubin = std::fs::read(path).unwrap(); + let file = object::File::parse(cubin.as_slice()).unwrap(); + let symbol = file + .symbols() + .find(|s| s.name().ok() == Some("primitives_grid_constant")) + .unwrap() + .index() + .0 as u32; + let info = file.section_by_name(".nv.info").unwrap().data().unwrap(); + let mut offset = 0; + let mut frames = vec![]; + while offset < info.len() { + let format = info[offset]; + let attribute = info[offset + 1]; + let size = usize::from(read_u16(info, offset + 2).unwrap()); + if attribute == 0x11 && read_u32(info, offset + 4).unwrap() == symbol { + assert_eq!((format, size), (4, 8)); + frames.push(read_u32(info, offset + 8).unwrap()); + } + offset += 4 + if format == 4 { size } else { 0 }; + } + assert_eq!( + frames, + [0], + "typed grid-constant descriptor was copied to a local frame" + ); +} + +#[test] +#[ignore = "requires pinned CUTLASS 4.7; no GPU"] +fn cutlass_primitives_kernel_discovery() { + let mut config = CutlassBackendConfig::new( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set pinned compiler library") + .into(), + ); + config.profile = cuda_oxide_mlir_export::CutlassPrimitivesMlir22::NAME.into(); + let temp = tempfile::tempdir().unwrap(); + for (label, attributes) in [ + ("marked", "cute.kernel, gpu.kernel"), + ("gpu_only", "gpu.kernel"), + ] { + let mlir = format!( + r#"module attributes {{gpu.container_module}} {{ + gpu.module @kernels {{ + "func.func"() <{{function_type = (!llvm.ptr) -> (), sym_name = "primitives_discovery"}}> ({{ + ^bb0(%out: !llvm.ptr): + %value = arith.constant 19 : i64 + llvm.store %value, %out : i64, !llvm.ptr + func.return + }}) {{{attributes}}} : () -> () + }} +}}"# + ); + let result = compile_to_cubin( + &config, + &mlir, + "sm_100a", + &["primitives_discovery".into()], + &temp.path().join(format!("{label}.cubin")), + None, + ); + if label == "marked" { + result.unwrap(); + } else { + let error = result.unwrap_err().to_string(); + assert!( + error.contains("missing kernel code section .text.primitives_discovery"), + "{error}" + ); + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs new file mode 100644 index 0000000000..e336c4056b --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/sm100.rs @@ -0,0 +1,197 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile the actual semantic SM100 mapping family through the pinned consumer. +//! This fixture is never launched: warp/cluster participation is tested by the +//! example's GPU smoke tests. In particular, successful MLIR parsing alone does +//! not validate CUTLASS 4.7's integer-register ABI for native TMEM loads. + +use super::{ + CUTLASS_COMPILER_LIBRARY_SHA256, CUTLASS_PRECOMPILED_PIPELINE, extract_kernels_binary, +}; +use crate::cutlass_compiler::CutlassCompilerLibrary; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_cute::{ + attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}, + layout::{ComposedLayout, OffsetUnit, Swizzle}, + sm100_ops::*, +}; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +fn shared_layout(rows: u32, columns: u32, bits: u32) -> CuteComposedLayoutAttr { + CuteComposedLayoutAttr( + ComposedLayout::new( + Swizzle::new(bits, 3, 3), + 0, + format!("({rows},{columns}):({columns},1)").parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ) +} +fn append(ctx: &Context, block: Ptr, op: O) { + op.get_operation().insert_at_back(block, ctx); +} + +#[test] +#[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] +fn cutlass_sm100_semantic_pack_compiles() { + let path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let u32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let f32: TypeHandle = FP32Type::get(&ctx).into(); + let f16: TypeHandle = MirFP16Type::get(&ctx).into(); + let halfptr: TypeHandle = MirPtrType::get(&mut ctx, f16, true, address_space::SHARED).into(); + let barptr: TypeHandle = MirPtrType::get(&mut ctx, u64, true, address_space::SHARED).into(); + let tokenptr: TypeHandle = MirPtrType::get(&mut ctx, u32, true, address_space::SHARED).into(); + let descptr: TypeHandle = MirPtrType::get(&mut ctx, u64, false, address_space::GENERIC).into(); + // TMEM, A/B/C shared tiles, predicate, full/empty barriers, stage/phase, + // bias, allocation token, descriptor, and cluster rank. + let args = vec![ + u32, halfptr, halfptr, halfptr, i1, barptr, barptr, u32, u32, f32, tokenptr, descptr, u32, + ]; + let module = ModuleOp::new(&mut ctx, "sm100_semantic_smoke".try_into().unwrap()); + let ft = FunctionType::get(&ctx, args.clone(), vec![]); + let func = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, func, TypeAttr::new(ft.into())); + function.set_symbol_name(&mut ctx, "sm100_semantic_probe".try_into().unwrap()); + func.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, func, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(func.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + let tmem = CuteSm100TmemPlanAttr { + columns: 512, + cta_group: 2, + }; + let mma = CuteSm100MmaPlanAttr { + m: 256, + n0: 192, + n1: 160, + k: 64, + tmem_columns: 512, + cta_group: 2, + collector: CuteSm100CollectorAttr::AFillLastUse, + a_layout: shared_layout(128, 64, 3), + b0_layout: shared_layout(96, 64, 3), + b1_layout: shared_layout(80, 64, 3), + }; + let pipeline = CuteSm100PipelinePlanAttr { + stages: 5, + transaction_bytes: 77_824, + consumer_arrivals: 1, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + }; + let accumulator = CuteSm100PipelinePlanAttr { + stages: 1, + transaction_bytes: 0, + consumer_arrivals: 8, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + }; + let epilogue = CuteSm100EpiloguePlanAttr { + layout: shared_layout(128, 32, 2), + rows: 128, + columns: 32, + }; + let store = CuteTmaStorePipelineAttr { stages: 2 }; + macro_rules! emit {($ty:ty,[$($operand:expr),*],$plan:expr)=>{{ + let op=<$ty>::new(&mut ctx,vec![$($operand),*],$plan);append(&ctx,block,op); + }};} + emit!(CuteSm100TmemAllocOp, [a[10]], tmem); + emit!(CuteSm100PipelineInitOp, [a[5], a[6]], pipeline); + emit!(CuteSm100AccumulatorInitOp, [a[6], a[5]], accumulator); + emit!(CuteSm100PipelineAcquireOp, [a[6], a[7], a[8]], pipeline); + emit!(CuteSm100PipelineExpectOp, [a[5], a[7]], pipeline); + for (ptr, layout) in [ + (a[1], mma.a_layout.clone()), + (a[2], mma.b0_layout.clone()), + (a[3], mma.b1_layout.clone()), + ] { + let copy = CuteSm100ClusterCopyPlanAttr { + layout, + cta_group: 2, + leader_rank: 0, + }; + emit!( + CuteSm100ClusterTmaLoadOp, + [a[11], ptr, a[7], a[7], a[5], a[12]], + copy + ); + } + emit!(CuteSm100PipelineWaitOp, [a[5], a[7], a[8]], pipeline); + emit!(CuteSm100AccumulatorAcquireOp, [a[6], a[8]], accumulator); + emit!(CuteSm100TiledMmaOp, [a[0], a[1], a[2], a[3], a[4]], mma); + emit!(CuteSm100PipelineReleaseOp, [a[6], a[7]], pipeline); + emit!(CuteSm100AccumulatorCommitOp, [a[5]], accumulator); + emit!(CuteSm100AccumulatorWaitOp, [a[5], a[8]], accumulator); + emit!( + CuteSm100TmemEpilogueOp, + [a[0], a[1], a[7], a[9]], + epilogue.clone() + ); + emit!(CuteSm100TmaStoreOp, [a[11], a[1], a[7], a[7]], epilogue); + emit!(CuteSm100StoreCommitOp, [], store); + emit!(CuteSm100StoreAcquireOp, [], store); + emit!(CuteSm100StoreTailOp, [], store); + emit!(CuteSm100AccumulatorReleaseOp, [a[6]], accumulator); + emit!(CuteSm100TmemDeallocOp, [a[0]], tmem); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let mlir = cuda_oxide_mlir_export::render_module(&target).unwrap(); + assert!(mlir.contains("cute_nvgpu.arch.copy.SM100.tmem_load")); + assert!(mlir.contains("vector<32xi32>")); + assert_eq!(mlir.matches("cute_nvgpu.make_umma_smem_desc").count(), 3); + let library = CutlassCompilerLibrary::load(path, CUTLASS_COMPILER_LIBRARY_SHA256).unwrap(); + let object = library + .compile_precompiled_mlir_to_object(&mlir, "sm_100a", CUTLASS_PRECOMPILED_PIPELINE) + .unwrap_or_else(|error| panic!("semantic SM100 MLIR failed compilation: {error}\n{mlir}")); + let image = extract_kernels_binary(&object.elf).unwrap(); + assert!(image.len() > 64); +} diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs new file mode 100644 index 0000000000..7f29a8baf6 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/tests/cutlass/tcgen05.rs @@ -0,0 +1,1328 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compiler-only regressions for the TCGen05 and cluster/TMA translation boundary. +//! This fixture is never launched; collective participation is the example's +//! runtime contract, which requires separate testing on an SM100 device. + +use super::{ + CUTLASS_COMPILER_LIBRARY_SHA256, CUTLASS_PRECOMPILED_PIPELINE, extract_kernels_binary, +}; +use crate::cutlass_compiler::CutlassCompilerLibrary; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{MirCastOp, MirConstantOp, MirFuncOp, MirReturnOp, MirStoreOp}, + types::{MirPointerKind, MirPtrType, address_space}, +}; +use dialect_nvvm::ops::{ + BarrierCtaSyncCountOp, ClusterBarrierModeAttr, ClusterBarrierOp, CpAsyncBulkCommitGroupOp, + CpAsyncBulkTensorG2sTile2dMulticastCg2Op, CpAsyncBulkTensorS2gTile2dOp, + CpAsyncBulkWaitGroupReadOp, CvtaGenericToSharedOffsetOp, ElectSyncOp, + FenceAcquireSharedClusterClusterOp, FenceMbarrierInitReleaseClusterOp, + FenceProxyAsyncSharedCtaOp, MapaSharedClusterOp, MbarrierArriveClusterOp, + MbarrierArriveClusterRelaxedOp, MbarrierArriveExpectTxClusterOp, MbarrierInitSharedOp, + MbarrierTryWaitParityClusterOp, MbarrierTryWaitParityClusterRelaxedOp, PrefetchTensorMapOp, + Tcgen05AllocCg2Op, Tcgen05CommitMulticastCg2Op, Tcgen05DeallocCg2Op, + Tcgen05FenceAfterThreadSyncOp, Tcgen05FenceBeforeThreadSyncOp, Tcgen05Ld32x32bX32RawOp, + Tcgen05LoadWaitOp, Tcgen05MmaCollectorAAttr, Tcgen05MmaCtaGroupAttr, Tcgen05MmaFormAttr, + Tcgen05MmaKindAttr, Tcgen05MmaOp, Tcgen05RelinquishAllocPermitCg2Op, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::{TypeHandle, TypedHandle}, + utils::apint::APInt, + value::Value, +}; +use std::num::NonZeroUsize; + +fn append( + ctx: &mut Context, + block: Ptr, + operands: Vec, + results: Vec, + marker: Option<&str>, +) -> Ptr { + let op = Operation::new(ctx, O::get_concrete_op_info(), results, operands, vec![], 0); + if let Some(marker) = marker { + op.deref_mut(ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + } + op.insert_at_back(block, ctx); + op +} + +/// Run explicitly after installing the fingerprint-pinned CUTLASS compiler: +/// +/// ```text +/// CUDA_OXIDE_TEST_CUTLASS_COMPILER=/absolute/path/libCutlassCompiler.so \ +/// cargo test -p cuda-oxide-codegen cutlass_tcgen05_pack_compiles -- --ignored +/// ``` +#[test] +#[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] +fn cutlass_tcgen05_pack_compiles() { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "tcgen05_smoke".try_into().unwrap()); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let i16: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let i64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let pointer: TypeHandle = MirPtrType::get(&mut ctx, i32, true, address_space::GENERIC).into(); + let arguments = vec![pointer, i32, i64, i64, i32, i32, i1, i16, pointer]; + let function_type = FunctionType::get(&ctx, arguments.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, "tcgen05_probe".try_into().unwrap()); + function_op.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, arguments); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + append::(&mut ctx, block, vec![a[0]], vec![i64], None); + append::(&mut ctx, block, vec![a[0], a[5]], vec![], Some("v1:i0359")); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0361")); + for collector in [ + Tcgen05MmaCollectorAAttr::Fill, + Tcgen05MmaCollectorAAttr::LastUse, + ] { + let op = append::( + &mut ctx, + block, + vec![a[1], a[2], a[3], a[4], a[6]], + vec![], + Some("v1:i0763"), + ); + let mma = Tcgen05MmaOp::new(op); + mma.set_attr_nvvm_tcgen05_mma_form(&ctx, Tcgen05MmaFormAttr::Shared); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&ctx, Tcgen05MmaCtaGroupAttr::Cg2); + mma.set_attr_nvvm_tcgen05_mma_collector_a(&ctx, collector); + } + append::( + &mut ctx, + block, + vec![a[0], a[7]], + vec![], + Some("v1:i0365"), + ); + let load = append::( + &mut ctx, + block, + vec![a[1]], + vec![i32; 32], + Some("v1:i0664"), + ); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0357")); + let last_register = load.deref(&ctx).get_result(31); + append::(&mut ctx, block, vec![a[8], last_register], vec![], None); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0346")); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0347")); + append::(&mut ctx, block, vec![a[1], a[5]], vec![], Some("v1:i0360")); + append::(&mut ctx, block, vec![], vec![], None); + + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let mlir = cuda_oxide_mlir_export::render_module(&target).unwrap(); + let library = + CutlassCompilerLibrary::load(compiler_path, CUTLASS_COMPILER_LIBRARY_SHA256).unwrap(); + let object = library + .compile_precompiled_mlir_to_object(&mlir, "sm_100a", CUTLASS_PRECOMPILED_PIPELINE) + .unwrap_or_else(|error| { + panic!("TCGen05 translated MLIR failed compilation: {error}\n{mlir}") + }); + let image = extract_kernels_binary(&object.elf).unwrap(); + assert!( + image.len() > 64, + "compiler did not produce a nonempty device image" + ); +} + +#[test] +#[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] +fn cutlass_attention_tmem_pack_compiles() { + // Group mode must be consistent throughout a kernel. Probe separate modules. + for group in [1, 2] { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "tcgen05_smoke".try_into().unwrap()); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let i16: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let i64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let pointer: TypeHandle = + MirPtrType::get(&mut ctx, i32, true, address_space::GENERIC).into(); + let arguments = vec![pointer, i32, i64, i64, i32, i32, i1, i16, pointer]; + let function_type = FunctionType::get(&ctx, arguments.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, "tcgen05_probe".try_into().unwrap()); + function_op.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, arguments); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + if group == 1 { + append::( + &mut ctx, + block, + vec![a[0], a[5]], + vec![], + Some("v1:i0343"), + ); + append::( + &mut ctx, + block, + vec![], + vec![], + Some("v1:i0345"), + ); + } else { + append::( + &mut ctx, + block, + vec![a[0], a[5]], + vec![], + Some("v1:i0359"), + ); + append::( + &mut ctx, + block, + vec![], + vec![], + Some("v1:i0361"), + ); + } + for tensor in [false, true] { + let op = append::( + &mut ctx, + block, + vec![a[1], if tensor { a[1] } else { a[2] }, a[3], a[4], a[6]], + vec![], + Some(if tensor { "v1:i0764" } else { "v1:i0763" }), + ); + let mma = Tcgen05MmaOp::new(op); + mma.set_attr_nvvm_tcgen05_mma_form( + &ctx, + if tensor { + Tcgen05MmaFormAttr::Tensor + } else { + Tcgen05MmaFormAttr::Shared + }, + ); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group( + &ctx, + if group == 1 { + Tcgen05MmaCtaGroupAttr::Cg1 + } else { + Tcgen05MmaCtaGroupAttr::Cg2 + }, + ); + mma.set_attr_nvvm_tcgen05_mma_collector_a(&ctx, Tcgen05MmaCollectorAAttr::Discard); + } + if group == 1 { + append::( + &mut ctx, + block, + vec![a[0]], + vec![], + Some("v1:i0348"), + ); + } else { + append::( + &mut ctx, + block, + vec![a[0]], + vec![], + Some("v1:i0363"), + ); + } + macro_rules! memory_pair { + ($load:ty,$store:ty,$width:literal,$load_marker:literal,$store_marker:literal) => {{ + let ld = append::<$load>( + &mut ctx, + block, + vec![a[1]], + vec![i32; $width], + Some($load_marker), + ); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0357")); + let regs = ld.deref(&ctx).results().collect::>(); + // Retained scalar uses exercise extraction independently of the + // native load/store vector round trip. + for &value in ®s { + let store = + append::(&mut ctx, block, vec![a[8], value], vec![], None); + MirStoreOp::new(store).set_volatile(&mut ctx, true); + } + append::<$store>( + &mut ctx, + block, + std::iter::once(a[1]).chain(regs).collect(), + vec![], + Some($store_marker), + ); + }}; + } + use dialect_nvvm::ops::*; + memory_pair!( + Tcgen05Ld32x32bX1RawOp, + Tcgen05St32x32bX1RawOp, + 1, + "v1:i0654", + "v1:i0712" + ); + memory_pair!( + Tcgen05Ld32x32bX1Pack16Op, + Tcgen05St32x32bX1Unpack16Op, + 1, + "v1:i0655", + "v1:i0713" + ); + memory_pair!( + Tcgen05Ld32x32bX2RawOp, + Tcgen05St32x32bX2RawOp, + 2, + "v1:i0656", + "v1:i0714" + ); + memory_pair!( + Tcgen05Ld32x32bX2Pack16Op, + Tcgen05St32x32bX2Unpack16Op, + 2, + "v1:i0657", + "v1:i0715" + ); + memory_pair!( + Tcgen05Ld32x32bX4RawOp, + Tcgen05St32x32bX4RawOp, + 4, + "v1:i0658", + "v1:i0716" + ); + memory_pair!( + Tcgen05Ld32x32bX4Pack16Op, + Tcgen05St32x32bX4Unpack16Op, + 4, + "v1:i0659", + "v1:i0717" + ); + memory_pair!( + Tcgen05Ld32x32bX8RawOp, + Tcgen05St32x32bX8RawOp, + 8, + "v1:i0660", + "v1:i0718" + ); + memory_pair!( + Tcgen05Ld32x32bX8Pack16Op, + Tcgen05St32x32bX8Unpack16Op, + 8, + "v1:i0661", + "v1:i0719" + ); + memory_pair!( + Tcgen05Ld32x32bX16RawOp, + Tcgen05St32x32bX16RawOp, + 16, + "v1:i0662", + "v1:i0720" + ); + memory_pair!( + Tcgen05Ld32x32bX16Pack16Op, + Tcgen05St32x32bX16Unpack16Op, + 16, + "v1:i0663", + "v1:i0721" + ); + memory_pair!( + Tcgen05Ld32x32bX32RawOp, + Tcgen05St32x32bX32RawOp, + 32, + "v1:i0664", + "v1:i0722" + ); + memory_pair!( + Tcgen05Ld32x32bX32Pack16Op, + Tcgen05St32x32bX32Unpack16Op, + 32, + "v1:i0665", + "v1:i0723" + ); + memory_pair!( + Tcgen05Ld32x32bX64RawOp, + Tcgen05St32x32bX64RawOp, + 64, + "v1:i0666", + "v1:i0724" + ); + memory_pair!( + Tcgen05Ld32x32bX64Pack16Op, + Tcgen05St32x32bX64Unpack16Op, + 64, + "v1:i0667", + "v1:i0725" + ); + memory_pair!( + Tcgen05Ld32x32bX128RawOp, + Tcgen05St32x32bX128RawOp, + 128, + "v1:i0668", + "v1:i0726" + ); + memory_pair!( + Tcgen05Ld32x32bX128Pack16Op, + Tcgen05St32x32bX128Unpack16Op, + 128, + "v1:i0669", + "v1:i0727" + ); + append::( + &mut ctx, + block, + vec![], + vec![], + Some("v1:i0358"), + ); + if group == 1 { + append::( + &mut ctx, + block, + vec![a[1], a[5]], + vec![], + Some("v1:i0344"), + ); + } else { + append::( + &mut ctx, + block, + vec![a[1], a[5]], + vec![], + Some("v1:i0360"), + ); + } + append::(&mut ctx, block, vec![], vec![], None); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let mlir = cuda_oxide_mlir_export::render_module(&target).unwrap(); + let library = + CutlassCompilerLibrary::load(compiler_path, CUTLASS_COMPILER_LIBRARY_SHA256).unwrap(); + let object = library + .compile_precompiled_mlir_to_object(&mlir, "sm_100a", CUTLASS_PRECOMPILED_PIPELINE) + .unwrap_or_else(|error| { + panic!("TCGen05 translated MLIR failed compilation: {error}\n{mlir}") + }); + let image = extract_kernels_binary(&object.elf).unwrap(); + assert!( + image.len() > 64, + "compiler did not produce a nonempty device image" + ); + } +} + +/// Compile the actual exported cluster/TMA mappings against the pinned consumer. +/// In particular, CUTLASS 4.7's two G2S intrinsic paths require different pointer +/// address spaces; MLIR parsing alone cannot catch the LLVM intrinsic ABI error. +#[test] +#[ignore = "requires the installed official CUTLASS 4.7 compiler; does not require a GPU"] +#[allow(clippy::disallowed_methods)] // Synthetic Rust raw pointer provenance for mapa's source contract. +fn cutlass_cluster_tma_pack_compiles() { + let compiler_path = std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER") + .expect("set CUDA_OXIDE_TEST_CUTLASS_COMPILER to the pinned compiler library"); + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "cluster_tma_smoke".try_into().unwrap()); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let u16 = IntegerType::get(&ctx, 16, Signedness::Unsigned); + let u32 = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64 = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let pointer: TypeHandle = MirPtrType::get_with_kind( + &mut ctx, + u32.into(), + true, + address_space::GENERIC, + MirPointerKind::RawMut, + ) + .into(); + let predicate_pointer: TypeHandle = MirPtrType::get_with_kind( + &mut ctx, + i1, + true, + address_space::GENERIC, + MirPointerKind::RawMut, + ) + .into(); + // Shared source, tensor descriptor, local barrier, lane result, wait result. + let arguments = vec![pointer, pointer, pointer, pointer, predicate_pointer]; + let function_type = FunctionType::get(&ctx, arguments.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, "cluster_tma_probe".try_into().unwrap()); + function_op.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, arguments); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + + let zero = integer_constant(&mut ctx, block, u32, 0, 32); + let one = integer_constant(&mut ctx, block, u32, 1, 32); + let threads = integer_constant(&mut ctx, block, u32, 128, 32); + let bytes = integer_constant(&mut ctx, block, u32, 77824, 32); + let warp_mask = integer_constant(&mut ctx, block, u32, u32::MAX.into(), 32); + let multicast = integer_constant(&mut ctx, block, u16, 3, 16); + let cache_hint = integer_constant(&mut ctx, block, u64, 0, 64); + + append::(&mut ctx, block, vec![a[1]], vec![], Some("v1:i0887")); + append::(&mut ctx, block, vec![a[2], one], vec![], Some("v1:i0097")); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0313")); + for (mode, marker) in [ + (ClusterBarrierModeAttr::ArriveRelaxed, "v1:i0279"), + (ClusterBarrierModeAttr::Wait, "v1:i0281"), + ] { + let op = append::(&mut ctx, block, vec![], vec![], Some(marker)); + ClusterBarrierOp::new(op).set_attr_nvvm_cluster_barrier_mode(&ctx, mode); + } + append::( + &mut ctx, + block, + vec![a[2], bytes], + vec![u64.into()], + Some("v1:i0307"), + ); + let mapped = MapaSharedClusterOp::build(&mut ctx, a[2], zero); + mapped.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0320".into()), + ); + mapped.insert_at_back(block, &ctx); + let remote_barrier = mapped.deref(&ctx).get_result(0); + // Native semantic order: dst, barrier, descriptor, x, y, mask, cache hint. + append::( + &mut ctx, + block, + vec![ + a[0], + remote_barrier, + a[1], + zero, + zero, + multicast, + cache_hint, + ], + vec![], + Some("v1:i0331"), + ); + let waited = append::( + &mut ctx, + block, + vec![a[2], zero], + vec![i1], + Some("v1:i0311"), + ); + let ready = waited.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![a[4], ready], vec![], None); + let election = append::( + &mut ctx, + block, + vec![warp_mask], + vec![u32.into(), i1], + Some("v1:i0367"), + ); + let lane = election.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![a[3], lane], vec![], None); + append::( + &mut ctx, + block, + vec![one, threads], + vec![], + Some("v1:i0883"), + ); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0312")); + append::( + &mut ctx, + block, + vec![a[0], a[1], zero, zero], + vec![], + Some("v1:i0336"), + ); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0340")); + append::(&mut ctx, block, vec![zero], vec![], Some("v1:i0342")); + let remote_address = append::( + &mut ctx, + block, + vec![remote_barrier], + vec![u64.into()], + None, + ); + MirCastOp::new(remote_address).set_attr_cast_kind(&ctx, MirCastKindAttr::PointerExposeAddress); + let address = remote_address.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![address], vec![], Some("v1:i0308")); + append::( + &mut ctx, + block, + vec![address], + vec![], + Some("v1:i1040"), + ); + let relaxed_wait = append::( + &mut ctx, + block, + vec![a[2], zero], + vec![i1], + Some("v1:i1041"), + ); + let ready = relaxed_wait.deref(&ctx).get_result(0); + let store = append::(&mut ctx, block, vec![a[4], ready], vec![], None); + MirStoreOp::new(store).set_volatile(&mut ctx, true); + append::(&mut ctx, block, vec![], vec![], Some("v1:i1042")); + append::(&mut ctx, block, vec![], vec![], None); + + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let mlir = cuda_oxide_mlir_export::render_module(&target).unwrap(); + let library = + CutlassCompilerLibrary::load(compiler_path, CUTLASS_COMPILER_LIBRARY_SHA256).unwrap(); + let dump = tempfile::tempdir().unwrap(); + let ptx_prefix = dump.path().join("cluster_barriers"); + let pipeline = CUTLASS_PRECOMPILED_PIPELINE.replace( + "cubin-format=bin", + &format!("cubin-format=bin dump-ptx-path={}", ptx_prefix.display()), + ); + let object = library + .compile_precompiled_mlir_to_object(&mlir, "sm_100a", &pipeline) + .unwrap_or_else(|error| { + panic!("Cluster/TMA translated MLIR failed compilation: {error}\n{mlir}") + }); + let ptx = std::fs::read_to_string(ptx_prefix.with_extension("sm_100a.ptx")).unwrap(); + let prefetch_count = ptx + .lines() + .filter(|line| line.split_whitespace().next() == Some("prefetch.tensormap")) + .count(); + assert_eq!(prefetch_count, 1, "{ptx}"); + assert!(mlir.contains("\"nvvm.prefetch\""), "{mlir}"); + assert!( + ptx.contains("mbarrier.arrive.release.cluster.shared::cluster.b64"), + "{ptx}" + ); + assert!( + ptx.contains("mbarrier.arrive.relaxed.cluster.shared::cluster.b64"), + "{ptx}" + ); + assert!( + ptx.contains("mbarrier.try_wait.parity.relaxed.cluster"), + "{ptx}" + ); + assert!( + ptx.contains("fence.acquire.sync_restrict::shared::cluster.cluster;"), + "{ptx}" + ); + assert!(!mlir.contains("inline_ptx"), "{mlir}"); + assert!( + extract_kernels_binary(&object.elf).unwrap().len() > 64, + "compiler did not produce a nonempty device image" + ); +} + +fn integer_constant( + ctx: &mut Context, + block: Ptr, + ty: TypedHandle, + value: u64, + bits: usize, +) -> Value { + let op = append::(ctx, block, vec![], vec![ty.into()], None); + MirConstantOp::new(op).set_attr_value( + ctx, + IntegerAttr::new(ty, APInt::from_u64(value, NonZeroUsize::new(bits).unwrap())), + ); + op.deref(ctx).get_result(0) +} + +mod cutlass_attention_semantic_pack_compiles_module { + use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; + use dialect_cute::sm100_attention_ops::*; + use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType}, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, + }; + + #[test] + #[ignore = "requires the pinned CUTLASS compiler; no GPU"] + fn cutlass_attention_semantic_pack_compiles() { + for group in [1, 2] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "attention".try_into().unwrap()); + let f16 = MirFP16Type::get(&ctx).into(); + let u32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let b: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get(&mut ctx, f16, true, 3).into(); + let desc: TypeHandle = MirPtrType::get(&mut ctx, u64, false, 0).into(); + let barrier: TypeHandle = MirPtrType::get(&mut ctx, u64, true, 3).into(); + let args = vec![u32, shared, shared, b, u32, desc, barrier, u32, u32]; + let ty = FunctionType::get(&ctx, args.clone(), vec![]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, f, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "attention_probe".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(f.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + let plan = CuteSm100MmaAttr { + m: 128, + n: 128, + k: 128, + cta_group: group, + a_mn: false, + b_mn: false, + a_layout: shared_layout(128 / group, 128, false).unwrap(), + b_layout: shared_layout(128 / group, 128, false).unwrap(), + }; + CuteSm100MmaSsOp::new( + &mut ctx, + vec![a[0], a[1], a[2], a[3], a[7], a[8]], + plan.clone(), + ) + .get_operation() + .insert_at_back(block, &ctx); + if group == 1 { + let mut pv = plan.clone(); + pv.b_mn = true; + pv.b_layout = shared_layout(128, 128, true).unwrap(); + CuteSm100MmaTsOp::new(&mut ctx, vec![a[0], a[4], a[2], a[3], a[8]], pv) + .get_operation() + .insert_at_back(block, &ctx); + } + let copy = CuteSm100TmaAttr { + layout: plan.a_layout, + cta_group: group, + }; + CuteSm100TmaLoadOp::new( + &mut ctx, + vec![a[5], a[1], a[4], a[4], a[4], a[6], a[4]], + copy.clone(), + ) + .get_operation() + .insert_at_back(block, &ctx); + CuteSm100TmaStoreOp::new(&mut ctx, vec![a[5], a[1], a[4], a[4], a[4]], copy) + .get_operation() + .insert_at_back(block, &ctx); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + f.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + pliron::builtin::attributes::StringAttr::new("true".into()), + ); + let target = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let text = cuda_oxide_mlir_export::render_module(&target).unwrap(); + let library = super::CutlassCompilerLibrary::load( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER").expect("pinned compiler path"), + super::CUTLASS_COMPILER_LIBRARY_SHA256, + ) + .unwrap(); + let dump = tempfile::tempdir().unwrap(); + let prefix = dump.path().join("attention_descriptors"); + let pipeline = super::CUTLASS_PRECOMPILED_PIPELINE.replace( + "cubin-format=bin", + &format!("cubin-format=bin dump-ptx-path={}", prefix.display()), + ); + let object = library + .compile_precompiled_mlir_to_object(&text, "sm_100a", &pipeline) + .unwrap_or_else(|e| panic!("attention compiler probe: {e}\n{text}")); + assert!(super::extract_kernels_binary(&object.elf).unwrap().len() > 64); + let ptx = std::fs::read_to_string(prefix.with_extension("sm_100a.ptx")).unwrap(); + let copies = ptx + .lines() + .filter(|line| { + line.contains("cp.async.bulk.tensor") + && line.contains(".mbarrier::complete_tx::bytes") + }) + .collect::>(); + assert_eq!(copies.len(), 1, "{ptx}"); + assert!(!copies[0].contains(".multicast"), "{ptx}"); + assert_eq!(copies[0].contains(".cta_group::2"), group == 2, "{ptx}"); + + assert_eq!( + text.matches("\"nvvm.tcgen05.mma\"").count(), + if group == 1 { 16 } else { 8 } + ); + assert!(text.contains("#cute_nvgpu.major"), "{text}"); + assert!(text.contains("num_cta = "), "{text}"); + assert!(text.contains("1, 1, 1, 3,"), "{text}"); + if group == 1 { + assert!(text.contains("#cute_nvgpu.major"), "{text}"); + } + } + } +} + +mod cutlass_attention_numeric_pack_compiles_module { + /* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + + use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; + use dialect_mir::ops::{MirFuncOp, MirReturnOp}; + use dialect_nvvm::ops::*; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::{TypeHandle, Typed}, + }; + + fn mark(ctx: &mut Context, op: Ptr, block: Ptr, marker: &str) { + op.deref_mut(ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + op.insert_at_back(block, ctx); + if op.deref(ctx).get_num_results() == 1 { + let value = op.deref(ctx).get_result(0); + let args = block.deref(ctx).arguments().collect::>(); + let ptr = args + .iter() + .rev() + .find(|arg| { + arg.get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|p| p.pointee == value.get_type(ctx)) + }) + .copied() + .expect("output pointer"); + let store = Operation::new( + ctx, + dialect_mir::ops::MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr, value], + vec![], + 0, + ); + dialect_mir::ops::MirStoreOp::new(store).set_volatile(ctx, true); + store.insert_at_back(block, ctx); + } + } + + #[test] + #[ignore = "requires the pinned CUTLASS compiler; no GPU"] + fn cutlass_attention_numeric_pack_compiles() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "numeric".try_into().unwrap()); + let f32: TypeHandle = FP32Type::get(&ctx).into(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let pred = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let shared = dialect_mir::types::MirPtrType::get(&mut ctx, i32, true, 3).into(); + let global = dialect_mir::types::MirPtrType::get(&mut ctx, i32, false, 1).into(); + let u64 = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let bar = dialect_mir::types::MirPtrType::get(&mut ctx, u64, true, 3).into(); + let f_out = dialect_mir::types::MirPtrType::get(&mut ctx, f32, true, 1).into(); + let u_out = dialect_mir::types::MirPtrType::get(&mut ctx, i32, true, 1).into(); + let boolean = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let b_out = dialect_mir::types::MirPtrType::get(&mut ctx, boolean, true, 1).into(); + let packed_out = dialect_mir::types::MirPtrType::get(&mut ctx, u64, true, 1).into(); + let args = vec![ + f32, f32, f32, i32, i32, pred, shared, global, bar, f_out, u_out, b_out, u64, u64, u64, + packed_out, + ]; + let fty = FunctionType::get(&ctx, args.clone(), vec![]); + let op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let f = MirFuncOp::new(&mut ctx, op, TypeAttr::new(fty.into())); + for (key, value) in [("maxntid", 128), ("minctasm", 1)] { + let integer = pliron::builtin::attributes::IntegerAttr::new( + IntegerType::get(&ctx, 32, Signedness::Unsigned), + pliron::utils::apint::APInt::from_u32( + value, + std::num::NonZeroUsize::new(32).unwrap(), + ), + ); + op.deref_mut(&ctx) + .attributes + .set(key.try_into().unwrap(), integer); + } + + f.set_symbol_name(&mut ctx, "numeric_probe".try_into().unwrap()); + module.append_operation(&mut ctx, op, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(op.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + // Distinct dynamic packed words retain both FP32 lanes through native arithmetic. + let add = AddF32x2Op::build(&mut ctx, a[12], a[13]); + mark(&mut ctx, add, block, "v1:i0995"); + let mul = MulF32x2Op::build(&mut ctx, a[12], a[13]); + mark(&mut ctx, mul, block, "v1:i0999"); + let fma = FmaF32x2Op::build(&mut ctx, a[12], a[13], a[14]); + mark(&mut ctx, fma, block, "v1:i1001"); + let shuffle = ShflSyncDownF32Op::build(&mut ctx, a[3], a[0], a[4]); + mark(&mut ctx, shuffle, block, "v1:i0056"); + let shuffle = ShflSyncBflyI32Op::build(&mut ctx, a[3], a[4], a[4]); + mark(&mut ctx, shuffle, block, "v1:i0051"); + // Ordinary Rust selection and saturation survive the same consumer boundary. + use dialect_mir::{ + ops::{MirCallOp, MirStoreOp}, + rust_intrinsics::{CALLEE_SATURATING_SUB, CALLEE_SELECT_UNPREDICTABLE}, + }; + let maximum = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![f32], + vec![a[0], a[1]], + vec![], + 0, + ); + MirCallOp::new(maximum).set_attr_callee( + &ctx, + StringAttr::new(dialect_mir::rust_intrinsics::CALLEE_MAXNUM_NSZ_F32.into()), + ); + maximum.insert_at_back(block, &ctx); + let max_value = maximum.deref(&ctx).get_result(0); + let max_store = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![a[9], max_value], + vec![], + 0, + ); + MirStoreOp::new(max_store).set_volatile(&mut ctx, true); + max_store.insert_at_back(block, &ctx); + let select = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![i32], + vec![a[5], a[3], a[4]], + vec![], + 0, + ); + MirCallOp::new(select) + .set_attr_callee(&ctx, StringAttr::new(CALLEE_SELECT_UNPREDICTABLE.into())); + select.insert_at_back(block, &ctx); + let selected = select.deref(&ctx).get_result(0); + let sub = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![i32], + vec![selected, a[4]], + vec![], + 0, + ); + MirCallOp::new(sub).set_attr_callee(&ctx, StringAttr::new(CALLEE_SATURATING_SUB.into())); + sub.insert_at_back(block, &ctx); + let value = sub.deref(&ctx).get_result(0); + let store = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![a[10], value], + vec![], + 0, + ); + MirStoreOp::new(store).set_volatile(&mut ctx, true); + store.insert_at_back(block, &ctx); + + for (kind, ftz, marker) in [ + (ScalarMathOperationAttr::Ex2, false, "v1:i0819"), + (ScalarMathOperationAttr::Lg2, false, "v1:i0786"), + (ScalarMathOperationAttr::Rcp, true, "v1:i0788"), + ] { + let op = ScalarMathOp::build( + &mut ctx, + a[0], + ScalarMathFormatAttr::F32, + kind, + ScalarMathPrecisionAttr::Approx, + if ftz { + ScalarMathSubnormalAttr::Ftz + } else { + ScalarMathSubnormalAttr::Preserve + }, + ); + mark(&mut ctx, op, block, marker); + } + let op = ScalarArithmeticOp::build( + &mut ctx, + a[..3].to_vec(), + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Fma, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + mark(&mut ctx, op, block, "v1:i0418"); + let op = CvtF16x2F32Op::build(&mut ctx, a[0], a[1]); + mark(&mut ctx, op, block, "v1:i0081"); + let op = ReduxSyncFmaxOp::build(&mut ctx, a[3], a[0]); + mark(&mut ctx, op, block, "v1:i1007"); + let op = SetMaxNRegIncOp::build(&mut ctx, 192); + mark(&mut ctx, op, block, "v1:i0915"); + let op = CpAsyncCa4Op::build(&mut ctx, a[6], a[7]); + mark(&mut ctx, op, block, "v1:i0086"); + let op = CpAsyncMbarrierArriveNoIncOp::build(&mut ctx, a[8]); + mark(&mut ctx, op, block, "v1:i0103"); + let op = VoteSyncAllOp::build(&mut ctx, a[3], a[5]); + mark(&mut ctx, op, block, "v1:i0040"); + let op = BarWarpSyncOp::build(&mut ctx, a[3]); + mark(&mut ctx, op, block, "v1:i0049"); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + block + .deref(&ctx) + .get_parent_region() + .unwrap() + .deref(&ctx) + .get_parent_op() + .deref_mut(&ctx) + .attributes + .set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + let result = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let text = cuda_oxide_mlir_export::render_module(&result).unwrap(); + let library = super::CutlassCompilerLibrary::load( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER").expect("pinned compiler path"), + super::CUTLASS_COMPILER_LIBRARY_SHA256, + ) + .unwrap(); + let dump = tempfile::tempdir().unwrap(); + let ptx_prefix = dump.path().join("attention_numeric"); + let pipeline = super::CUTLASS_PRECOMPILED_PIPELINE.replace( + "cubin-format=bin", + &format!("cubin-format=bin dump-ptx-path={}", ptx_prefix.display()), + ); + let object = library + .compile_precompiled_mlir_to_object(&text, "sm_100a", &pipeline) + .unwrap_or_else(|e| panic!("attention compiler probe: {e}\n{text}")); + assert!(super::extract_kernels_binary(&object.elf).unwrap().len() > 64); + let ptx = std::fs::read_to_string(ptx_prefix.with_extension("sm_100a.ptx")).unwrap(); + assert!(ptx.contains(".maxntid 128, 1, 1"), "{ptx}"); + assert!(ptx.contains(".minnctapersm 1"), "{ptx}"); + assert!(ptx.contains("setmaxnreg.inc.sync.aligned.u32"), "{ptx}"); + assert!( + ptx.lines() + .any(|line| line.split_whitespace().next() == Some("max.f32")), + "{ptx}" + ); + assert!( + !ptx.contains("max.NaN.f32") && !ptx.contains("max.ftz.f32"), + "{ptx}" + ); + + for instruction in ["add.rn.f32x2", "mul.rn.f32x2", "fma.rn.f32x2"] { + assert!( + ptx.contains(instruction), + "packed instruction lost: {instruction}\n{ptx}" + ); + } + + for name in [ + "llvm.nvvm.ex2.approx.f32", + "llvm.nvvm.lg2.approx.f", + "llvm.nvvm.rcp.approx.ftz.f", + "llvm.nvvm.fma.rn.f", + "nvvm.redux.sync", + "nvvm.cp.async.shared.global", + "nvvm.cp.async.mbarrier.arrive", + "nvvm.vote.sync", + "nvvm.bar.warp.sync", + "llvm.nvvm.ff2f16x2.rn", + "nvvm.add.packed.f32x2", + "nvvm.mul.packed.f32x2", + "nvvm.fma.packed.f32x2", + ] { + assert!(text.contains(name), "{text}"); + } + assert!(text.contains("192 : i32"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); + } +} + +/// The source-level CuTe-RS hint reuses the catalog-owned full-warp shuffle. +/// Compare that existing native route with the pinned CuTe operation itself. +#[test] +#[ignore = "requires the pinned CUTLASS compiler; no GPU"] +fn cutlass_warp_uniform_matches_cute_native_hint() { + use dialect_mir::ops::MirShrOp; + use dialect_nvvm::ops::{ReadPtxSregTidXOp, ShflSyncIdxI32Op}; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "warp_uniform".try_into().unwrap()); + let word = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let pointer: TypeHandle = MirPtrType::get(&mut ctx, word.into(), true, 1).into(); + let ty = FunctionType::get(&ctx, vec![pointer], vec![]); + let func = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, func, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "warp_uniform_probe".try_into().unwrap()); + func.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, func, 0); + let block = BasicBlock::new(&mut ctx, None, vec![pointer]); + block.insert_at_back(func.deref(&ctx).get_region(0), &ctx); + let output = block.deref(&ctx).arguments().next().unwrap(); + let tid = + append::(&mut ctx, block, vec![], vec![word.into()], Some("v1:i0001")); + let tid = tid.deref(&ctx).get_result(0); + let five = integer_constant(&mut ctx, block, word, 5, 32); + let warp = append::(&mut ctx, block, vec![tid, five], vec![word.into()], None); + let warp = warp.deref(&ctx).get_result(0); + let full_mask = integer_constant(&mut ctx, block, word, u32::MAX as u64, 32); + let zero = integer_constant(&mut ctx, block, word, 0, 32); + let uniform = append::( + &mut ctx, + block, + vec![full_mask, warp, zero], + vec![word.into()], + Some("v1:i0050"), + ); + let value = uniform.deref(&ctx).get_result(0); + let store = append::(&mut ctx, block, vec![output, value], vec![], None); + MirStoreOp::new(store).set_volatile(&mut ctx, true); + append::(&mut ctx, block, vec![], vec![], None); + let target = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let native = cuda_oxide_mlir_export::render_module(&target).unwrap(); + assert!(native.contains("llvm.nvvm.shfl.sync.idx.i32"), "{native}"); + let cute = r#"module { + gpu.module @kernels { + llvm.func @warp_uniform_probe(%out: !llvm.ptr<1>) attributes {nvvm.kernel} { + %tid = nvvm.read.ptx.sreg.tid.x : i32 + %five = arith.constant 5 : i32 + %warp = arith.shrui %tid, %five : i32 + %uniform = "cute_nvgpu.arch.make_warp_uniform"(%warp) : (i32) -> i32 + llvm.store %uniform, %out {volatile_} : i32, !llvm.ptr<1> + llvm.return + } + } + }"#; + let library = CutlassCompilerLibrary::load( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER").expect("pinned compiler path"), + CUTLASS_COMPILER_LIBRARY_SHA256, + ) + .unwrap(); + let dump = tempfile::tempdir().unwrap(); + for (name, source) in [("native", native.as_str()), ("cute", cute)] { + let prefix = dump.path().join(name); + let pipeline = CUTLASS_PRECOMPILED_PIPELINE.replace( + "cubin-format=bin", + &format!("cubin-format=bin dump-ptx-path={}", prefix.display()), + ); + let object = library + .compile_precompiled_mlir_to_object(source, "sm_100a", &pipeline) + .unwrap_or_else(|e| panic!("{name} uniform probe: {e}\n{source}")); + assert!(extract_kernels_binary(&object.elf).unwrap().len() > 64); + let ptx = std::fs::read_to_string(prefix.with_extension("sm_100a.ptx")).unwrap(); + let shuffle = ptx + .lines() + .filter(|line| line.contains("shfl.sync")) + .collect::>(); + assert_eq!(shuffle.len(), 1, "{name}: {ptx}"); + assert!(shuffle[0].contains("shfl.sync.idx.b32"), "{name}: {ptx}"); + assert!( + shuffle[0].trim_end().ends_with(", 0, 31, -1;"), + "{name}: {ptx}" + ); + } +} + +/// A semantic register tile must remain one cache-aware 256-bit global store. +#[test] +#[ignore = "requires the pinned CUTLASS compiler; no GPU"] +fn cutlass_global_store_no_allocate_is_one_256_bit_transaction() { + assert!(CutlassFullCuteMlir22::new("sm_90a").is_err()); + use dialect_cute::copy_ops::CuteStoreGlobalNoAllocateOp; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "wide_store".try_into().unwrap()); + let word: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let output: TypeHandle = MirPtrType::get(&mut ctx, word, true, address_space::GLOBAL).into(); + let source: TypeHandle = MirPtrType::get(&mut ctx, word, false, address_space::GENERIC).into(); + let ty = FunctionType::get(&ctx, vec![output, source], vec![]); + let func = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, func, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "wide_store_probe".try_into().unwrap()); + func.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, func, 0); + let block = BasicBlock::new(&mut ctx, None, vec![output, source]); + block.insert_at_back(func.deref(&ctx).get_region(0), &ctx); + let a: Vec<_> = block.deref(&ctx).arguments().collect(); + let store = CuteStoreGlobalNoAllocateOp::new(&mut ctx, a[0], a[1], 8); + store.get_operation().insert_at_back(block, &ctx); + append::(&mut ctx, block, vec![], vec![], None); + let target = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let text = cuda_oxide_mlir_export::render_module(&target).unwrap(); + assert!( + text.contains("atom.r2g"), + "{text}" + ); + assert!(text.contains("gmem, align<32>"), "{text}"); + assert!(text.contains("rmem, align<32>"), "{text}"); + let compiler = CutlassCompilerLibrary::load( + std::env::var_os("CUDA_OXIDE_TEST_CUTLASS_COMPILER").expect("pinned compiler path"), + CUTLASS_COMPILER_LIBRARY_SHA256, + ) + .unwrap(); + let dump = tempfile::tempdir().unwrap(); + let prefix = dump.path().join("wide-store"); + let pipeline = CUTLASS_PRECOMPILED_PIPELINE.replace( + "cubin-format=bin", + &format!("cubin-format=bin dump-ptx-path={}", prefix.display()), + ); + let object = compiler + .compile_precompiled_mlir_to_object(&text, "sm_100a", &pipeline) + .unwrap_or_else(|e| panic!("{e}\n{text}")); + assert!(extract_kernels_binary(&object.elf).unwrap().len() > 64); + let ptx = std::fs::read_to_string(prefix.with_extension("sm_100a.ptx")).unwrap(); + let stores: Vec<_> = ptx + .lines() + .filter(|line| line.contains("st.global")) + .collect(); + assert_eq!(stores.len(), 1, "{ptx}"); + assert!( + stores[0].contains("L1::no_allocate") + && (stores[0].contains(".v8.b32") || stores[0].contains(".v4.b64")), + "{ptx}" + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/Cargo.toml b/cuda-oxide/crates/cuda-oxide-mlir-export/Cargo.toml new file mode 100644 index 0000000000..58c070b672 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "cuda-oxide-mlir-export" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "CUDA Oxide mapping packs and profiles for textual MLIR export" +readme = "README.md" + +[dependencies] +pliron = { workspace = true } +pliron-mlir-interop = { workspace = true } +pliron-derive = { workspace = true } +dialect-mir = { workspace = true } +dialect-nvvm = { workspace = true } +dialect-cute = { workspace = true } +cute-layout = { workspace = true } +llvm-export = { workspace = true } +reserved-oxide-symbols = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +sha2 = { workspace = true } diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/README.md b/cuda-oxide/crates/cuda-oxide-mlir-export/README.md new file mode 100644 index 0000000000..b3b91b0124 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/README.md @@ -0,0 +1,207 @@ +# cuda-oxide MLIR export + +This crate maps cuda-oxide operations to MLIR accepted by the pinned CUTLASS 4.7 +compiler: explicit builtin, MIR, NVVM and CuTe mapping packs, source checks, typed +translation plans and consumer profiles. The generic printer comes from upstream +`pliron-mlir-interop`. + +```text +shared Pliron module (dialect-mir, dialect-cute, dialect-nvvm) + │ CuTe whole-module verification and the architecture gate + ▼ +typed target plan explicit mapping recipes, checked value and block connections + │ local ToMlirOp, ToMlirType and ToMlirAttr adapters + ▼ +Pliron target graph ──► upstream GenericOp and MlirPrinter ──► MLIR text + │ + ▼ +official CUTLASS 4.7 compiler ──► cubin, validated and embedded in the host binary +``` + +`render_module` returns `Result`; an unsupported mapping or a +malformed plan is an error. The workspace pins crates.io `pliron-mlir-interop = "=0.1.0"` +with Pliron `=0.20.0`; no upstream source is vendored. Two compatibility pieces live in +the local adapter under `src/interop/`: typed discardable attributes render through the +upstream `MlirPrinter` and pass to `GenericOp::attr_raw`, and block-argument locations are +restored into complete block-header lines derived from the target graph, where each +header must match a single block or export fails rather than dropping a location. The +export manifest names this contract `pliron-mlir-interop-0.1.0+cuda-oxide.2`, so cached +artifacts cannot reuse an older printer contract. MIR, CuTe and NVVM mappings stay local; +upstream's bundled LLVM source-dialect translations are not used. + +cuda-oxide installs the exact CUTLASS release archive, verifies the archive and library +digests, and loads that versioned library explicitly. The backend extracts and validates +the cubin and embeds it for the ordinary launch path. Set `CUDA_OXIDE_MLIR_OUTPUT=` +to keep the MLIR text; select `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir` to compile it. + +The SM100 primitive pack additionally maps two-CTA FP16 TCGen05 MMA with A +collector selectors, tensor-memory allocation/load/commit/wait operations, +cluster synchronization, elected-lane results and multicast TMA. The +[`fp16_gemm_256x352_cute`](../../cute-rs-preview/examples/fp16_gemm_256x352_cute) +example exercises these together. It also preserves fixed cluster dimensions +in `nvvm.cluster_dim` and the scalar order of compiler-created register arrays. +Unsupported primitive selectors and intrinsic identities fail at export. + +## Primitive-only profile + +Select `CUDA_OXIDE_MLIR_PROFILE=cutlass-primitives-mlir22` together with the +CUTLASS backend for Rust kernels using `ctm_rs`. This profile +accepts ordinary MIR plus the reviewed NVVM subset and rejects CuTe operations, +types and attributes recursively. It targets `sm_100a` and `sm_120a`; TCGen05 +requires `sm_100a`, while MXF4 register MMA requires `sm_120a`. The default full +CuTe profile remains available. Profile identity participates in build caching. + +Real Rust fixtures also exercise checked arithmetic, padded tuple carriers, +Direct/Single enum storage, static shared globals and Rust pointer-address +semantics. Explicit NVVM shared offsets preserve AS3/AS7; Rust address casts +convert named pointers to generic addresses first. See the +[ctm-rs README](../../cute-rs-preview/ctm-rs/README.md). + +See [`hoist_enum_scratch`](src/mir/checked.rs) for why enum scratch allocations +move to the function entry while their loads and stores stay in place. + +## Scalar code and control flow + +These mappings handle the ordinary scalar code around a CuTe kernel: + +```text +mir.func / mir.return → func.func / func.return +mir.goto / mir.cond_br → cf.br / cf.cond_br +mir.assert → cf.assert +mir.constant → arith.constant +mir.add / sub / mul → arith.add* / sub* / mul* +mir.div / rem → signed or unsigned arith operation +mir.lt / le / gt / ge → arith.cmpi or arith.cmpf +mir.eq / ne → arith.cmpi or arith.cmpf +mir.cast → the matching arith conversion +``` + +Rust's `usize` is a 64-bit unsigned integer on the CUDA target. MLIR integer +types do not write signedness in the type, so it becomes `i64`: + +```text +Pliron MLIR +builtin.integer ui64 i64 +mir.div on that ui64 arith.divui +mir.lt on that ui64 arith.cmpi ult +``` + +The operation name keeps the part that matters. A signed `i64` division uses +`arith.divsi`; an unsigned `u64` division uses `arith.divui`. The translator +looks at the Pliron type before both types become MLIR `i64`. + +Control flow maps to `cf`, not `scf`. Rust MIR is already a graph of basic +blocks and can jump in shapes that are not a neat `if` or `for` loop: + +```text + mir.cond_br + / \ + block A block B + \ / + join + + ↓ + + cf.cond_br + / \ + block A block B + \ / + join +``` + +Later MLIR passes may rebuild structured loops. The exporter does not guess +that structure while translating the graph. + +## Pointers and small Rust values + +The elementwise kernel still contains ordinary Rust values around its CuTe +views. They keep their physical shape when they cross into MLIR: + +```text +mir.ptr → !llvm.ptr +mir.slice → !llvm.struct<(!llvm.ptr, i64)> +mir.array → !llvm.array<4 x T> +mir.struct → !llvm.struct<(A, B)> + +mir.alloca / load / store → llvm.alloca / load / store +mir.field_addr → llvm.getelementptr +mir.extract_field → llvm.extractvalue +``` + +Rust can reorder fields and insert padding. The translator reads the MIR +layout and rebuilds that exact physical order. It refuses packed or otherwise +different by-value layouts instead of silently changing an address. + +For `#[grid_constant] value: &T`, the host passes `T` by value. Every thread +borrows the same read-only kernel parameter storage: + +```text +host T -> one by-value launch argument -> device &T shared by the grid +``` + +The exporter records the value's size with `llvm.byval`, its Rust alignment +with `llvm.align`, and shared read-only storage with `nvvm.grid_constant`. +For a TMA descriptor this means 128 bytes aligned to 64 bytes. Slices expand +into pointer and length arguments, so these attributes follow the descriptor +to its new argument index. + +The backend checks the cubin's parameter sizes and offsets before saving it. +Missing or inconsistent metadata fails compilation. The caller still owns +the [pointer validity and lifetime requirements](../../cuda-oxide-book/gpu-programming/kernels-and-device-functions.md#grid-constant-parameters) +for data accessed through the value. + +The CUDA kernel marker also has an exact mapping: + +```text +gpu_kernel = "true" → cute.kernel +anything else → translation error +``` + +## Tensor operations + +The elementwise path keeps tensor layouts visible until CUTLASS lowers them: + +```text +make tensor + │ + ▼ +divide it into tiles + │ + ▼ +pick this thread's tile + │ + ├── full tile ──► cute.copy + │ + └── edge tile ──► checked scalar load/store +``` + +The profile maps `cute.tensor_*` operations to CUTLASS's `cute` and +`cute_nvgpu` operations. First, the shared verifier checks tensor origins, +pipeline state, and operation ordering across the module. The verified CuTe +operations then go directly to MLIR translation. + +Thread, block, and grid coordinates are direct NVVM mappings. For example: + +```text +nvvm.read_ptx_sreg_tid_x → nvvm.read.ptx.sreg.tid.x +nvvm.read_ptx_sreg_ntid_y → nvvm.read.ptx.sreg.ntid.y +nvvm.read_ptx_sreg_nctaid_z → nvvm.read.ptx.sreg.nctaid.z +``` + +Only the reviewed coordinate reads are registered. Other NVVM operations fail +with a missing-mapping error until they get their own recipe. + +## Mapping organization + +The public registration functions and consumer profiles remain at the crate +root. Private `mir/` and `nvvm/` modules own their dialect mappings. The `cute/` +hierarchy groups tensor construction and GEMV, cooperative and TMA copy, MMA, +pipeline synchronization, epilogues, and scheduling. SM100 and SM120 recipes +live under the capability they implement; the independent SM100 attention +recipes retain their own contracts. + +The full CuTe profile keeps registration sequencing explicit. The GEMM +foundation owns shared attributes and the MIR insert-field and extern-shared +adapters, including for SM100. The tensor profile owns the combined tensor-view +translator, which delegates elementwise views to their existing translator. +Shared registration is performed once and is not filtered by architecture. diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common.rs new file mode 100644 index 0000000000..e934e8c9f6 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common.rs @@ -0,0 +1,445 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared static handles, MIR adapters, and scalar operation builders. + +use crate::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationRegistry, TranslationSession, TypeTranslation, +}; +use dialect_mir::{ + ops::{MirExternSharedOp, MirInsertFieldOp}, + types::MirUnionType, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Represent source-only verified layout contracts with an empty target tuple. +pub(in crate::cute) struct StaticContractTypeTranslation; + +impl TypeTranslation for StaticContractTypeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + Ok(MlirType::Tuple(vec![])) + } +} + +/// Map MIR field insertion to llvm.insertvalue at its verified field index. +pub(in crate::cute) struct InsertFieldTranslation; + +impl OperationTranslation for InsertFieldTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = MirInsertFieldOp::new(source); + semantic.verify(ctx).map_err(stringify)?; + expect_shape("mir.insert_field", &input, 1, 2)?; + let index = semantic + .get_attr_insert_index(ctx) + .ok_or_else(|| "mir.insert_field has no insert_index".to_owned())? + .0 as i64; + if source + .deref(ctx) + .get_operand(0) + .get_type(ctx) + .deref(ctx) + .is::() + { + return crate::mir::storage::insert_union_field( + ctx, + source, + input, + session, + index as usize, + ); + } + let mut target = operation( + "llvm.insertvalue", + input.results, + input.operands, + &input.location, + )?; + target + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![index])); + Ok(vec![target]) + } +} + +/// Obtain dynamic shared memory and apply the source byte offset and alignment. +pub(in crate::cute) struct ExternSharedTranslation; + +impl OperationTranslation for ExternSharedTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = MirExternSharedOp::new(source); + semantic.verify(ctx).map_err(stringify)?; + input.attributes.remove("extern_byte_offset"); + input.attributes.remove("extern_alignment"); + expect_shape("mir.extern_shared", &input, 1, 0)?; + let offset = semantic.get_byte_offset_value(ctx); + let alignment = semantic.get_alignment_value(ctx); + let cute_pointer = MlirType::dialect(format!("!cute.ptr>"))?; + let (raw_result, raw) = fresh(session, cute_pointer); + let raw_op = operation( + "cute_nvgpu.arch.get_dyn_smem", + vec![raw_result], + vec![], + &input.location, + )?; + let (base_result, base) = if offset == 0 { + ( + input.results[0].clone(), + MlirValueUse { + id: input.results[0].id, + ty: input.results[0].ty.clone(), + }, + ) + } else { + fresh(session, shared_pointer_type()?) + }; + let bridge = operation( + "builtin.unrealized_conversion_cast", + vec![base_result], + vec![raw], + &input.location, + )?; + if offset == 0 { + return Ok(vec![raw_op, bridge]); + } + + let index = constant_value( + session, + MlirType::Integer(64), + offset as i128, + &input.location, + )?; + let mut gep = operation( + "llvm.getelementptr", + input.results, + vec![base, index.1], + &input.location, + )?; + gep_properties(&mut gep, MlirType::Integer(8), vec![i32::MIN]); + Ok(vec![raw_op, bridge, index.0, gep]) + } +} + +/// Preserve a one-input static handle through an unrealized conversion cast. +pub(in crate::cute) struct IdentityTranslation( + #[doc = "Source operation spelling used when its one-to-one contract is violated."] + pub(in crate::cute) &'static str, +); + +impl OperationTranslation for IdentityTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape(self.0, &input, 1, 1)?; + Ok(vec![operation( + "builtin.unrealized_conversion_cast", + input.results, + input.operands, + &input.location, + )?]) + } +} + +/// Preserve the first runtime operand while erasing verified static operands. +pub(in crate::cute) struct IdentityFirstOperandTranslation( + #[doc = "Source operation spelling used when its first-operand forwarding contract is violated."] + pub(in crate::cute) &'static str, +); + +impl OperationTranslation for IdentityFirstOperandTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + if input.results.len() != 1 || input.operands.is_empty() { + return Err(format!( + "{} needs one result and at least one operand", + self.0 + )); + } + if !input.successors.is_empty() || !input.regions.is_empty() || !input.attributes.is_empty() + { + return Err(format!("{} retained source-only structure", self.0)); + } + Ok(vec![operation( + "builtin.unrealized_conversion_cast", + input.results, + vec![input.operands[0].clone()], + &input.location, + )?]) + } +} + +// The remaining definitions in this file are kept below the scalar foundation +// so each official-parser rung can be tested independently. + +/// Erase listed static attributes and forward each operand to its corresponding result. +pub(in crate::cute) struct StaticMultiIdentityTranslation { + /// Source operation spelling used in contract diagnostics. + pub(in crate::cute) name: &'static str, + /// Source-only attributes consumed before forwarding runtime operands. + pub(in crate::cute) attributes: &'static [&'static str], +} + +impl OperationTranslation for StaticMultiIdentityTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + for attribute in self.attributes { + input.attributes.remove(*attribute); + } + if input.results.len() != input.operands.len() { + return Err(format!( + "{} expected matching result and operand counts", + self.name + )); + } + expect_shape(self.name, &input, input.results.len(), input.operands.len())?; + input + .results + .into_iter() + .zip(input.operands) + .map(|(result, operand)| { + operation( + "builtin.unrealized_conversion_cast", + vec![result], + vec![operand], + &input.location, + ) + }) + .collect() + } +} + +/// Append the requested integer conversion and return its typed result. +pub(in crate::cute) fn cast_int( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: MlirValueUse, + target: MlirType, + name: &str, + location: &MlirLocation, +) -> Result { + let (result, use_) = fresh(session, target); + operations.push(operation(name, vec![result], vec![value], location)?); + Ok(use_) +} + +/// Append an LLVM pointer offset measured in elements of the supplied type. +pub(in crate::cute) fn gep_value( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base: MlirValueUse, + index: MlirValueUse, + element: MlirType, + location: &MlirLocation, +) -> Result { + let (result, use_) = fresh(session, base.ty.clone()); + let mut gep = operation( + "llvm.getelementptr", + vec![result], + vec![base, index], + location, + )?; + gep_properties(&mut gep, element, vec![i32::MIN]); + operations.push(gep); + Ok(use_) +} + +/// Set GEP element and index metadata without promising no-wrap behavior. +pub(in crate::cute) fn gep_properties( + operation: &mut MlirOperation, + element: MlirType, + indices: Vec, +) { + operation + .properties + .insert("elem_type".into(), MlirAttribute::Type(element)); + operation.properties.insert( + "noWrapFlags".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(32), + }, + ); + operation.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(indices), + ); +} + +/// Append an i64 constant while preserving the unsigned input's bit pattern. +pub(in crate::cute) fn emit_i64_constant( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: u64, + location: &MlirLocation, +) -> Result { + let (operation_, use_) = + constant_value(session, MlirType::Integer(64), u64_bits(value), location)?; + operations.push(operation_); + Ok(use_) +} + +/// Append an i32 constant and return its SSA operand handle. +pub(in crate::cute) fn emit_i32_constant( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: i128, + location: &MlirLocation, +) -> Result { + let (operation_, use_) = constant_value(session, MlirType::Integer(32), value, location)?; + operations.push(operation_); + Ok(use_) +} + +/// Allocate a typed integer constant and return its operation and SSA use. +pub(in crate::cute) fn constant_value( + session: &mut TranslationSession<'_>, + ty: MlirType, + value: i128, + location: &MlirLocation, +) -> Result<(MlirOperation, MlirValueUse), String> { + let (result, use_) = fresh(session, ty); + Ok((integer_constant(result, value, location)?, use_)) +} + +/// Append a two-operand operation whose result has the left operand's type. +pub(in crate::cute) fn emit_binary( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + name: &str, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + let (result, use_) = fresh(session, left.ty.clone()); + operations.push(binary(name, result, left, right, location)?); + Ok(use_) +} + +/// Build a two-operand target operation using the supplied result. +pub(in crate::cute) fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +/// Build an arith.constant using the supplied result's integer type. +pub(in crate::cute) fn integer_constant( + result: MlirResult, + value: i128, + location: &MlirLocation, +) -> Result { + let mut operation = operation("arith.constant", vec![result.clone()], vec![], location)?; + operation.properties.insert( + "value".into(), + MlirAttribute::Integer { + value, + ty: result.ty, + }, + ); + Ok(operation) +} + +/// Require exact result/operand counts and no leftover attributes, regions, or successors. +pub(in crate::cute) fn expect_shape( + name: &str, + input: &OperationInput, + results: usize, + operands: usize, +) -> Result<(), String> { + if input.results.len() != results || input.operands.len() != operands { + return Err(format!( + "{name} expected {results} results and {operands} operands, got {} and {}", + input.results.len(), + input.operands.len() + )); + } + if !input.successors.is_empty() || !input.regions.is_empty() || !input.attributes.is_empty() { + return Err(format!("{name} retained source-only target structure")); + } + Ok(()) +} + +/// Build a target operation with the supplied operands, results, and source location. +pub(in crate::cute) fn operation( + name: &str, + results: Vec, + operands: Vec, + location: &MlirLocation, +) -> Result { + let mut operation = MlirOperation::new(name)?; + operation.results = results; + operation.operands = operands; + operation.location = location.clone(); + Ok(operation) +} + +/// Allocate one SSA identifier and return matching result and operand handles. +pub(in crate::cute) fn fresh( + session: &mut TranslationSession<'_>, + ty: MlirType, +) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +/// Return the opaque LLVM pointer type for generic address space zero. +pub(in crate::cute) fn generic_pointer_type() -> Result { + MlirType::dialect("!llvm.ptr") +} + +/// Return the opaque LLVM pointer type for CTA shared address space three. +pub(in crate::cute) fn shared_pointer_type() -> Result { + MlirType::dialect("!llvm.ptr<3>") +} + +/// Preserve all 64 bits while spelling an unsigned value as a signed MLIR literal. +pub(in crate::cute) fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +/// Preserve a source diagnostic as the string expected by translation recipes. +pub(in crate::cute) fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common_sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common_sm100.rs new file mode 100644 index 0000000000..f4acedc738 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/common_sm100.rs @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared input contract for independent SM100 attention operations. + +use crate::OperationInput; + +/// Reject results, attributes, regions, or successors on independent SM100 operations. +pub(in crate::cute) fn contract(input: &OperationInput) -> Result<(), String> { + if !input.results.is_empty() + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err( + "SM100 operation retained unexpected results, attributes or control flow".into(), + ); + } + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs new file mode 100644 index 0000000000..c17c707d4c --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Realize verified CuTe cooperative-copy and matrix-load plans for attention. +//! +//! The source copy ABI explicitly requires partial-atom source-size zero filling; +//! it is not an ordinary predicated CuTe copy (which can leave shared bytes +//! untouched). NVVM expresses that contract directly. Both operations retain the +//! verified CuTe thread/value assignment and composed shared layout, including +//! element-to-byte swizzle conversion. Commit, wait, and synchronization stay +//! caller-owned, as do the unsafe pointer, capacity, and participation promises. + +use crate::cute::copy::emit::{Emitter, integer, require_input}; +use crate::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationSession, +}; +use cute_layout::{IntTuple, validate_cooperative_copy_plan}; +use dialect_cute::ops::CuteCopyG2SOp; +use dialect_mir::types::MirFP16Type; +use pliron::{ + builtin::types::IntegerType, + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::TypeHandle, +}; + +/// Emit cooperative two-byte storage copies with partial-atom zero filling. +pub(in crate::cute) struct CopyTranslation; +impl OperationTranslation for CopyTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let copy = CuteCopyG2SOp::wrap(source); + copy.verify(ctx).map_err(|e| e.to_string())?; + let elem: TypeHandle = copy.get_attr_copy_g2s_elem(ctx).unwrap().clone().into(); + let element = if elem.deref(ctx).is::() { + MlirType::Float(crate::MlirFloatType::F16) + } else if elem + .deref(ctx) + .downcast_ref::() + .is_some_and(|ty| ty.width() == 16) + { + // cp.async transports unchanged bytes. Integer carriers preserve + // BF16 storage without introducing an FP16 numerical conversion. + MlirType::Integer(16) + } else { + return Err("CuTe cooperative copy supports f16 or 16-bit integer storage only".into()); + }; + input.attributes.remove("copy_g2s_elem"); + require_input(&input, 9)?; + let atom = copy.get_attr_atom_bytes(ctx).unwrap().0; + let plan = validate_cooperative_copy_plan( + atom, + ©.get_attr_thread_layout(ctx).unwrap().0, + ©.get_attr_value_layout(ctx).unwrap().0, + ©.get_attr_tile_layout(ctx).unwrap().0, + ©.get_attr_smem_layout(ctx).unwrap().0, + 2, + ) + .map_err(|e| e.to_string())?; + let mut e = Emitter::new(session, input.location); + let global = e.pointer(input.operands[0].clone(), 1)?; + let shared = e.pointer(input.operands[6].clone(), 3)?; + // The cooperative-copy contract supplies exactly the verified thread + // IDs. Preserve that range in the scalar address calculation too. + let thread_count = e.constant(plan.thread_count, 32)?; + let thread = e.binary("arith.remui", input.operands[8].clone(), thread_count)?; + let thread = e.value("arith.extui", MlirType::Integer(64), vec![thread])?; + let tv_modes = plan.tv_layout.modes(); + let thread_cell = e.index(&tv_modes[0], thread)?; + let tile_rows = e.c64(plan.tiler_shape[0])?; + let tile_columns = e.c64(plan.tiler_shape[1])?; + let row_base = e.binary("arith.muli", input.operands[4].clone(), tile_rows.clone())?; + let column_base = e.binary( + "arith.muli", + input.operands[5].clone(), + tile_columns.clone(), + )?; + // Reject overflowing tile origins through the bounds predicate before + // choosing a source address. Out-of-range coordinates still zero-fill. + let last_row_tile = + e.binary("arith.divui", input.operands[1].clone(), tile_rows.clone())?; + let last_column_tile = e.binary("arith.divui", input.operands[2].clone(), tile_columns)?; + let row_tile_valid = e.compare(7, input.operands[4].clone(), last_row_tile)?; // ule + let column_tile_valid = e.compare(7, input.operands[5].clone(), last_column_tile)?; + let tile_valid = e.binary("arith.andi", row_tile_valid, column_tile_valid)?; + let zero = e.c64(0)?; + let two = e.c64(2)?; + let values_per_atom = i64::from(atom) / 2; + let atom_values = e.c64(values_per_atom)?; + for value_index in (0..plan.values_per_thread).step_by(values_per_atom as usize) { + let value_cell = tv_modes[1] + .checked_call(&IntTuple::Leaf(value_index)) + .ok_or_else(|| "verified CuTe value assignment is out of range".to_owned())?; + let value_cell = e.c64(value_cell)?; + let cell = e.binary("arith.addi", thread_cell.clone(), value_cell)?; + let row = e.binary("arith.remui", cell.clone(), tile_rows.clone())?; + let column = e.binary("arith.divui", cell, tile_rows.clone())?; + let shared_offset = e.composed(&plan.smem_byte_layout, row.clone(), column.clone())?; + let destination = e.gep(shared.clone(), shared_offset, MlirType::Integer(8))?; + let row = e.binary("arith.addi", row_base.clone(), row)?; + let column = e.binary("arith.addi", column_base.clone(), column)?; + let row_valid = e.compare(6, row.clone(), input.operands[1].clone())?; // ult + let column_valid = e.compare(6, column.clone(), input.operands[2].clone())?; + let valid = e.binary("arith.andi", row_valid, column_valid)?; + let valid = e.binary("arith.andi", valid, tile_valid.clone())?; + let remaining = e.binary("arith.subi", input.operands[2].clone(), column.clone())?; + let remaining = e.binary("arith.minui", remaining, atom_values.clone())?; + let valid_elements = e.select(valid.clone(), remaining, zero.clone())?; + let valid_bytes = e.binary("arith.muli", valid_elements, two.clone())?; + let valid_bytes = e.value("arith.trunci", MlirType::Integer(32), vec![valid_bytes])?; + let row_offset = e.binary("arith.muli", row, input.operands[3].clone())?; + let source_offset = e.binary("arith.addi", row_offset, column)?; + let source_offset = e.select(valid, source_offset, zero.clone())?; + let source = e.gep(global.clone(), source_offset, element.clone())?; + let mut op = e.operation( + "nvvm.cp.async.shared.global", + vec![destination, source, valid_bytes], + )?; + op.properties + .insert("size".into(), integer(i128::from(atom), 32)); + op.properties.insert( + "modifier".into(), + MlirAttribute::dialect(if atom == 16 { + "#nvvm" + } else { + "#nvvm" + })?, + ); + e.operations.push(op); + } + Ok(e.operations) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/emit.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/emit.rs new file mode 100644 index 0000000000..27d8e4d239 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/emit.rs @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Layout-aware operation builder shared by cooperative copy and matrix loads. + +use crate::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + TranslationSession, +}; +use cute_layout::{ComposedLayout, Layout, Swizzle}; + +/// Require the expected operands and no results or untranslated source structure. +pub(in crate::cute) fn require_input( + input: &OperationInput, + operands: usize, +) -> Result<(), String> { + if input.operands.len() != operands + || !input.results.is_empty() + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("CuTe attention operation retained unexpected structure".into()); + } + Ok(()) +} + +/// Build a signless integer attribute with an explicit bit width. +pub(in crate::cute) fn integer(value: i128, width: u32) -> MlirAttribute { + MlirAttribute::Integer { + value, + ty: MlirType::Integer(width), + } +} + +/// Accumulate scalar MLIR for verified layout indexing and memory accesses. +pub(in crate::cute) struct Emitter<'a, 'ctx> { + /// Translation session providing unique SSA and block identifiers. + pub(in crate::cute) session: &'a mut TranslationSession<'ctx>, + /// Source location attached to every operation emitted by this builder. + pub(in crate::cute) location: MlirLocation, + /// Target operations accumulated in execution order. + pub(in crate::cute) operations: Vec, +} +impl<'a, 'ctx> Emitter<'a, 'ctx> { + /// Start an empty operation sequence using the source location and shared SSA allocator. + pub(in crate::cute) fn new( + session: &'a mut TranslationSession<'ctx>, + location: MlirLocation, + ) -> Self { + Self { + session, + location, + operations: vec![], + } + } + /// Allocate one SSA identifier and return matching result and operand handles. + pub(in crate::cute) fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = self.session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) + } + /// Build a result-free operation at this builder's source location without appending it. + pub(in crate::cute) fn operation( + &self, + name: &str, + operands: Vec, + ) -> Result { + let mut op = MlirOperation::new(name)?; + op.location = self.location.clone(); + op.operands = operands; + Ok(op) + } + /// Append an operation with one newly allocated result of the requested type. + pub(in crate::cute) fn value( + &mut self, + name: &str, + ty: MlirType, + operands: Vec, + ) -> Result { + let (result, use_) = self.fresh(ty); + let mut op = self.operation(name, operands)?; + op.results.push(result); + self.operations.push(op); + Ok(use_) + } + /// Append a signed 64-bit integer constant. + pub(in crate::cute) fn c64(&mut self, value: i64) -> Result { + self.constant(value, 64) + } + /// Append an integer constant with the requested width in bits. + pub(in crate::cute) fn constant( + &mut self, + value: i64, + width: u32, + ) -> Result { + let (result, use_) = self.fresh(MlirType::Integer(width)); + let mut op = self.operation("arith.constant", vec![])?; + op.results.push(result); + op.properties + .insert("value".into(), integer(i128::from(value), width)); + self.operations.push(op); + Ok(use_) + } + /// Append a binary operation preserving the left operand's type. + pub(in crate::cute) fn binary( + &mut self, + name: &str, + a: MlirValueUse, + b: MlirValueUse, + ) -> Result { + self.value(name, a.ty.clone(), vec![a, b]) + } + /// Append an integer comparison using the MLIR predicate number, returning i1. + pub(in crate::cute) fn compare( + &mut self, + predicate: i128, + a: MlirValueUse, + b: MlirValueUse, + ) -> Result { + let value = self.value("arith.cmpi", MlirType::Integer(1), vec![a, b])?; + self.operations + .last_mut() + .unwrap() + .properties + .insert("predicate".into(), integer(predicate, 64)); + Ok(value) + } + /// Append an arith.select with the chosen operands' common type. + pub(in crate::cute) fn select( + &mut self, + predicate: MlirValueUse, + yes: MlirValueUse, + no: MlirValueUse, + ) -> Result { + self.value("arith.select", yes.ty.clone(), vec![predicate, yes, no]) + } + /// Convert an LLVM pointer to the requested address space when necessary. + pub(in crate::cute) fn pointer( + &mut self, + pointer: MlirValueUse, + space: u32, + ) -> Result { + let ty = MlirType::dialect(format!("!llvm.ptr<{space}>"))?; + if pointer.ty == ty { + Ok(pointer) + } else { + self.value("llvm.addrspacecast", ty, vec![pointer]) + } + } + /// Append a dynamic element-indexed GEP without no-wrap assumptions. + pub(in crate::cute) fn gep( + &mut self, + base: MlirValueUse, + index: MlirValueUse, + element: MlirType, + ) -> Result { + let value = self.value("llvm.getelementptr", base.ty.clone(), vec![base, index])?; + let op = self.operations.last_mut().unwrap(); + op.properties + .insert("elem_type".into(), MlirAttribute::Type(element)); + op.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![i32::MIN]), + ); + op.properties.insert("noWrapFlags".into(), integer(0, 32)); + Ok(value) + } + /// A flat CuTe coordinate expands in column-major order through its modes. + pub(in crate::cute) fn index( + &mut self, + layout: &Layout, + mut coordinate: MlirValueUse, + ) -> Result { + let sizes = layout.shape.flatten(); + let strides = layout.stride.flatten(); + let MlirType::Integer(width) = coordinate.ty else { + return Err("CuTe layout coordinate must be an integer".into()); + }; + let mut offset = self.constant(0, width)?; + for (i, (&size, &stride)) in sizes.iter().zip(&strides).enumerate() { + let size = self.constant(size, width)?; + let digit = if i + 1 == sizes.len() { + coordinate.clone() + } else { + self.binary("arith.remui", coordinate.clone(), size.clone())? + }; + let stride = self.constant(stride, width)?; + let term = self.binary("arith.muli", digit, stride)?; + offset = self.binary("arith.addi", offset, term)?; + if i + 1 != sizes.len() { + coordinate = self.binary("arith.divui", coordinate, size)?; + } + } + Ok(offset) + } + /// Sum row and column offsets from a verified two-mode layout. + pub(in crate::cute) fn linear_pair( + &mut self, + layout: &Layout, + row: MlirValueUse, + column: MlirValueUse, + ) -> Result { + let modes = layout.modes(); + if modes.len() != 2 { + return Err("CuTe attention shared layout requires two modes".into()); + } + let row = self.index(&modes[0], row)?; + let column = self.index(&modes[1], column)?; + self.binary("arith.addi", row, column) + } + /// Apply the outer row/column layout, the fixed offset, and the inner swizzle in order. + pub(in crate::cute) fn composed( + &mut self, + layout: &ComposedLayout, + row: MlirValueUse, + column: MlirValueUse, + ) -> Result { + let MlirType::Integer(width) = row.ty else { + return Err("CuTe layout coordinate must be an integer".into()); + }; + let plain = self.linear_pair(layout.outer(), row, column)?; + let offset = self.constant(layout.offset(), width)?; + let plain = self.binary("arith.addi", plain, offset)?; + self.swizzle(layout.inner(), plain) + } + /// Apply the layout's XOR bit permutation without changing the offset's integer width. + pub(in crate::cute) fn swizzle( + &mut self, + swizzle: Swizzle, + plain: MlirValueUse, + ) -> Result { + let MlirType::Integer(width) = plain.ty else { + return Err("CuTe swizzle offset must be an integer".into()); + }; + if swizzle.num_bits == 0 { + return Ok(plain); + } + let mask = self.constant(swizzle.y_mask(), width)?; + let bits = self.binary("arith.andi", plain.clone(), mask)?; + let shift = self.constant(i64::from(swizzle.num_shift.unsigned_abs()), width)?; + let moved = self.binary( + if swizzle.num_shift >= 0 { + "arith.shrui" + } else { + "arith.shli" + }, + bits, + shift, + )?; + self.binary("arith.xori", plain, moved) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/mod.rs new file mode 100644 index 0000000000..4e677dff24 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/mod.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Cooperative and TMA transfers, including their layout-aware address builders. + +pub(in crate::cute) mod cooperative; +pub(in crate::cute) mod emit; +pub(in crate::cute) mod registry; +pub(in crate::cute) mod sm100_attention; +pub(in crate::cute) mod tma; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/registry.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/registry.rs new file mode 100644 index 0000000000..b645d795a4 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/registry.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Single ownership of cooperative-copy and matrix-load profile registrations. + +use crate::cute::{copy::cooperative::CopyTranslation, mma::ldmatrix::MatrixTranslation}; +use crate::{DropAttribute, TranslationError, TranslationRegistry}; +use dialect_cute::{ + attributes::{ + CuteComposedLayoutAttr, CuteCopyAtomAttr, CuteDivisibilityAttr, CuteLayoutAttr, + CuteMatrixRoleAttr, + }, + ops::{CuteCopyG2SOp, CuteLdmatrixOp}, +}; + +/// Register cooperative-copy and matrix-load operations plus their consumed plan attributes. +pub(crate) fn register_cute_attention_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(CopyTranslation)?; + registry.register_operation::(MatrixTranslation)?; + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/sm100_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/sm100_attention.rs new file mode 100644 index 0000000000..adf3e93f02 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/sm100_attention.rs @@ -0,0 +1,89 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Independent SM100 attention TMA operations and their descriptor contracts. + +use crate::cute::{ + common_sm100::contract, + emit::{Emit, integer_property, llvm_ptr_type, property, segments}, +}; +use crate::{MlirOperation, MlirType, OperationInput, OperationTranslation, TranslationSession}; +use dialect_cute::sm100_attention_ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Emit independent SM100 attention loads or stores using the verified descriptor plan. +pub(in crate::cute) struct Tma( + #[doc = "True selects shared-to-global stores; false selects global-to-shared loads."] + pub(in crate::cute) bool, +); +impl OperationTranslation for Tma { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let p = if self.0 { + let o = CuteSm100TmaStoreOp::wrap(source); + o.verify(ctx).map_err(|e| e.to_string())?; + o.plan(ctx) + } else { + let o = CuteSm100TmaLoadOp::wrap(source); + o.verify(ctx).map_err(|e| e.to_string())?; + o.plan(ctx) + } + .ok_or("missing SM100 TMA plan")?; + contract(&input)?; + let mut e = Emit { + session, + location: input.location, + ops: vec![], + }; + let a = input.operands; + let desc = e.cute_ptr( + a[0].clone(), + "!cute_nvgpu.tma_descriptor_tiled", + "generic", + 64, + None, + )?; + let data = e.cute_ptr(a[1].clone(), "f16", "smem", 128, None)?; + if self.0 { + let mut op = e.operation( + "cute_nvgpu.arch.copy.SM100.tma_store", + vec![desc, data, a[2].clone(), a[3].clone(), a[4].clone()], + )?; + property(&mut op, "mode", "#cute_nvgpu.tma_store_mode")?; + segments(&mut op, &[1, 1, 3, 0]); + e.ops.push(op); + } else { + let local = e.llvm_ptr(a[5].clone(), 3)?; + let bar = if p.cta_group == 2 { + let zero = e.constant(32, 0)?; + let remote = e.value("nvvm.mapa", vec![local, zero], llvm_ptr_type(7)?)?; + let bits = e.cast("llvm.ptrtoint", remote, MlirType::Integer(32))?; + e.cast("llvm.inttoptr", bits, llvm_ptr_type(3)?)? + } else { + local + }; + let bar = e.cute_ptr(bar, "i64", "smem", 8, None)?; + // Each issuing CTA copies to its own shared tile. Group2 changes + // completion routing, not the number of destination CTAs; the + // ordinary two-CTA TMA operation therefore has no multicast mask. + let args = vec![desc, data, bar, a[2].clone(), a[3].clone(), a[4].clone()]; + let mut op = e.operation("cute_nvgpu.arch.copy.SM100.tma_load", args)?; + property(&mut op, "mode", "#cute_nvgpu.tma_load_mode")?; + integer_property(&mut op, "num_cta", 32, p.cta_group as i128); + segments(&mut op, &[1, 1, 1, 3, 0, 0, 0]); + e.ops.push(op); + } + Ok(e.ops) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/tma.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/tma.rs new file mode 100644 index 0000000000..a85bf384ed --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/tma.rs @@ -0,0 +1,207 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Two-dimensional TMA view and transfer recipes shared by GEMM capabilities. +//! +//! Source coordinates count tiles in `(row, column)` order. NVVM takes element +//! coordinates in `(column, row)` order. The verified source view supplies tile +//! extents even though its translated type is just a pointer. For example: +//! +//! ``` +//! let (rows, columns) = (128, 64); +//! let (tile_row, tile_column) = (2, 3); +//! let nvvm_coordinates = (tile_column * columns, tile_row * rows); +//! assert_eq!(nvvm_coordinates, (192, 256)); +//! ``` + +use crate::cute::common::{ + emit_binary, emit_i64_constant, expect_shape, fresh, generic_pointer_type, operation, + shared_pointer_type, stringify, +}; +use crate::{ + MlirAttribute, MlirLocation, MlirOperation, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationRegistry, TranslationSession, TypeTranslation, +}; +use dialect_cute::{ + attributes::CuteTensorAddressSpaceAttr, epilogue_ops::CuteTmaStore2dSemanticOp, + gemm_tma_ops::CuteTmaCopy2dOp, types::CuteTmaViewType, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Erase a verified TMA view to a generic or CTA-shared pointer according to its space. +pub(in crate::cute) struct TmaViewTypeTranslation; + +impl TypeTranslation for TmaViewTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let view = source_ref + .downcast_ref::() + .ok_or_else(|| "expected cute.tma_view".to_owned())?; + view.verify(ctx).map_err(stringify)?; + let tensor = view + .tensor_view(ctx) + .ok_or_else(|| "cute.tma_view lost its tensor facts".to_owned())?; + match tensor.space { + CuteTensorAddressSpaceAttr::Gmem => generic_pointer_type(), + CuteTensorAddressSpaceAttr::Smem => shared_pointer_type(), + } + } +} + +/// Emit a two-dimensional TMA load with tile coordinates converted to element coordinates. +pub(in crate::cute) struct TmaLoadTranslation; + +impl OperationTranslation for TmaLoadTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteTmaCopy2dOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.tma_copy_2d", &input, 0, 5)?; + let (rows, cols) = tma_tile_extents(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + let mut operations = vec![]; + let column = tma_coordinate( + session, + &mut operations, + input.operands[3].clone(), + cols, + &input.location, + )?; + let row = tma_coordinate( + session, + &mut operations, + input.operands[2].clone(), + rows, + &input.location, + )?; + let mut copy = operation( + "nvvm.cp.async.bulk.tensor.shared.cluster.global", + vec![], + vec![ + input.operands[1].clone(), + input.operands[0].clone(), + column, + row, + input.operands[4].clone(), + ], + &input.location, + )?; + copy.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 1, 0, 0, 0, 0]), + ); + copy.attributes + .insert("isCTAOnly".into(), MlirAttribute::Bool(true)); + operations.push(copy); + Ok(operations) + } +} + +/// Emit a two-dimensional TMA store with NVVM's column-before-row coordinate order. +pub(in crate::cute) struct TmaStoreTranslation; + +impl OperationTranslation for TmaStoreTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteTmaStore2dSemanticOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.tma_store_2d", &input, 0, 4)?; + let (rows, cols) = tma_tile_extents(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + let mut operations = vec![]; + let column = tma_coordinate( + session, + &mut operations, + input.operands[3].clone(), + cols, + &input.location, + )?; + let row = tma_coordinate( + session, + &mut operations, + input.operands[2].clone(), + rows, + &input.location, + )?; + let mut copy = operation( + "nvvm.cp.async.bulk.tensor.global.shared.cta", + vec![], + vec![ + input.operands[1].clone(), + input.operands[0].clone(), + column, + row, + ], + &input.location, + )?; + copy.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 0, 0]), + ); + operations.push(copy); + Ok(operations) + } +} + +/// Read positive row and column element counts from the source view's shared layout. +fn tma_tile_extents(ctx: &Context, ty: TypeHandle) -> Result<(u64, u64), String> { + let ty_ref = ty.deref(ctx); + let view = ty_ref + .downcast_ref::() + .ok_or_else(|| "expected a cute.tma_view operand".to_owned())?; + let modes = view.smem_layout.0.outer().modes(); + if modes.len() != 2 { + return Err(format!("TMA mapping needs two modes, got {}", modes.len())); + } + let rows = modes[0] + .checked_size() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| "TMA row extent is invalid".to_owned())?; + let cols = modes[1] + .checked_size() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| "TMA column extent is invalid".to_owned())?; + Ok((rows, cols)) +} + +/// Multiply a tile index by its extent and produce NVVM's i32 element coordinate. +fn tma_coordinate( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + tile: MlirValueUse, + extent: u64, + location: &MlirLocation, +) -> Result { + let extent = emit_i64_constant(session, operations, extent, location)?; + let elements = emit_binary(session, operations, "arith.muli", tile, extent, location)?; + let (result, use_) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "arith.trunci", + vec![result], + vec![elements], + location, + )?); + Ok(use_) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/emit.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/emit.rs new file mode 100644 index 0000000000..8388fa430c --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/emit.rs @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Operation builder shared by SM100 recipes and NVVM numeric mappings. + +use crate::{ + MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, MlirType, + MlirValueUse, TranslationSession, +}; +use dialect_cute::attributes::CuteComposedLayoutAttr; + +/// Build SM100 and numeric MLIR sequences while sharing SSA allocation and source location. +pub(crate) struct Emit<'a, 'b> { + /// Translation session providing unique SSA and block identifiers. + pub(crate) session: &'a mut TranslationSession<'b>, + /// Source location attached to every operation emitted by this builder. + pub(crate) location: MlirLocation, + /// Target operations accumulated in execution order. + pub(crate) ops: Vec, +} +impl Emit<'_, '_> { + /// Build a result-free operation at this builder's location without appending it. + pub(crate) fn operation( + &self, + name: &str, + operands: Vec, + ) -> Result { + let mut op = MlirOperation::new(name)?; + op.operands = operands; + op.location = self.location.clone(); + Ok(op) + } + /// Allocate one SSA identifier and return matching result and operand handles. + pub(crate) fn fresh(&mut self, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = self.session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) + } + /// Append a target operation with no results. + pub(in crate::cute) fn emit( + &mut self, + name: &str, + operands: Vec, + ) -> Result<(), String> { + let op = self.operation(name, operands)?; + self.ops.push(op); + Ok(()) + } + /// Append an operation with one fresh result of the supplied type. + pub(crate) fn value( + &mut self, + name: &str, + operands: Vec, + ty: MlirType, + ) -> Result { + let (r, v) = self.fresh(ty); + let mut op = self.operation(name, operands)?; + op.results.push(r); + self.ops.push(op); + Ok(v) + } + /// Append a named one-operand conversion to the requested type. + pub(crate) fn cast( + &mut self, + name: &str, + operand: MlirValueUse, + ty: MlirType, + ) -> Result { + self.value(name, vec![operand], ty) + } + /// Append a binary operation with the left operand's result type. + pub(crate) fn binary( + &mut self, + name: &str, + a: MlirValueUse, + b: MlirValueUse, + ) -> Result { + let ty = a.ty.clone(); + self.value(name, vec![a, b], ty) + } + /// Append a signless integer constant with the requested bit width. + pub(crate) fn constant(&mut self, bits: u32, value: i128) -> Result { + let ty = MlirType::Integer(bits); + let (r, v) = self.fresh(ty.clone()); + let mut op = self.operation("arith.constant", vec![])?; + op.results.push(r); + op.properties + .insert("value".into(), MlirAttribute::Integer { value, ty }); + self.ops.push(op); + Ok(v) + } + /// Convert a pointer to the requested LLVM address space only when needed. + pub(crate) fn llvm_ptr( + &mut self, + ptr: MlirValueUse, + space: u32, + ) -> Result { + let ty = llvm_ptr_type(space)?; + if ptr.ty == ty { + Ok(ptr) + } else { + self.cast("llvm.addrspacecast", ptr, ty) + } + } + /// Bridge an LLVM pointer to a CuTe pointer with byte alignment and optional byte swizzle. + pub(crate) fn cute_ptr( + &mut self, + ptr: MlirValueUse, + element: &str, + space: &str, + align: u32, + swizzle: Option, + ) -> Result { + let ptr = self.llvm_ptr(ptr, if space == "smem" { 3 } else { 0 })?; + let suffix = swizzle.map(|s| format!(",{s}")).unwrap_or_default(); + self.cast( + "builtin.unrealized_conversion_cast", + ptr, + dialect(&format!( + "!cute.ptr<{element},{space},align<{align}>{suffix}>" + ))?, + ) + } + /// Append a pointer offset measured in elements of the supplied LLVM type. + pub(in crate::cute) fn gep( + &mut self, + ptr: MlirValueUse, + index: MlirValueUse, + elem: MlirType, + ) -> Result { + let (r, v) = self.fresh(ptr.ty.clone()); + let mut op = self.operation("llvm.getelementptr", vec![ptr, index])?; + op.results.push(r); + op.properties + .insert("elem_type".into(), MlirAttribute::Type(elem)); + op.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![i32::MIN]), + ); + self.ops.push(op); + Ok(v) + } + /// Return a CTA-shared barrier pointer, optionally offset by eight-byte barrier slots. + pub(in crate::cute) fn barrier( + &mut self, + ptr: MlirValueUse, + stage: Option, + ) -> Result { + let ptr = self.llvm_ptr(ptr, 3)?; + match stage { + Some(stage) => self.gep(ptr, stage, MlirType::Integer(64)), + None => Ok(ptr), + } + } + /// Build a K-major FP16 shared descriptor, converting element swizzles to byte swizzles. + pub(in crate::cute) fn descriptor( + &mut self, + ptr: MlirValueUse, + layout: &CuteComposedLayoutAttr, + rows: u32, + ) -> Result { + // Source composed layouts count FP16 elements. CuTe pointers carry + // BYTE-address swizzles, so the base advances by log2(sizeof(f16)). + let swizzle = layout.0.inner(); + let ptr = self.cute_ptr( + ptr, + "f16", + "smem", + 128, + Some(format!( + "S<{},{},{}>", + swizzle.num_bits, + swizzle.num_base + 1, + swizzle.num_shift + )), + )?; + let (result, desc) = self.fresh(dialect("!cute_nvgpu.smem_desc")?); + let mut op = self.operation("cute_nvgpu.make_umma_smem_desc", vec![ptr])?; + op.results.push(result); + property( + &mut op, + "layout", + &format!("#cute.layout<\"(({rows},16),1):((64,1),0)\">"), + )?; + property(&mut op, "major", "#cute_nvgpu.major")?; + self.ops.push(op); + self.cast( + "builtin.unrealized_conversion_cast", + desc, + MlirType::Integer(64), + ) + } + /// Append the selected before/after TCGen05 ordering fence. + pub(in crate::cute) fn fence(&mut self, kind: &str) -> Result<(), String> { + let mut op = self.operation("nvvm.tcgen05.fence", vec![])?; + property(&mut op, "kind", &format!("#nvvm.tcgen05_fence<{kind}>"))?; + self.ops.push(op); + Ok(()) + } + /// Wrap operations in one fresh argument-free block and region. + pub(in crate::cute) fn region(&mut self, operations: Vec) -> MlirRegion { + MlirRegion { + blocks: vec![MlirBlock { + id: self.session.fresh_block(), + arguments: vec![], + operations, + }], + } + } + /// Retry bounded parity waits until completion while preserving cluster acquire semantics. + pub(in crate::cute) fn wait_cluster( + &mut self, + barrier: MlirValueUse, + phase: MlirValueUse, + ) -> Result<(), String> { + // Preserve cluster acquire semantics and retry on a bounded try-wait; + // the older high-level NVVM blocking operation only has CTA scope. + let start = self.ops.len(); + let (r, done) = self.fresh(MlirType::Integer(1)); + let mut wait = self.operation("nvvm.mbarrier.wait.parity", vec![barrier, phase])?; + wait.results.push(r); + property(&mut wait, "kind", "#nvvm.mbar_wait")?; + property(&mut wait, "scope", "#nvvm.mbar_scope")?; + property(&mut wait, "order", "#nvvm.mem_order")?; + self.ops.push(wait); + let yes = self.constant(1, 1)?; + let again = self.binary("arith.xori", done, yes)?; + self.emit("scf.condition", vec![again])?; + let before = self.ops.split_off(start); + let yield_op = self.operation("scf.yield", vec![])?; + let mut loop_op = self.operation("scf.while", vec![])?; + loop_op.regions = vec![self.region(before), self.region(vec![yield_op])]; + self.ops.push(loop_op); + Ok(()) + } +} +/// Wrap a checked type token; the target backend parses its dialect syntax and meaning. +pub(in crate::cute) fn dialect(s: &str) -> Result { + MlirType::dialect(s) +} +/// Spell an opaque LLVM pointer in the requested numeric address space. +pub(crate) fn llvm_ptr_type(space: u32) -> Result { + dialect(&if space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{space}>") + }) +} +/// Attach a checked textual dialect attribute as a target property. +pub(crate) fn property(op: &mut MlirOperation, key: &str, value: &str) -> Result<(), String> { + op.properties + .insert(key.into(), MlirAttribute::dialect(value)?); + Ok(()) +} +/// Attach an integer property with the requested bit width. +pub(crate) fn integer_property(op: &mut MlirOperation, key: &str, bits: u32, value: i128) { + op.properties.insert( + key.into(), + MlirAttribute::Integer { + value, + ty: MlirType::Integer(bits), + }, + ); +} +/// Record flattened operand segment lengths for the target operation. +pub(crate) fn segments(op: &mut MlirOperation, values: &[i32]) { + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(values.to_vec()), + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/mod.rs new file mode 100644 index 0000000000..b584c57570 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/mod.rs @@ -0,0 +1,9 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Architecture-specific accumulator conversion and shared output storage. + +pub(in crate::cute) mod sm100; +pub(in crate::cute) mod sm120; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm100.rs new file mode 100644 index 0000000000..ad9735169e --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm100.rs @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 tensor-memory accumulator load, bias, conversion, and shared stores. + +use crate::cute::emit::{Emit, dialect, integer_property}; +use crate::{MlirAttribute, MlirFloatType, MlirType, MlirValueUse}; +use dialect_cute::sm100_ops::*; + +/// Load a TMEM vector, add scalar bias, convert to FP16, and store swizzled shared chunks. +pub(in crate::cute) fn emit_epilogue( + e: &mut Emit<'_, '_>, + a: &[MlirValueUse], + p: &CuteSm100EpiloguePlanAttr, +) -> Result<(), String> { + let ptr = e.cast( + "builtin.unrealized_conversion_cast", + a[0].clone(), + dialect("!cute.ptr>")?, + )?; + let (result, regs) = e.fresh(MlirType::Vector { + shape: vec![32], + element: Box::new(MlirType::Integer(32)), + }); + let mut load = e.operation("cute_nvgpu.arch.copy.SM100.tmem_load", vec![ptr])?; + load.results.push(result); + for name in ["num_dp", "num_b", "num_rep"] { + integer_property(&mut load, name, 32, 32); + } + e.ops.push(load); + // The native CuTe TMEM load is a synchronous copy operation: CUTLASS + // inserts its tcgen05.wait::ld before making the register vector available. + // Keep the register tile vector-valued through the bias and conversion. + // Scalar extraction here prevents packed FP16 conversion and shared stores. + let floats = MlirType::Vector { + shape: vec![32], + element: Box::new(MlirType::Float(MlirFloatType::F32)), + }; + let values = e.cast("arith.bitcast", regs, floats.clone())?; + let bias = e.cast("vector.broadcast", a[3].clone(), floats)?; + let biased = e.binary("arith.addf", values, bias)?; + let halves = e.cast( + "arith.truncf", + biased, + MlirType::Vector { + shape: vec![32], + element: Box::new(MlirType::Float(MlirFloatType::F16)), + }, + )?; + let dst = e.llvm_ptr(a[1].clone(), 3)?; + let stride = e.constant(32, p.columns as i128)?; + let row = e.binary("arith.muli", a[2].clone(), stride)?; + let swizzle = p.layout.0.inner(); + let mask = e.constant( + 32, + (((1u32 << swizzle.num_bits) - 1) << swizzle.num_base) as i128, + )?; + let shift = e.constant(32, swizzle.num_shift as i128)?; + // The verified B64 layout has element swizzle S<2,3,3>: it only + // changes bits 3 and 4 using bits 6 and 7. Within each eight-element + // chunk the low three bits therefore remain ordered and contiguous. + // Swizzle each chunk's base once, then store its eight FP16 values. + for col in (0..p.columns).step_by(8) { + let (result, chunk) = e.fresh(MlirType::Vector { + shape: vec![8], + element: Box::new(MlirType::Float(MlirFloatType::F16)), + }); + let mut extract = e.operation("vector.shuffle", vec![halves.clone(), halves.clone()])?; + extract.results.push(result); + extract.properties.insert( + "mask".into(), + MlirAttribute::DenseI64Array((col..col + 8).map(i64::from).collect()), + ); + e.ops.push(extract); + let c = e.constant(32, col as i128)?; + let plain = e.binary("arith.addi", row.clone(), c)?; + let shifted = e.binary("arith.shrui", plain.clone(), shift.clone())?; + let swizzled = e.binary("arith.andi", shifted, mask.clone())?; + let offset = e.binary("arith.xori", plain, swizzled)?; + let addr = e.gep(dst.clone(), offset, MlirType::Float(MlirFloatType::F16))?; + let mut store = e.operation("llvm.store", vec![chunk, addr])?; + // The raw public pointer only promises FP16 alignment. Preserve that + // contract: stronger alignment can be inferred from an aligned shared + // allocation and the chunk's multiple-of-eight element offset. + integer_property(&mut store, "alignment", 64, 2); + e.ops.push(store); + } + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm120.rs new file mode 100644 index 0000000000..8012053785 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm120.rs @@ -0,0 +1,548 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM120 fragment stores, TMA-store completion, and epilogue synchronization. +//! +//! Each thread's 64 accumulator values are ordered as two M bands, eight N slots, +//! and four registers per slot. The N offsets select two 128x64 shared halves; +//! their bases are 8192 f16 elements apart. This illustrates the register and +//! half selection, before the warp/lane offsets and shared-memory swizzle: +//! +//! ``` +//! let n_offsets = [0, 8, 32, 40, 64, 72, 96, 104]; +//! let (m_band, n_slot) = (0, 4); +//! let first_register = (m_band * 8 + n_slot) * 4; +//! let half_base = (n_offsets[n_slot] / 64) * 8192; +//! assert_eq!((first_register, half_base), (16, 8192)); +//! ``` + +use crate::cute::{ + common::{ + cast_int, emit_binary, emit_i32_constant, emit_i64_constant, expect_shape, fresh, + gep_properties, gep_value, integer_constant, operation, stringify, + }, + mma::sm120::{CarrierLayout, extract_carrier}, +}; +use crate::{ + MlirAttribute, MlirLocation, MlirOperation, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationSession, +}; +use dialect_cute::{ + attributes::{CuteEpilogueSyncPhaseAttr, CuteTmaStorePipelineAttr}, + epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, CuteTmaStoreAcquireOp, + CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::Typed, +}; + +/// NVVM operation spelling for a barrier with an explicit participant count. +const COUNTED_CTA_SYNC_OP: &str = "nvvm.barrier.cta.sync"; +/// Lower store acquire/tail to a bulk-group wait with the verified pending-group limit. +pub(in crate::cute) struct TmaStoreWaitTranslation { + /// Source operation selecting queue acquisition or complete queue draining. + pub(in crate::cute) name: &'static str, +} + +impl OperationTranslation for TmaStoreWaitTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let group = match self.name { + "cute.tma_store_acquire" => { + let semantic = CuteTmaStoreAcquireOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + semantic + .pipeline(ctx) + .and_then(CuteTmaStorePipelineAttr::max_pending) + .ok_or_else(|| "tma_store_acquire has no valid pipeline".to_owned())? + } + "cute.tma_store_tail" => { + CuteTmaStoreTailOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + 0 + } + _ => return Err(format!("unknown TMA store wait mapping {}", self.name)), + }; + expect_shape(self.name, &input, 0, 0)?; + let mut wait = operation( + "nvvm.cp.async.bulk.wait_group", + vec![], + vec![], + &input.location, + )?; + wait.properties.insert( + "group".into(), + MlirAttribute::Integer { + value: i128::from(group), + ty: MlirType::Integer(32), + }, + ); + wait.attributes.insert("read".into(), MlirAttribute::Unit); + Ok(vec![wait]) + } +} + +/// Map a verified result-free, operand-free source operation to its target spelling. +pub(in crate::cute) struct NoOperandRenameTranslation( + #[doc = "Target operation spelling for the verified zero-operand source operation."] + pub(in crate::cute) &'static str, +); + +impl OperationTranslation for NoOperandRenameTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + if self.0 == "nvvm.cp.async.bulk.commit.group" { + CuteTmaStoreCommitOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + } + expect_shape("no-operand rename", &input, 0, 0)?; + Ok(vec![operation(self.0, vec![], vec![], &input.location)?]) + } +} + +/// Select a shared FP16 half and return its pointer plus element count. +pub(in crate::cute) struct EpilogueHalfTranslation; + +impl OperationTranslation for EpilogueHalfTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = CuteEpilogueHalfOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + input.attributes.remove("epilogue_half_tile"); + expect_shape("cute.epilogue_half", &input, 2, 1)?; + let tile = semantic + .tile(ctx) + .ok_or_else(|| "epilogue_half has no tile".to_owned())?; + let half = semantic + .half(ctx) + .ok_or_else(|| "epilogue_half has no index".to_owned())?; + let half_elements = tile + .half_elements() + .ok_or_else(|| "epilogue half size overflowed".to_owned())?; + let mut operations = vec![]; + if half.0 == 0 { + operations.push(operation( + "builtin.unrealized_conversion_cast", + vec![input.results[0].clone()], + vec![input.operands[0].clone()], + &input.location, + )?); + } else { + let index = + emit_i64_constant(session, &mut operations, half_elements, &input.location)?; + let mut gep = operation( + "llvm.getelementptr", + vec![input.results[0].clone()], + vec![input.operands[0].clone(), index], + &input.location, + )?; + gep_properties(&mut gep, epilogue_gep_element_type(), vec![i32::MIN]); + operations.push(gep); + } + operations.push(integer_constant( + input.results[1].clone(), + i128::from(half_elements), + &input.location, + )?); + Ok(operations) + } +} + +/// Publish shared stores when required, then synchronize the counted epilogue participants. +pub(in crate::cute) struct EpilogueSyncTranslation; + +impl OperationTranslation for EpilogueSyncTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = CuteEpilogueSyncOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + input.attributes.remove("epilogue_sync_tile"); + expect_shape("cute.epilogue_sync", &input, 0, 1)?; + let phase = semantic + .phase(ctx) + .ok_or_else(|| "epilogue_sync has no phase".to_owned())?; + let barrier = semantic + .barrier(ctx) + .ok_or_else(|| "epilogue_sync has no barrier".to_owned())?; + let threads = barrier + .participant_threads() + .ok_or_else(|| "epilogue barrier thread count overflowed".to_owned())?; + let mut operations = vec![]; + let id = emit_i32_constant( + session, + &mut operations, + i128::from(barrier.barrier_id), + &input.location, + )?; + let threads = emit_i32_constant( + session, + &mut operations, + i128::from(threads), + &input.location, + )?; + operations.extend(epilogue_sync_operations( + phase, + id, + threads, + &input.location, + )?); + Ok(operations) + } +} + +/// Pack the verified 64-register accumulator and store its two swizzled shared halves. +pub(in crate::cute) struct EpilogueStoreTranslation; + +impl OperationTranslation for EpilogueStoreTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteEpilogueStoreFragmentOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + input.attributes.remove("epilogue_store_fragment_tile"); + expect_shape("cute.epilogue_store_fragment", &input, 0, 4)?; + let accumulator_layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_operand(3).get_type(ctx), + session, + 0, + )?; + let mut operations = vec![]; + let mut accumulator = vec![]; + extract_carrier( + session, + &mut operations, + input.operands[3].clone(), + &accumulator_layout, + &mut accumulator, + &input.location, + )?; + if accumulator.len() != 64 { + return Err("cute.epilogue_store_fragment needs 64 f32 leaves".into()); + } + + let warp_i32 = cast_int( + session, + &mut operations, + input.operands[1].clone(), + MlirType::Integer(32), + "arith.trunci", + &input.location, + )?; + let three = emit_i32_constant(session, &mut operations, 3, &input.location)?; + let warp_m = emit_binary( + session, + &mut operations, + "arith.andi", + warp_i32.clone(), + three, + &input.location, + )?; + let two = emit_i32_constant(session, &mut operations, 2, &input.location)?; + let warp_n = emit_binary( + session, + &mut operations, + "arith.shrui", + warp_i32, + two, + &input.location, + )?; + let lane = input.operands[2].clone(); + let fifteen = emit_i32_constant(session, &mut operations, 15, &input.location)?; + let lane15 = emit_binary( + session, + &mut operations, + "arith.andi", + lane.clone(), + fifteen, + &input.location, + )?; + let warp_row_stride = emit_i32_constant(session, &mut operations, 1024, &input.location)?; + let warp_rows = emit_binary( + session, + &mut operations, + "arith.muli", + warp_m, + warp_row_stride, + &input.location, + )?; + let lane_row_stride = emit_i32_constant(session, &mut operations, 64, &input.location)?; + let lane_rows = emit_binary( + session, + &mut operations, + "arith.muli", + lane15, + lane_row_stride, + &input.location, + )?; + let row_zero = emit_binary( + session, + &mut operations, + "arith.addi", + warp_rows, + lane_rows, + &input.location, + )?; + let band_stride = emit_i32_constant(session, &mut operations, 4096, &input.location)?; + let row_one = emit_binary( + session, + &mut operations, + "arith.addi", + row_zero.clone(), + band_stride, + &input.location, + )?; + let columns_per_warp = emit_i32_constant(session, &mut operations, 16, &input.location)?; + let column_base = emit_binary( + session, + &mut operations, + "arith.muli", + warp_n, + columns_per_warp, + &input.location, + )?; + let seven = emit_i32_constant(session, &mut operations, 7, &input.location)?; + let lane7 = emit_binary( + session, + &mut operations, + "arith.andi", + lane, + seven, + &input.location, + )?; + let eight = emit_i32_constant(session, &mut operations, 8, &input.location)?; + let swizzle = emit_binary( + session, + &mut operations, + "arith.muli", + lane7, + eight, + &input.location, + )?; + let half_stride = emit_i32_constant(session, &mut operations, 8192, &input.location)?; + + /// Logical column offsets for the eight accumulator slots within each M band. + const N_OFFSETS: [i128; 8] = [0, 8, 32, 40, 64, 72, 96, 104]; + for m_band in 0..2_usize { + for (n_slot, n_offset) in N_OFFSETS.into_iter().enumerate() { + let local = + emit_i32_constant(session, &mut operations, n_offset % 64, &input.location)?; + let column = emit_binary( + session, + &mut operations, + "arith.addi", + column_base.clone(), + local, + &input.location, + )?; + let logical = emit_binary( + session, + &mut operations, + "arith.addi", + if m_band == 0 { + row_zero.clone() + } else { + row_one.clone() + }, + column, + &input.location, + )?; + let mut physical = emit_binary( + session, + &mut operations, + "arith.xori", + logical, + swizzle.clone(), + &input.location, + )?; + if n_offset >= 64 { + physical = emit_binary( + session, + &mut operations, + "arith.addi", + physical, + half_stride.clone(), + &input.location, + )?; + } + let pointer = gep_value( + session, + &mut operations, + input.operands[0].clone(), + physical, + // `physical` is an f16 element offset. The source path + // performs `base.add(physical)` before casting to u8 for + // stmatrix, preserving each 16-byte row address. + epilogue_gep_element_type(), + &input.location, + )?; + let cell = (m_band * 8 + n_slot) * 4; + let top = emit_f16x2_pack( + session, + &mut operations, + accumulator[cell].clone(), + accumulator[cell + 1].clone(), + &input.location, + )?; + let bottom = emit_f16x2_pack( + session, + &mut operations, + accumulator[cell + 2].clone(), + accumulator[cell + 3].clone(), + &input.location, + )?; + let mut store = operation( + "nvvm.stmatrix", + vec![], + vec![pointer, top, bottom], + &input.location, + )?; + store.properties.insert( + "layout".into(), + MlirAttribute::dialect("#nvvm.mma_layout")?, + ); + operations.push(store); + } + } + Ok(operations) + } +} + +/// Convert low/high FP32 lanes to one packed FP16 word using reversed intrinsic operands. +fn emit_f16x2_pack( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + low: MlirValueUse, + high: MlirValueUse, + location: &MlirLocation, +) -> Result { + let vector_type = MlirType::Vector { + shape: vec![2], + element: Box::new(MlirType::Float(crate::MlirFloatType::F16)), + }; + let (call_result, converted) = fresh(session, vector_type); + let mut call = operation( + "llvm.call_intrinsic", + vec![call_result], + // `ff2f16x2` follows PTX source order: its first f32 becomes the + // high half. CuTe's `cvt_f16x2_f32(lo, hi)` contract is low-first, + // so preserve the shared low-first semantic contract by reversing the + // intrinsic operands. + f16x2_intrinsic_operands(low, high).into(), + location, + )?; + call.properties.insert( + "intrin".into(), + MlirAttribute::String("llvm.nvvm.ff2f16x2.rn".into()), + ); + call.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + call.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![2, 0]), + ); + operations.push(call); + let (bits_result, bits) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "llvm.bitcast", + vec![bits_result], + vec![converted], + location, + )?); + Ok(bits) +} + +/// Put the high lane first to preserve the source API's low-first packed result. +pub(in crate::cute) fn f16x2_intrinsic_operands( + low: MlirValueUse, + high: MlirValueUse, +) -> [MlirValueUse; 2] { + [high, low] +} + +/// Build a CTA synchronization using the explicit barrier identifier and thread count. +fn counted_cta_sync( + barrier_id: MlirValueUse, + participant_threads: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation( + COUNTED_CTA_SYNC_OP, + vec![], + vec![barrier_id, participant_threads], + location, + ) +} + +/// Publish generic CTA-shared stores to the asynchronous proxy used by TMA. +fn async_shared_proxy_fence(location: &MlirLocation) -> Result { + let mut fence = operation("nvvm.fence.proxy", vec![], vec![], location)?; + fence.attributes.insert( + "kind".into(), + MlirAttribute::dialect("#nvvm.proxy_kind")?, + ); + fence.attributes.insert( + "space".into(), + MlirAttribute::dialect("#nvvm.shared_space")?, + ); + Ok(fence) +} + +/// For ReadyForTma emit a proxy fence before counted sync; Reusable needs only sync. +pub(in crate::cute) fn epilogue_sync_operations( + phase: CuteEpilogueSyncPhaseAttr, + barrier_id: MlirValueUse, + participant_threads: MlirValueUse, + location: &MlirLocation, +) -> Result, String> { + let mut operations = Vec::with_capacity(match phase { + CuteEpilogueSyncPhaseAttr::Reusable => 1, + CuteEpilogueSyncPhaseAttr::ReadyForTma => 2, + }); + if phase == CuteEpilogueSyncPhaseAttr::ReadyForTma { + operations.push(async_shared_proxy_fence(location)?); + } + operations.push(counted_cta_sync(barrier_id, participant_threads, location)?); + Ok(operations) +} + +/// Return f16 because epilogue pointer offsets count elements rather than bytes. +pub(in crate::cute) fn epilogue_gep_element_type() -> MlirType { + MlirType::Float(crate::MlirFloatType::F16) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/independent_sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/independent_sm100.rs new file mode 100644 index 0000000000..d5cbac4f10 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/independent_sm100.rs @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Independent SM100 attention pack registration. + +use crate::cute::{copy::sm100_attention::Tma, mma::sm100_attention::Mma}; +use crate::{DropAttribute, TranslationError, TranslationRegistry}; +use dialect_cute::sm100_attention_ops::*; + +/// Register standalone SM100 attention MMA and TMA operations without GEMM-plan ownership. +pub(crate) fn register_cute_sm100_attention_pack( + r: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + r.register_attribute::(DropAttribute)?; + r.register_attribute::(DropAttribute)?; + r.register_operation::(Mma(false))?; + r.register_operation::(Mma(true))?; + r.register_operation::(Tma(false))?; + r.register_operation::(Tma(true))?; + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/ldmatrix.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/ldmatrix.rs new file mode 100644 index 0000000000..8ddc89425e --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/ldmatrix.rs @@ -0,0 +1,290 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Matrix-load address plans and ldmatrix instruction emission. + +use crate::cute::copy::emit::{Emitter, integer, require_input}; +use crate::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationSession, +}; +use cute_layout::{ComposedLayout, IntTuple, OffsetUnit}; +use dialect_cute::{attributes::CuteMatrixRoleAttr, ops::CuteLdmatrixOp}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Emit lane-specific ldmatrix addresses and the verified transpose/register-count variant. +pub(in crate::cute) struct MatrixTranslation; + +/// Bound every intermediate in the byte-offset calculation, not just its final +/// value: negative strides/offsets could otherwise cancel an overflowing term. +/// The verified layout and the unsafe matrix-load API promise a hardware lane +/// (0..32) and an in-bounds selected matrix window. Thus the window coordinates, +/// their row/column expansions, and every nested-layout digit are bounded by +/// the positive shape products. This is not a check of dynamic caller inputs. +pub(in crate::cute) fn matrix_offset_fits_i32(layout: &ComposedLayout) -> bool { + let limit = i64::from(i32::MAX); + if layout.unit() != OffsetUnit::Bytes + || layout.outer().modes().len() != 2 + || !layout + .outer() + .checked_size() + .is_some_and(|size| size <= limit) + || !(0..=limit).contains(&layout.offset()) + { + return false; + } + let mut maximum = layout.offset(); + for (&extent, &stride) in layout + .outer() + .shape + .flatten() + .iter() + .zip(&layout.outer().stride.flatten()) + { + if !(0..=limit).contains(&stride) { + return false; + } + let Some(next) = (extent - 1) + .checked_mul(stride) + .and_then(|term| maximum.checked_add(term)) + else { + return false; + }; + if next > limit { + return false; + } + maximum = next; + } + // Keeping both fields below the sign bit bounds the mask, shifted bits, + // and final XOR even when swizzling raises the maximum physical offset. + let swizzle = layout.inner(); + swizzle.is_valid() + && (0..=limit).contains(&swizzle.y_mask()) + && (0..=limit).contains(&swizzle.z_mask()) +} + +/// A proved split between each window base and its lane-dependent byte offsets. +pub(in crate::cute) struct MatrixAddressPlan { + /// Low address bits up to the highest bit touched by the swizzle. + low_mask: i64, + /// Remaining nonnegative i32 address bits kept in the uniform window base. + high_mask: i64, +} + +/// Prove a carry-free split between window-only and lane-only offsets. Nested +/// layout indexing need not distribute over coordinate addition, and a fixed +/// offset may overlap lane bits, so compare the factored result with the full +/// checked map for every legal window/lane. The existing static-validation cap +/// bounds this proof to at most one lane address per four tile elements. +pub(in crate::cute) fn matrix_address_plan( + layout: &ComposedLayout, + role: CuteMatrixRoleAttr, +) -> Option { + if !matrix_offset_fits_i32(layout) + || layout.outer().checked_size()? > cute_layout::validation::MAX_STATIC_VALIDATION_ELEMENTS + { + return None; + } + let modes = layout.outer().modes(); + let rows = modes[0].checked_size()?; + let columns = modes[1].checked_size()?; + let window_columns = if role == CuteMatrixRoleAttr::B { 8 } else { 16 }; + if rows < 16 || columns < window_columns { + return None; + } + let index = |mode: usize, coordinate| modes[mode].checked_call(&IntTuple::Leaf(coordinate)); + let swizzle = layout.inner(); + let active = swizzle.y_mask() | swizzle.z_mask(); + let low_mask = (1i64 << (i64::BITS - active.leading_zeros())) - 1; + let high_mask = i64::from(i32::MAX) ^ low_mask; + for window_row in 0..rows / 16 { + for window_column in 0..columns / window_columns { + let tile_row = window_row * 16; + let tile_column = window_column * window_columns; + let uniform = layout.offset() + index(0, tile_row)? + index(1, tile_column)?; + for lane in 0..32 { + let (local_row, local_column) = if role == CuteMatrixRoleAttr::BPairKMajor { + (lane % 8 + lane / 16 * 8, lane / 8 % 2 * 8) + } else { + ( + lane % 16, + if role == CuteMatrixRoleAttr::B { + 0 + } else { + lane / 16 * 8 + }, + ) + }; + let local = index(0, local_row)? + index(1, local_column)?; + if uniform & local != 0 { + return None; + } + let high = uniform & high_mask; + let low = swizzle.apply(uniform & low_mask) ^ swizzle.apply(local); + let coordinate = tile_row + local_row + rows * (tile_column + local_column); + let original = layout.checked_call(&IntTuple::Leaf(coordinate))?; + if high & low != 0 || high.checked_add(low)? != original { + return None; + } + } + } + } + Some(MatrixAddressPlan { + low_mask, + high_mask, + }) +} + +impl OperationTranslation for MatrixTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let load = CuteLdmatrixOp::wrap(source); + load.verify(ctx).map_err(|e| e.to_string())?; + input.attributes.remove("ldmatrix_elem"); + require_input(&input, 5)?; + let role = *load.get_attr_matrix_role(ctx).unwrap(); + let (count, columns, transpose) = match role { + CuteMatrixRoleAttr::A | CuteMatrixRoleAttr::BPairKMajor => (4, 16, false), + CuteMatrixRoleAttr::B => (2, 8, true), + CuteMatrixRoleAttr::BPair => (4, 16, true), + }; + let layout = load + .get_attr_ldmatrix_smem_layout(ctx) + .unwrap() + .0 + .to_byte_offsets(2) + .map_err(|e| e.to_string())?; + let mut e = Emitter::new(session, input.location); + let shared = e.pointer(input.operands[0].clone(), 3)?; + let width = if matrix_offset_fits_i32(&layout) { + 32 + } else { + 64 + }; + let (window_row, window_column, lane) = if width == 32 { + ( + e.value( + "arith.trunci", + MlirType::Integer(32), + vec![input.operands[1].clone()], + )?, + e.value( + "arith.trunci", + MlirType::Integer(32), + vec![input.operands[2].clone()], + )?, + input.operands[3].clone(), + ) + } else { + ( + input.operands[1].clone(), + input.operands[2].clone(), + e.value( + "arith.extui", + MlirType::Integer(64), + vec![input.operands[3].clone()], + )?, + ) + }; + let eight = e.constant(8, width)?; + let two = e.constant(2, width)?; + let sixteen = e.constant(16, width)?; + let column_extent = e.constant(columns, width)?; + let submatrix = e.binary("arith.divui", lane.clone(), eight.clone())?; + let lane_row = e.binary("arith.remui", lane, eight.clone())?; + // Each group of eight lanes supplies one 8x8 matrix to ldmatrix.x4. + // K-major B pairs need [top-left, top-right, bottom-left, bottom-right] + // so adjacent result registers form an MMA B operand. A fragments + // and transposed B pairs retain their original column-first ordering. + let (row_half_op, column_half_op) = if role == CuteMatrixRoleAttr::BPairKMajor { + ("arith.divui", "arith.remui") + } else { + ("arith.remui", "arith.divui") + }; + let row_half = e.binary(row_half_op, submatrix.clone(), two.clone())?; + let row_half = e.binary("arith.muli", row_half, eight.clone())?; + let local_row = e.binary("arith.addi", row_half, lane_row)?; + let window_row = e.binary("arith.muli", window_row, sixteen)?; + let window_column = e.binary("arith.muli", window_column, column_extent)?; + let local_column = if count == 4 { + let half = e.binary(column_half_op, submatrix, two)?; + e.binary("arith.muli", half, eight)? + } else { + e.constant(0, width)? + }; + let source = if let Some(plan) = matrix_address_plan(&layout, role) { + let uniform = e.linear_pair(layout.outer(), window_row, window_column)?; + let offset = e.constant(layout.offset(), 32)?; + let uniform = e.binary("arith.addi", uniform, offset)?; + let local = e.linear_pair(layout.outer(), local_row, local_column)?; + let high_mask = e.constant(plan.high_mask, 32)?; + let high = e.binary("arith.andi", uniform.clone(), high_mask)?; + let low_mask = e.constant(plan.low_mask, 32)?; + let low = e.binary("arith.andi", uniform, low_mask)?; + let low = e.swizzle(layout.inner(), low)?; + let local = e.swizzle(layout.inner(), local)?; + let low = e.binary("arith.xori", low, local)?; + // Keep the high window offset additive so constant matrix-row + // displacements can remain GEP/immediate offsets after lowering. + // The allocation pointer itself never enters the swizzle. + let high = e.value("arith.extui", MlirType::Integer(64), vec![high])?; + let low = e.value("arith.extui", MlirType::Integer(64), vec![low])?; + let base = e.gep(shared, high, MlirType::Integer(8))?; + e.gep(base, low, MlirType::Integer(8))? + } else { + let row = e.binary("arith.addi", window_row, local_row)?; + let column = e.binary("arith.addi", window_column, local_column)?; + let offset = e.composed(&layout, row, column)?; + // Preserve 64-bit pointer arithmetic in both general paths. + let offset = if width == 32 { + e.value("arith.extui", MlirType::Integer(64), vec![offset])? + } else { + offset + }; + e.gep(shared, offset, MlirType::Integer(8))? + }; + let aggregate_type = + MlirType::dialect(format!("!llvm.struct<({})>", vec!["i32"; count].join(", ")))?; + let (result, aggregate) = e.fresh(aggregate_type); + let mut load = e.operation("nvvm.ldmatrix", vec![source])?; + load.results.push(result); + load.properties + .insert("num".into(), integer(count as i128, 32)); + load.properties.insert( + "layout".into(), + MlirAttribute::dialect(if transpose { + "#nvvm.mma_layout" + } else { + "#nvvm.mma_layout" + })?, + ); + e.operations.push(load); + for index in 0..count { + let (result, register) = e.fresh(MlirType::Integer(32)); + let mut extract = e.operation("llvm.extractvalue", vec![aggregate.clone()])?; + extract.results.push(result); + extract.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + e.operations.push(extract); + let index = e.c64(index as i64)?; + let slot = e.gep(input.operands[4].clone(), index, MlirType::Integer(32))?; + let mut store = e.operation("llvm.store", vec![register, slot])?; + store.properties.insert("alignment".into(), integer(4, 64)); + e.operations.push(store); + } + Ok(e.operations) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/mod.rs new file mode 100644 index 0000000000..d9c79ce299 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/mod.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Matrix-load address plans and architecture-specific tensor-core instructions. + +pub(in crate::cute) mod ldmatrix; +pub(in crate::cute) mod sm100; +pub(in crate::cute) mod sm100_attention; +pub(in crate::cute) mod sm120; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100.rs new file mode 100644 index 0000000000..c2b8e07f50 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100.rs @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Paired SM100 GEMM MMA with explicit A-collector lifetime sequencing. + +use crate::MlirValueUse; +use crate::cute::emit::{Emit, llvm_ptr_type, property, segments}; +use dialect_cute::sm100_ops::*; + +/// Emit each K16 MMA step with layout-derived descriptors and explicit A-collector reuse. +pub(in crate::cute) fn emit_mma( + e: &mut Emit<'_, '_>, + a: &[MlirValueUse], + p: &CuteSm100MmaPlanAttr, +) -> Result<(), String> { + let ad = e.descriptor(a[1].clone(), &p.a_layout, p.m / 2)?; + let b0d = e.descriptor(a[2].clone(), &p.b0_layout, p.n0 / 2)?; + let b1d = e.descriptor(a[3].clone(), &p.b1_layout, p.n1 / 2)?; + let n0 = e.constant(32, p.n0 as i128)?; + let second = e.binary("arith.addi", a[0].clone(), n0)?; + let d0 = e.cast("llvm.inttoptr", a[0].clone(), llvm_ptr_type(6)?)?; + let d1 = e.cast("llvm.inttoptr", second, llvm_ptr_type(6)?)?; + let id0 = e.constant(32, instruction_descriptor(p.m, p.n0) as i128)?; + let id1 = e.constant(32, instruction_descriptor(p.m, p.n1) as i128)?; + let yes = e.constant(1, 1)?; + // Each pair owns the A collector's complete lifetime. The descriptor + // advance is K16 FP16 elements expressed in 16-byte descriptor units. + for kstep in 0..p.k / 16 { + let offset = e.constant(64, (kstep * 2) as i128)?; + let av = e.binary("arith.addi", ad.clone(), offset.clone())?; + let bv0 = e.binary("arith.addi", b0d.clone(), offset.clone())?; + let bv1 = e.binary("arith.addi", b1d.clone(), offset)?; + let accumulate = if kstep == 0 { + a[4].clone() + } else { + yes.clone() + }; + for (dest, b, desc, collector) in [ + (d0.clone(), bv0, id0.clone(), "fill"), + (d1.clone(), bv1, id1.clone(), "lastuse"), + ] { + let mut op = e.operation( + "nvvm.tcgen05.mma", + vec![dest, av.clone(), b, desc, accumulate.clone()], + )?; + property(&mut op, "mmaKind", "#nvvm.tcgen05_mma_kind")?; + property(&mut op, "ctaGroup", "#nvvm.cta_group")?; + property( + &mut op, + "collectorOp", + &format!("#nvvm.tcgen05_mma_collectorop<{collector}>"), + )?; + segments(&mut op, &[1, 1, 1, 1, 1, 0, 0]); + e.ops.push(op); + } + } + Ok(()) +} + +/// Encode FP16 MMA shape fields for the verified M and N extents. +pub(in crate::cute) fn instruction_descriptor(m: u32, n: u32) -> u32 { + (1 << 4) | ((n >> 3) << 17) | ((m >> 4) << 24) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs new file mode 100644 index 0000000000..6be8b135dd --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm100_attention.rs @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Independent SM100 attention MMA and layout-derived shared descriptors. + +use crate::cute::{ + common_sm100::contract, + emit::{Emit, llvm_ptr_type, property, segments}, +}; +use crate::{ + MlirOperation, MlirType, MlirValueUse, OperationInput, OperationTranslation, TranslationSession, +}; +use dialect_cute::{attributes::CuteComposedLayoutAttr, sm100_attention_ops::*}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Emit independent SM100 attention MMA with shared or tensor-memory A operands. +pub(in crate::cute) struct Mma( + #[doc = "True for tensor-memory A and shared B; false for two shared-memory operands."] + pub(in crate::cute) bool, +); +impl OperationTranslation for Mma { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let p = if self.0 { + let o = CuteSm100MmaTsOp::wrap(source); + o.verify(ctx).map_err(|e| e.to_string())?; + o.plan(ctx) + } else { + let o = CuteSm100MmaSsOp::wrap(source); + o.verify(ctx).map_err(|e| e.to_string())?; + o.plan(ctx) + } + .ok_or("missing SM100 MMA plan")?; + contract(&input)?; + let mut e = Emit { + session, + location: input.location, + ops: vec![], + }; + let a = input.operands; + let ad = if self.0 { + None + } else { + Some(descriptor( + &mut e, + a[1].clone(), + a[4].clone(), + &p.a_layout, + p.m / p.cta_group, + p.k, + p.a_mn, + )?) + }; + let bd = descriptor( + &mut e, + a[2].clone(), + a[if self.0 { 4 } else { 5 }].clone(), + &p.b_layout, + p.n / p.cta_group, + p.k, + p.b_mn, + )?; + let dst = e.cast("llvm.inttoptr", a[0].clone(), llvm_ptr_type(6)?)?; + let instruction = (1u32 << 4) + | ((p.n >> 3) << 17) + | ((p.m >> 4) << 24) + | ((p.a_mn as u32) << 15) + | ((p.b_mn as u32) << 16); + let inst = e.constant(32, instruction as i128)?; + let yes = e.constant(1, 1)?; + for k in (0..p.k).step_by(16) { + let av = if let Some(ad) = &ad { + let offset = e.constant( + 64, + descriptor_offset(p.m / p.cta_group, p.k, p.a_mn, k) as i128, + )?; + e.binary("arith.addi", ad.clone(), offset)? + } else { + // Two FP16 probability columns occupy one 32-bit TMEM column. + let offset = e.constant(32, (k / 2) as i128)?; + let address = e.binary("arith.addi", a[1].clone(), offset)?; + e.cast("llvm.inttoptr", address, llvm_ptr_type(6)?)? + }; + let offset = e.constant( + 64, + descriptor_offset(p.n / p.cta_group, p.k, p.b_mn, k) as i128, + )?; + let bv = e.binary("arith.addi", bd.clone(), offset)?; + let mut mma = e.operation( + "nvvm.tcgen05.mma", + vec![ + dst.clone(), + av, + bv, + inst.clone(), + if k == 0 { a[3].clone() } else { yes.clone() }, + ], + )?; + property(&mut mma, "mmaKind", "#nvvm.tcgen05_mma_kind")?; + property( + &mut mma, + "ctaGroup", + &format!("#nvvm.cta_group", p.cta_group), + )?; + property( + &mut mma, + "collectorOp", + "#nvvm.tcgen05_mma_collectorop", + )?; + segments(&mut mma, &[1, 1, 1, 1, 1, 0, 0]); + e.ops.push(mma); + } + Ok(e.ops) + } +} +/// The address field is measured in 16-byte units, independently of swizzle. +fn descriptor_offset(rows: u32, k_extent: u32, mn: bool, k: u32) -> u32 { + if mn { + k * rows.min(64) / 8 + } else { + let span = k_extent.min(64); + ((k / span) * rows * span + k % span) / 8 + } +} +/// Build a shared-memory descriptor using the verified major mode and element-to-byte swizzle. +fn descriptor( + e: &mut Emit<'_, '_>, + ptr: MlirValueUse, + byte_offset: MlirValueUse, + layout: &CuteComposedLayoutAttr, + rows: u32, + k: u32, + mn: bool, +) -> Result { + let sw = layout.0.inner(); + let ptr = e.cute_ptr( + ptr, + "f16", + "smem", + 128, + Some(format!( + "S<{},{},{}>", + sw.num_bits, + sw.num_base + 1, + sw.num_shift + )), + )?; + let (result, value) = e.fresh(MlirType::dialect("!cute_nvgpu.smem_desc")?); + let mut op = e.operation("cute_nvgpu.make_umma_smem_desc", vec![ptr])?; + op.results.push(result); + let span = if mn { rows } else { k }.min(64); + let native = if mn { + format!( + "((({span},{}),16),1):(((1,{}),{span}),0)", + rows / span, + span * k + ) + } else { + format!("(({rows},16),1):(({span},1),0)") + }; + property(&mut op, "layout", &format!("#cute.layout<\"{native}\">"))?; + property( + &mut op, + "major", + if mn { + "#cute_nvgpu.major" + } else { + "#cute_nvgpu.major" + }, + )?; + e.ops.push(op); + let base = e.cast( + "builtin.unrealized_conversion_cast", + value, + MlirType::Integer(64), + )?; + // Preserve a loop-invariant allocation descriptor independently of the + // selected tile. Complete-tile offsets cannot change swizzle phase, and + // valid shared allocations cannot carry out of the address field. + let offset = e.cast("arith.extui", byte_offset, MlirType::Integer(64))?; + let shift = e.constant(64, 4)?; + let offset = e.binary("arith.shrui", offset, shift)?; + e.binary("arith.addi", base, offset) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs new file mode 100644 index 0000000000..b778846bbc --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs @@ -0,0 +1,1124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM120 block-scaled MMA carriers, shared loads, and instruction emission. + +use crate::cute::common::{ + cast_int, emit_binary, emit_i32_constant, emit_i64_constant, expect_shape, fresh, gep_value, + operation, stringify, +}; +use crate::{ + MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationSession, +}; +use dialect_cute::smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, CuteTiledGemmOp, +}; +use dialect_mir::types::{MirArrayType, MirStructType, MirTupleType}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Byte-address bits selected by the verified FP16 shared-layout swizzle. +const LDMATRIX_BYTE_SWIZZLE_MASK: u64 = 0x180; +/// Right shift mapping the selected byte-address bits onto their XOR destinations. +const LDMATRIX_BYTE_SWIZZLE_SHIFT: u64 = 3; +/// Recursive physical LLVM carrier layout retaining the logical scalar register order. +#[derive(Clone)] +pub(in crate::cute) struct CarrierLayout { + /// LLVM type used to transport this carrier node. + pub(in crate::cute) target: MlirType, + /// Whether this node is a scalar leaf or a list of physical aggregate slots. + kind: CarrierKind, +} + +/// Distinguish scalar registers from recursively nested aggregate storage. +#[derive(Clone)] +enum CarrierKind { + /// One scalar register with no nested fields. + Leaf, + /// Physical field indices paired with their recursively verified carrier layouts. + Aggregate( + #[doc = "Physical field indices paired with child carrier layouts in logical register order."] + Vec<(i64, CarrierLayout)>, + ), +} + +impl CarrierLayout { + /// Build a carrier tree from MIR aggregate layout, rejecting unsupported shapes or excessive nesting. + pub(in crate::cute) fn from_source( + ctx: &Context, + source: TypeHandle, + session: &TranslationSession<'_>, + depth: usize, + ) -> Result { + if depth > 64 { + return Err("MMA carrier aggregate is unexpectedly deep".into()); + } + let target = session.translate_type(source).map_err(stringify)?; + let source_ref = source.deref(ctx); + if let Some(array) = source_ref.downcast_ref::() { + let count = usize::try_from(array.size()) + .map_err(|_| "MMA carrier array is too large".to_owned())?; + let child = Self::from_source(ctx, array.element_type(), session, depth + 1)?; + return Ok(Self { + target, + kind: CarrierKind::Aggregate( + (0..count) + .map(|index| (index as i64, child.clone())) + .collect(), + ), + }); + } + + let fields = if let Some(tuple) = source_ref.downcast_ref::() { + Some(tuple.get_types().to_vec()) + } else { + source_ref + .downcast_ref::() + .map(|structure| structure.field_types.clone()) + }; + let Some(fields) = fields else { + return Ok(Self { + target, + kind: CarrierKind::Leaf, + }); + }; + drop(source_ref); + + let plan = crate::mir::memory::aggregate_plan(ctx, source, |ty| { + session.translate_type(ty).map_err(stringify) + })?; + if plan.ty != target { + return Err("MMA carrier aggregate disagrees with MIR type translation".into()); + } + let mut children = Vec::new(); + for (declaration, field) in fields.into_iter().enumerate() { + let Some(slot) = plan.decl_to_slot[declaration] else { + continue; + }; + children.push(( + slot as i64, + Self::from_source(ctx, field, session, depth + 1)?, + )); + } + Ok(Self { + target, + kind: CarrierKind::Aggregate(children), + }) + } + + /// Count scalar register leaves while ignoring aggregate nesting. + pub(in crate::cute) fn leaf_count(&self) -> usize { + match &self.kind { + CarrierKind::Leaf => 1, + CarrierKind::Aggregate(children) => { + children.iter().map(|(_, child)| child.leaf_count()).sum() + } + } + } +} + +/// Append LLVM extracts in logical register order through physical aggregate slots. +pub(in crate::cute) fn extract_carrier( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: MlirValueUse, + layout: &CarrierLayout, + leaves: &mut Vec, + location: &MlirLocation, +) -> Result<(), String> { + match &layout.kind { + CarrierKind::Leaf => { + leaves.push(value); + Ok(()) + } + CarrierKind::Aggregate(children) => { + for (slot, child) in children { + let (result, extracted) = fresh(session, child.target.clone()); + let mut extract = operation( + "llvm.extractvalue", + vec![result], + vec![value.clone()], + location, + )?; + extract + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![*slot])); + operations.push(extract); + extract_carrier(session, operations, extracted, child, leaves, location)?; + } + Ok(()) + } + } +} + +/// Reconstruct the original physical aggregate from ordered scalar register results. +fn rebuild_carrier( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + layout: &CarrierLayout, + leaves: &[MlirValueUse], + final_result: MlirResult, + location: &MlirLocation, +) -> Result<(), String> { + let mut cursor = 0; + build_carrier_node( + session, + operations, + layout, + leaves, + &mut cursor, + Some(final_result), + location, + )?; + if cursor != leaves.len() { + return Err(format!( + "MMA carrier consumed {cursor} leaves but received {}", + leaves.len() + )); + } + Ok(()) +} + +/// Recursively insert scalar leaves into one physical carrier node. +fn build_carrier_node( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + layout: &CarrierLayout, + leaves: &[MlirValueUse], + cursor: &mut usize, + final_result: Option, + location: &MlirLocation, +) -> Result { + if matches!(layout.kind, CarrierKind::Leaf) { + let leaf = leaves + .get(*cursor) + .cloned() + .ok_or_else(|| "MMA carrier did not receive enough leaves".to_owned())?; + *cursor += 1; + if let Some(result) = final_result { + let use_ = MlirValueUse { + id: result.id, + ty: result.ty.clone(), + }; + operations.push(operation( + "builtin.unrealized_conversion_cast", + vec![result], + vec![leaf], + location, + )?); + return Ok(use_); + } + return Ok(leaf); + } + + let CarrierKind::Aggregate(children) = &layout.kind else { + unreachable!() + }; + let mut built_children = Vec::with_capacity(children.len()); + for (slot, child) in children { + built_children.push(( + *slot, + build_carrier_node(session, operations, child, leaves, cursor, None, location)?, + )); + } + + let poison_result = if built_children.is_empty() { + final_result + .clone() + .unwrap_or_else(|| fresh(session, layout.target.clone()).0) + } else { + fresh(session, layout.target.clone()).0 + }; + let mut aggregate = MlirValueUse { + id: poison_result.id, + ty: poison_result.ty.clone(), + }; + operations.push(operation( + "llvm.mlir.poison", + vec![poison_result], + vec![], + location, + )?); + let child_count = built_children.len(); + for (index, (slot, child)) in built_children.into_iter().enumerate() { + let result = if index + 1 == child_count { + final_result + .clone() + .unwrap_or_else(|| fresh(session, layout.target.clone()).0) + } else { + fresh(session, layout.target.clone()).0 + }; + let next = MlirValueUse { + id: result.id, + ty: result.ty.clone(), + }; + let mut insert = operation( + "llvm.insertvalue", + vec![result], + vec![aggregate, child], + location, + )?; + insert + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![slot])); + operations.push(insert); + aggregate = next; + } + Ok(aggregate) +} + +/// Fill every scalar leaf of a verified accumulator carrier with the source value. +pub(in crate::cute) struct FragmentFillTranslation; + +impl OperationTranslation for FragmentFillTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteFragmentFillOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.fragment_fill", &input, 1, 1)?; + let layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_result(0).get_type(ctx), + session, + 0, + )?; + if layout.leaf_count() != 64 { + return Err("cute.fragment_fill needs exactly 64 f32 leaves".into()); + } + let leaves = vec![input.operands[0].clone(); 64]; + let mut operations = vec![]; + rebuild_carrier( + session, + &mut operations, + &layout, + &leaves, + input.results[0].clone(), + &input.location, + )?; + Ok(operations) + } +} + +/// Load and arrange scale words for the verified SM120 block-scaled MMA plan. +pub(in crate::cute) struct MmaLoadScalesTranslation; + +impl OperationTranslation for MmaLoadScalesTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteMmaLoadScalesOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + input.attributes.remove("mma_load_scales_a_view"); + input.attributes.remove("mma_load_scales_b_view"); + expect_shape("cute.mma_load_scales", &input, 1, 7)?; + let layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_result(0).get_type(ctx), + session, + 0, + )?; + if layout.leaf_count() != 10 { + return Err("cute.mma_load_scales needs exactly ten u32 leaves".into()); + } + + let mut operations = vec![]; + let lane = input.operands[0].clone(); + let wm = cast_int( + session, + &mut operations, + input.operands[5].clone(), + MlirType::Integer(32), + "arith.trunci", + &input.location, + )?; + let wn = cast_int( + session, + &mut operations, + input.operands[6].clone(), + MlirType::Integer(32), + "arith.trunci", + &input.location, + )?; + let four = emit_i32_constant(session, &mut operations, 4, &input.location)?; + let q = emit_binary( + session, + &mut operations, + "arith.divui", + lane.clone(), + four, + &input.location, + )?; + let three = emit_i32_constant(session, &mut operations, 3, &input.location)?; + let r = emit_binary( + session, + &mut operations, + "arith.andi", + lane, + three, + &input.location, + )?; + let one = emit_i32_constant(session, &mut operations, 1, &input.location)?; + let parity = emit_binary( + session, + &mut operations, + "arith.andi", + r, + one, + &input.location, + )?; + let eight = emit_i32_constant(session, &mut operations, 8, &input.location)?; + let parity_band = emit_binary( + session, + &mut operations, + "arith.muli", + parity, + eight, + &input.location, + )?; + let a_provider = emit_binary( + session, + &mut operations, + "arith.addi", + q.clone(), + parity_band, + &input.location, + )?; + let sixteen = emit_i32_constant(session, &mut operations, 16, &input.location)?; + let a_band = emit_binary( + session, + &mut operations, + "arith.muli", + wm, + sixteen.clone(), + &input.location, + )?; + let a_row0 = emit_binary( + session, + &mut operations, + "arith.addi", + a_band, + a_provider, + &input.location, + )?; + let sixty_four = emit_i32_constant(session, &mut operations, 64, &input.location)?; + let a_row1 = emit_binary( + session, + &mut operations, + "arith.addi", + a_row0.clone(), + sixty_four, + &input.location, + )?; + let b_band = emit_binary( + session, + &mut operations, + "arith.muli", + wn, + sixteen, + &input.location, + )?; + let b_row0 = emit_binary( + session, + &mut operations, + "arith.addi", + b_band, + q, + &input.location, + )?; + let mut rows = vec![a_row0, a_row1, b_row0.clone()]; + for delta in [8_i128, 32, 40, 64, 72, 96, 104] { + let delta = emit_i32_constant(session, &mut operations, delta, &input.location)?; + rows.push(emit_binary( + session, + &mut operations, + "arith.addi", + b_row0.clone(), + delta, + &input.location, + )?); + } + let mut words = Vec::with_capacity(10); + for (index, row) in rows.into_iter().enumerate() { + words.push(load_scale_word( + session, + &mut operations, + if index < 2 { + input.operands[1].clone() + } else { + input.operands[3].clone() + }, + row, + &input.location, + )?); + } + rebuild_carrier( + session, + &mut operations, + &layout, + &words, + input.results[0].clone(), + &input.location, + )?; + Ok(operations) + } +} + +/// Select a K slice and unpack its packed register fragments for the next MMA step. +pub(in crate::cute) struct FragmentSliceKTranslation; + +impl OperationTranslation for FragmentSliceKTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteFragmentSliceKOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.fragment_slice_k", &input, 1, 2)?; + let source_layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_operand(0).get_type(ctx), + session, + 0, + )?; + let result_layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_result(0).get_type(ctx), + session, + 0, + )?; + let mut operations = vec![]; + let mut words = vec![]; + extract_carrier( + session, + &mut operations, + input.operands[0].clone(), + &source_layout, + &mut words, + &input.location, + )?; + if words.len() != 10 || result_layout.leaf_count() != 10 { + return Err("cute.fragment_slice_k needs ten packed scale words".into()); + } + let sixteen = emit_i64_constant(session, &mut operations, 16, &input.location)?; + let shift64 = emit_binary( + session, + &mut operations, + "arith.muli", + input.operands[1].clone(), + sixteen, + &input.location, + )?; + let shift = cast_int( + session, + &mut operations, + shift64, + MlirType::Integer(32), + "arith.trunci", + &input.location, + )?; + let mask = emit_i32_constant(session, &mut operations, 0xffff, &input.location)?; + let mut selected = Vec::with_capacity(10); + for word in words { + let shifted = emit_binary( + session, + &mut operations, + "arith.shrui", + word, + shift.clone(), + &input.location, + )?; + selected.push(emit_binary( + session, + &mut operations, + "arith.andi", + shifted, + mask.clone(), + &input.location, + )?); + } + rebuild_carrier( + session, + &mut operations, + &result_layout, + &selected, + input.results[0].clone(), + &input.location, + )?; + Ok(operations) + } +} + +/// Load the A fragment using the verified shared layout and ldmatrix lane assignment. +pub(in crate::cute) struct MmaLoadATranslation; + +impl OperationTranslation for MmaLoadATranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteMmaLoadAOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + input.attributes.remove("mma_load_a_view"); + expect_shape("cute.mma_load_a", &input, 1, 5)?; + let layout = CarrierLayout::from_source( + ctx, + source.deref(ctx).get_result(0).get_type(ctx), + session, + 0, + )?; + if layout.leaf_count() != 8 { + return Err("cute.mma_load_a needs exactly eight u32 leaves".into()); + } + let mut operations = vec![]; + let first = emit_ldmatrix_x4( + session, + &mut operations, + input.operands[1].clone(), + input.operands[3].clone(), + input.operands[4].clone(), + input.operands[0].clone(), + &input.location, + )?; + let four = emit_i64_constant(session, &mut operations, 4, &input.location)?; + let second_row = emit_binary( + session, + &mut operations, + "arith.addi", + input.operands[3].clone(), + four, + &input.location, + )?; + let second = emit_ldmatrix_x4( + session, + &mut operations, + input.operands[1].clone(), + second_row, + input.operands[4].clone(), + input.operands[0].clone(), + &input.location, + )?; + let mut registers = first.to_vec(); + registers.extend(second); + rebuild_carrier( + session, + &mut operations, + &layout, + ®isters, + input.results[0].clone(), + &input.location, + )?; + Ok(operations) + } +} + +/// Compute swizzled FP16 shared addresses and emit a four-register ldmatrix load. +fn emit_ldmatrix_x4( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base: MlirValueUse, + warp_tile_row: MlirValueUse, + warp_tile_column: MlirValueUse, + lane: MlirValueUse, + location: &MlirLocation, +) -> Result<[MlirValueUse; 4], String> { + let warp_tile_row = cast_int( + session, + operations, + warp_tile_row, + MlirType::Integer(32), + "arith.trunci", + location, + )?; + let warp_tile_column = cast_int( + session, + operations, + warp_tile_column, + MlirType::Integer(32), + "arith.trunci", + location, + )?; + let eight = emit_i32_constant(session, operations, 8, location)?; + let submatrix = emit_binary( + session, + operations, + "arith.divui", + lane.clone(), + eight.clone(), + location, + )?; + let row_in = emit_binary( + session, + operations, + "arith.remui", + lane, + eight.clone(), + location, + )?; + let two = emit_i32_constant(session, operations, 2, location)?; + let parity = emit_binary( + session, + operations, + "arith.remui", + submatrix.clone(), + two.clone(), + location, + )?; + let row_offset = emit_binary( + session, + operations, + "arith.muli", + parity, + eight.clone(), + location, + )?; + let sixteen = emit_i32_constant(session, operations, 16, location)?; + let row = emit_binary( + session, + operations, + "arith.muli", + warp_tile_row, + sixteen.clone(), + location, + )?; + let row = emit_binary(session, operations, "arith.addi", row, row_offset, location)?; + let row = emit_binary(session, operations, "arith.addi", row, row_in, location)?; + + let column = emit_binary( + session, + operations, + "arith.muli", + warp_tile_column, + sixteen, + location, + )?; + let sub_half = emit_binary( + session, + operations, + "arith.divui", + submatrix, + two.clone(), + location, + )?; + let column_offset = emit_binary(session, operations, "arith.muli", sub_half, eight, location)?; + let column = emit_binary( + session, + operations, + "arith.addi", + column, + column_offset, + location, + )?; + + // The verified shared layout is f16 (128,32):(32,1), converted to byte + // strides (64,2), followed by S<2,3,3>. + let sixty_four = emit_i32_constant(session, operations, 64, location)?; + let row_bytes = emit_binary(session, operations, "arith.muli", row, sixty_four, location)?; + let column_bytes = emit_binary(session, operations, "arith.muli", column, two, location)?; + let plain = emit_binary( + session, + operations, + "arith.addi", + row_bytes, + column_bytes, + location, + )?; + // The source layout's element-unit S<2,3,3> becomes S<2,4,3> + // after converting f16 elements to byte offsets. Mask bits 8..7 and + // shift them onto bits 5..4, preserving the 16-byte ldmatrix segment. + let mask = emit_i32_constant( + session, + operations, + LDMATRIX_BYTE_SWIZZLE_MASK as i128, + location, + )?; + let source_bits = emit_binary( + session, + operations, + "arith.andi", + plain.clone(), + mask, + location, + )?; + let three = emit_i32_constant( + session, + operations, + LDMATRIX_BYTE_SWIZZLE_SHIFT as i128, + location, + )?; + let moved = emit_binary( + session, + operations, + "arith.shrui", + source_bits, + three, + location, + )?; + let byte_offset = emit_binary(session, operations, "arith.xori", plain, moved, location)?; + let pointer = gep_value( + session, + operations, + base, + byte_offset, + MlirType::Integer(8), + location, + )?; + + let result_type = MlirType::dialect("!llvm.struct<(i32, i32, i32, i32)>")?; + let (matrix_result, matrix) = fresh(session, result_type); + let mut load = operation( + "nvvm.ldmatrix", + vec![matrix_result], + vec![pointer], + location, + )?; + load.properties.insert( + "layout".into(), + MlirAttribute::dialect("#nvvm.mma_layout")?, + ); + load.properties.insert( + "num".into(), + MlirAttribute::Integer { + value: 4, + ty: MlirType::Integer(32), + }, + ); + operations.push(load); + let mut registers = Vec::with_capacity(4); + for index in 0..4 { + let (result, use_) = fresh(session, MlirType::Integer(32)); + let mut extract = operation( + "llvm.extractvalue", + vec![result], + vec![matrix.clone()], + location, + )?; + extract + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![index])); + operations.push(extract); + registers.push(use_); + } + registers + .try_into() + .map_err(|_| "ldmatrix did not produce four registers".to_owned()) +} + +/// Emit the verified block-scaled MMA sequence and rebuild its accumulator carrier. +pub(in crate::cute) struct TiledGemmTranslation; + +impl OperationTranslation for TiledGemmTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let gemm = CuteTiledGemmOp::wrap(source); + gemm.verify(ctx).map_err(stringify)?; + let plan = gemm + .plan(ctx) + .ok_or_else(|| "cute.tiled_gemm missing plan".to_owned())?; + let is_nvf4 = + plan.atom.scale_format == dialect_cute::attributes::CuteTensorFormatAttr::E4M3; + input.attributes.remove("tiled_gemm_b_view"); + expect_shape("cute.tiled_gemm", &input, 1, 8)?; + let operation_ref = source.deref(ctx); + let a_layout = CarrierLayout::from_source( + ctx, + operation_ref.get_operand(1).get_type(ctx), + session, + 0, + )?; + let scales_layout = CarrierLayout::from_source( + ctx, + operation_ref.get_operand(6).get_type(ctx), + session, + 0, + )?; + let accumulator_layout = CarrierLayout::from_source( + ctx, + operation_ref.get_operand(7).get_type(ctx), + session, + 0, + )?; + let result_layout = + CarrierLayout::from_source(ctx, operation_ref.get_result(0).get_type(ctx), session, 0)?; + drop(operation_ref); + + let mut operations = vec![]; + let mut a = vec![]; + extract_carrier( + session, + &mut operations, + input.operands[1].clone(), + &a_layout, + &mut a, + &input.location, + )?; + let mut scales = vec![]; + extract_carrier( + session, + &mut operations, + input.operands[6].clone(), + &scales_layout, + &mut scales, + &input.location, + )?; + let mut accumulator = vec![]; + extract_carrier( + session, + &mut operations, + input.operands[7].clone(), + &accumulator_layout, + &mut accumulator, + &input.location, + )?; + if a.len() != 8 + || scales.len() != 10 + || accumulator.len() != 64 + || result_layout.leaf_count() != 64 + { + return Err("cute.tiled_gemm carrier shape changed after verification".into()); + } + + for pair in 0..4_usize { + let pair_offset = + emit_i64_constant(session, &mut operations, (pair * 2) as u64, &input.location)?; + let b_row = emit_binary( + session, + &mut operations, + "arith.addi", + input.operands[4].clone(), + pair_offset, + &input.location, + )?; + let loaded = emit_ldmatrix_x4( + session, + &mut operations, + input.operands[2].clone(), + b_row, + input.operands[5].clone(), + input.operands[0].clone(), + &input.location, + )?; + let b_fragments = [ + [loaded[0].clone(), loaded[2].clone()], + [loaded[1].clone(), loaded[3].clone()], + ]; + for m in 0..2_usize { + for (n_in_pair, b_fragment) in b_fragments.iter().enumerate() { + let n = pair * 2 + n_in_pair; + let cell = m * 8 + n; + let updated = emit_mxf4_mma( + session, + &mut operations, + &a[m * 4..m * 4 + 4], + b_fragment, + &accumulator[cell * 4..cell * 4 + 4], + scales[m].clone(), + scales[2 + n].clone(), + is_nvf4, + &input.location, + )?; + accumulator[cell * 4..cell * 4 + 4].clone_from_slice(&updated); + } + } + } + rebuild_carrier( + session, + &mut operations, + &result_layout, + &accumulator, + input.results[0].clone(), + &input.location, + )?; + Ok(operations) + } +} + +/// Emit one MXF4/NVF4 register MMA with A, B, accumulator, and scale operands in NVVM order. +#[allow(clippy::too_many_arguments)] +fn emit_mxf4_mma( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + a: &[MlirValueUse], + b: &[MlirValueUse], + accumulator: &[MlirValueUse], + scale_a: MlirValueUse, + scale_b: MlirValueUse, + is_nvf4: bool, + location: &MlirLocation, +) -> Result<[MlirValueUse; 4], String> { + if a.len() != 4 || b.len() != 2 || accumulator.len() != 4 { + return Err("SM120 MXF4 MMA received the wrong register count".into()); + } + let scale_a = if is_nvf4 { + scale_a + } else { + cast_int( + session, + operations, + scale_a, + MlirType::Integer(16), + "arith.trunci", + location, + )? + }; + let scale_b = if is_nvf4 { + scale_b + } else { + cast_int( + session, + operations, + scale_b, + MlirType::Integer(16), + "arith.trunci", + location, + )? + }; + let mut operands = Vec::with_capacity(12); + operands.extend_from_slice(a); + operands.extend_from_slice(b); + operands.extend_from_slice(accumulator); + operands.extend([scale_a, scale_b]); + + let result_type = MlirType::Float(crate::MlirFloatType::F32); + let mut results = Vec::with_capacity(4); + let mut uses = Vec::with_capacity(4); + for _ in 0..4 { + let (result, use_) = fresh(session, result_type.clone()); + results.push(result); + uses.push(use_); + } + let mut operation_ = operation( + "cute_nvgpu.arch.mma.SM120.block_scaled", + results, + operands, + location, + )?; + operation_.properties.insert( + "a_type".into(), + MlirAttribute::Type(MlirType::Float(crate::MlirFloatType::Other( + "f4E2M1FN".into(), + ))), + ); + operation_.properties.insert( + "b_type".into(), + MlirAttribute::Type(MlirType::Float(crate::MlirFloatType::Other( + "f4E2M1FN".into(), + ))), + ); + operation_.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![4, 2, 4, 1, 1, 0, 0]), + ); + operation_.properties.insert( + "sf_type".into(), + MlirAttribute::Type(MlirType::Float(crate::MlirFloatType::Other( + if is_nvf4 { "f8E4M3FN" } else { "f8E8M0FNU" }.into(), + ))), + ); + operation_.properties.insert( + "shape_MNK".into(), + MlirAttribute::dialect("#cute.shape<\"(16,8,64)\">")?, + ); + operation_.properties.insert( + "thread_id_a".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(16), + }, + ); + operation_.properties.insert( + "thread_id_b".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(16), + }, + ); + operation_.properties.insert( + "vec_size".into(), + MlirAttribute::Integer { + value: if is_nvf4 { 16 } else { 32 }, + ty: MlirType::Integer(32), + }, + ); + operations.push(operation_); + uses.try_into() + .map_err(|_| "SM120 MXF4 MMA did not produce four accumulators".to_owned()) +} + +/// Select a packed scale word using the source lane and scale-column assignment. +fn load_scale_word( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base: MlirValueUse, + row: MlirValueUse, + location: &MlirLocation, +) -> Result { + let mask31 = emit_i32_constant(session, operations, 31, location)?; + let low = emit_binary( + session, + operations, + "arith.andi", + row.clone(), + mask31, + location, + )?; + let four = emit_i32_constant(session, operations, 4, location)?; + let low = emit_binary(session, operations, "arith.muli", low, four, location)?; + let five = emit_i32_constant(session, operations, 5, location)?; + let quadrant = emit_binary(session, operations, "arith.shrui", row, five, location)?; + let three = emit_i32_constant(session, operations, 3, location)?; + let quadrant = emit_binary(session, operations, "arith.andi", quadrant, three, location)?; + let offset = emit_binary(session, operations, "arith.addi", low, quadrant, location)?; + let pointer = gep_value( + session, + operations, + base, + offset, + MlirType::Integer(32), + location, + )?; + let (result, use_) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "llvm.load", + vec![result], + vec![pointer], + location, + )?); + Ok(use_) +} + +/// Apply the verified byte-address XOR permutation while preserving 16-byte segments. +#[cfg(test)] +pub(in crate::cute) const fn ldmatrix_swizzled_byte_offset(plain: u64) -> u64 { + plain ^ ((plain & LDMATRIX_BYTE_SWIZZLE_MASK) >> LDMATRIX_BYTE_SWIZZLE_SHIFT) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mod.rs new file mode 100644 index 0000000000..2ea55fcef2 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mod.rs @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CuTe runtime mappings, organized by capability and architecture. +//! The full-profile registry owns shared registrations exactly once. + +mod common; +mod common_sm100; +pub(in crate::cute) mod copy; +mod emit; +pub(in crate::cute) mod epilogue; +mod independent_sm100; +pub(in crate::cute) mod mma; +pub(in crate::cute) mod pipeline; +mod registry; +pub(in crate::cute) mod scheduler; +mod sm100; +pub(in crate::cute) mod tensor; + +pub(crate) use copy::registry::register_cute_attention_pack; +pub(crate) use emit::{Emit, integer_property, property}; +pub(crate) use independent_sm100::register_cute_sm100_attention_pack; +pub(crate) use registry::register_cute_gemm_pack; +pub(crate) use sm100::register_cute_sm100_pack; +pub use tensor::elementwise::register_cute_elementwise_pack; +pub(crate) use tensor::gemv::register_cute_cutlass_profile_pack; + +#[cfg(test)] +#[path = "tests/gemm.rs"] +mod gemm_tests; + +#[cfg(test)] +#[path = "tests/cooperative.rs"] +mod cooperative_tests; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/mod.rs new file mode 100644 index 0000000000..7554d6630b --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/mod.rs @@ -0,0 +1,9 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Load and accumulator barriers, phase state, and completion ordering. + +pub(in crate::cute) mod sm100; +pub(in crate::cute) mod sm120; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm100.rs new file mode 100644 index 0000000000..135001be4e --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm100.rs @@ -0,0 +1,87 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 cluster load and accumulator pipeline synchronization. + +use crate::cute::{ + emit::{Emit, llvm_ptr_type, property}, + sm100::Kind, +}; +use crate::{MlirAttribute, MlirRegion, MlirType, MlirValueUse}; +use dialect_cute::sm100_ops::*; + +/// Emit cluster load/accumulator barrier actions and their required TCGen05 fences. +pub(in crate::cute) fn emit_pipeline( + e: &mut Emit<'_, '_>, + kind: Kind, + a: &[MlirValueUse], + p: &CuteSm100PipelinePlanAttr, +) -> Result<(), String> { + use Kind::*; + match kind { + PipelineInit | AccInit => { + for stage in 0..p.stages { + let slot = e.constant(32, stage as i128)?; + for (index, count) in if matches!(kind, PipelineInit) { + [(0, 1), (1, p.consumer_arrivals)] + } else { + [(0, p.consumer_arrivals), (1, 1)] + } { + let bar = e.barrier(a[index].clone(), Some(slot.clone()))?; + let count = e.constant(32, count as i128)?; + e.emit("nvvm.mbarrier.init", vec![bar, count])?; + } + } + } + PipelineAcquire | PipelineWait | AccAcquire | AccWait => { + let staged = matches!(kind, PipelineAcquire | PipelineWait); + let bar = e.barrier(a[0].clone(), staged.then(|| a[1].clone()))?; + let phase = a[if staged { 2 } else { 1 }].clone(); + e.wait_cluster(bar, phase)?; + if !matches!(kind, PipelineAcquire) { + e.fence("after")?; + } + } + PipelineExpect => { + let bar = e.barrier(a[0].clone(), Some(a[1].clone()))?; + let bytes = e.constant(32, p.transaction_bytes as i128)?; + let mut op = e.operation("nvvm.mbarrier.arrive.expect_tx", vec![bar, bytes])?; + property(&mut op, "scope", "#nvvm.mem_scope")?; + op.properties + .insert("relaxed".into(), MlirAttribute::Bool(true)); + e.ops.push(op); + } + PipelineRelease | AccCommit => { + let bar = e.barrier( + a[0].clone(), + matches!(kind, PipelineRelease).then(|| a[1].clone()), + )?; + let mask = e.constant(16, p.multicast_mask as i128)?; + let mut op = e.operation("nvvm.tcgen05.commit", vec![bar, mask])?; + property(&mut op, "group", "#nvvm.cta_group")?; + e.ops.push(op); + } + AccRelease => { + // Every lane orders its preceding TMEM reads before one lane of + // the warp releases this accumulator to the MMA producer. + e.fence("before")?; + let elected = e.value("nvvm.elect.sync", vec![], MlirType::Integer(1))?; + let start = e.ops.len(); + let bar = e.llvm_ptr(a[0].clone(), 3)?; + let leader = e.constant(32, p.leader_rank as i128)?; + let remote = e.value("nvvm.mapa", vec![bar, leader], llvm_ptr_type(7)?)?; + let mut arrive = e.operation("nvvm.mbarrier.arrive", vec![remote])?; + property(&mut arrive, "scope", "#nvvm.mem_scope")?; + e.ops.push(arrive); + e.emit("scf.yield", vec![])?; + let body = e.ops.split_off(start); + let mut conditional = e.operation("scf.if", vec![elected])?; + conditional.regions = vec![e.region(body), MlirRegion { blocks: vec![] }]; + e.ops.push(conditional); + } + _ => return Err("invalid SM100 pipeline operation".into()), + } + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm120.rs new file mode 100644 index 0000000000..0ff5d4225d --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/pipeline/sm120.rs @@ -0,0 +1,618 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Load-pipeline initialization, ring state, waits, and barrier operations. +//! +//! Shared storage holds `stages` full barriers, then `stages` empty barriers. +//! Consumers wait on full slots; producers acquire empty slots. Both track +//! `(slot, phase)`: consumers start at `(0, 0)`, producers at `(0, 1)` so the +//! initially empty slots are available. Advancing wraps the slot and flips the +//! phase only on wrap. For two stages, the state calculation is: +//! +//! ``` +//! let advance = |(slot, phase): (u32, u32)| { +//! let next = (slot + 1) % 2; +//! (next, phase ^ u32::from(next == 0)) +//! }; +//! let producer = advance((0, 1)); +//! assert_eq!(producer, (1, 1)); +//! assert_eq!(advance(producer), (0, 0)); +//! ``` + +use crate::cute::common::{ + emit_binary, emit_i32_constant, emit_i64_constant, expect_shape, fresh, gep_value, + integer_constant, operation, shared_pointer_type, stringify, +}; +use crate::{ + MlirAttribute, MlirBlock, MlirLocation, MlirOperation, MlirRegion, MlirResult, MlirType, + MlirValueUse, OperationInput, OperationTranslation, TranslationRegistry, TranslationSession, + TypeTranslation, +}; +use dialect_cute::{ + attributes::CutePipelineRoleAttr, + pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + }, + types::CuteTmaLoadPipelineType, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Bound for each hardware wait probe; the high-level NVVM wait retries to completion. +const PIPELINE_WAIT_TICKS: i128 = 10_000_000; +/// Erase a verified load-pipeline handle to its CTA-shared barrier pointer. +pub(in crate::cute) struct LoadPipelineTypeTranslation; + +impl TypeTranslation for LoadPipelineTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + source + .deref(ctx) + .downcast_ref::() + .ok_or_else(|| "expected cute.tma_load_pipeline".to_owned())? + .verify(ctx) + .map_err(stringify)?; + shared_pointer_type() + } +} + +/// Initialize full/empty barrier rings on the chosen thread, then publish and synchronize. +pub(in crate::cute) struct PipelineInitTranslation; + +impl OperationTranslation for PipelineInitTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteTmaLoadPipelineInitOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.tma_load_pipeline_init", &input, 0, 2)?; + let pipeline = load_pipeline_from_operand(ctx, source, 0)?; + let mut operations = vec![]; + + let (tid_result, tid) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "nvvm.read.ptx.sreg.tid.x", + vec![tid_result], + vec![], + &input.location, + )?); + let (initializer_result, initializer) = fresh(session, MlirType::Integer(1)); + let mut compare = operation( + "arith.cmpi", + vec![initializer_result], + vec![tid, input.operands[1].clone()], + &input.location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 0, // equal + ty: MlirType::Integer(64), + }, + ); + operations.push(compare); + + let mut init_operations = vec![]; + for stage in 0..pipeline.stages { + for (slot, count) in [ + (stage, 1i128), + (pipeline.stages + stage, i128::from(pipeline.consumer_warps)), + ] { + let barrier = barrier_pointer_constant( + session, + &mut operations, + input.operands[0].clone(), + slot, + &input.location, + )?; + let count = emit_i32_constant(session, &mut operations, count, &input.location)?; + init_operations.push(mbarrier_init(barrier, count, &input.location)?); + } + } + init_operations.push(operation("scf.yield", vec![], vec![], &input.location)?); + let mut conditional = operation("scf.if", vec![], vec![initializer], &input.location)?; + conditional.regions = vec![ + MlirRegion { + blocks: vec![MlirBlock { + id: session.fresh_block(), + arguments: vec![], + operations: init_operations, + }], + }, + MlirRegion { blocks: vec![] }, + ]; + operations.push(conditional); + operations.push(operation( + "nvvm.fence.mbarrier.init", + vec![], + vec![], + &input.location, + )?); + operations.push(operation("nvvm.barrier", vec![], vec![], &input.location)?); + Ok(operations) + } +} + +/// Create slot zero with producer or consumer initial parity from the source role. +pub(in crate::cute) struct PipelineStateNewTranslation; + +impl OperationTranslation for PipelineStateNewTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = CutePipelineStateNewOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + expect_shape("cute.pipeline_state_new", &input, 2, 0)?; + let state = semantic + .state(ctx) + .ok_or_else(|| "pipeline_state_new has no state".to_owned())?; + Ok(vec![ + integer_constant(input.results[0].clone(), 0, &input.location)?, + integer_constant( + input.results[1].clone(), + i128::from(state.role == CutePipelineRoleAttr::Producer), + &input.location, + )?, + ]) + } +} + +/// Forward the runtime slot while erasing the source pipeline-state contract. +pub(in crate::cute) struct PipelineStateSlotTranslation; + +impl OperationTranslation for PipelineStateSlotTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + CutePipelineStateSlotOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.pipeline_state_slot", &input, 1, 1)?; + Ok(vec![operation( + "arith.extui", + input.results, + input.operands, + &input.location, + )?]) + } +} + +/// Advance the slot and toggle its phase only when the ring wraps. +pub(in crate::cute) struct PipelineStateAdvanceTranslation; + +impl OperationTranslation for PipelineStateAdvanceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = CutePipelineStateAdvanceOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + expect_shape("cute.pipeline_state_advance", &input, 2, 2)?; + let stages = semantic + .state(ctx) + .ok_or_else(|| "pipeline_state_advance has no state".to_owned())? + .stages; + emit_pipeline_state_advance( + session, + input.operands[0].clone(), + input.operands[1].clone(), + stages, + Some([input.results[0].clone(), input.results[1].clone()]), + &input.location, + ) + .map(|(operations, _)| operations) + } +} + +/// Wait for completion in the selected full or empty barrier ring. +pub(in crate::cute) struct PipelineWaitTranslation { + /// Select the second barrier ring for producer acquisition. + pub(in crate::cute) empty_ring: bool, + /// Source wait operation used to select and validate its semantic contract. + pub(in crate::cute) name: &'static str, +} + +impl OperationTranslation for PipelineWaitTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + match self.name { + "cute.pipeline_producer_acquire" => CutePipelineProducerAcquireOp::wrap(source) + .verify(ctx) + .map_err(stringify)?, + "cute.pipeline_consumer_wait" => CutePipelineConsumerWaitOp::wrap(source) + .verify(ctx) + .map_err(stringify)?, + _ => return Err(format!("unknown pipeline wait mapping {}", self.name)), + } + expect_shape(self.name, &input, 0, 3)?; + let pipeline = load_pipeline_from_operand(ctx, source, 0)?; + let mut operations = vec![]; + let barrier = barrier_pointer_dynamic( + session, + &mut operations, + input.operands[0].clone(), + input.operands[1].clone(), + self.empty_ring.then_some(pipeline.stages), + &input.location, + )?; + let ticks = emit_i32_constant( + session, + &mut operations, + PIPELINE_WAIT_TICKS, + &input.location, + )?; + // This is the LLVM/NVVM high-level potentially-blocking operation. + // Its official definition lowers to a retry loop and guarantees phase + // completion; it is not one raw bounded PTX probe. + operations.push(operation( + "nvvm.mbarrier.try_wait.parity", + vec![], + vec![barrier, input.operands[2].clone(), ticks], + &input.location, + )?); + Ok(operations) + } +} + +/// Announce the verified transaction-byte count to the full-stage barrier. +pub(in crate::cute) struct PipelineExpectTxTranslation; + +impl OperationTranslation for PipelineExpectTxTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CutePipelineProducerExpectTxOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.pipeline_producer_expect_tx", &input, 1, 2)?; + let pipeline = load_pipeline_from_operand(ctx, source, 0)?; + let mut operations = vec![]; + let barrier = barrier_pointer_dynamic( + session, + &mut operations, + input.operands[0].clone(), + input.operands[1].clone(), + None, + &input.location, + )?; + let bytes = emit_i32_constant( + session, + &mut operations, + i128::from(pipeline.transaction_bytes), + &input.location, + )?; + operations.push(operation( + "nvvm.mbarrier.arrive.expect_tx", + vec![], + vec![barrier.clone(), bytes], + &input.location, + )?); + operations.push(operation( + "builtin.unrealized_conversion_cast", + input.results, + vec![barrier], + &input.location, + )?); + Ok(operations) + } +} + +/// Have the selected consumer lane arrive on the corresponding empty-stage barrier. +pub(in crate::cute) struct PipelineReleaseTranslation; + +impl OperationTranslation for PipelineReleaseTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CutePipelineConsumerReleaseOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.pipeline_consumer_release", &input, 0, 2)?; + let pipeline = load_pipeline_from_operand(ctx, source, 0)?; + let mut operations = vec![]; + let barrier = barrier_pointer_dynamic( + session, + &mut operations, + input.operands[0].clone(), + input.operands[1].clone(), + Some(pipeline.stages), + &input.location, + )?; + + let (lane_result, lane) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "nvvm.read.ptx.sreg.laneid", + vec![lane_result], + vec![], + &input.location, + )?); + let zero = emit_i32_constant(session, &mut operations, 0, &input.location)?; + let (elected_result, elected) = fresh(session, MlirType::Integer(1)); + let mut compare = operation( + "arith.cmpi", + vec![elected_result], + vec![lane, zero], + &input.location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(64), + }, + ); + operations.push(compare); + + let mut arrive = operation( + "nvvm.mbarrier.arrive", + vec![], + vec![barrier], + &input.location, + )?; + let yield_ = operation("scf.yield", vec![], vec![], &input.location)?; + arrive.location = input.location.clone(); + let mut conditional = operation("scf.if", vec![], vec![elected], &input.location)?; + conditional.regions = vec![ + MlirRegion { + blocks: vec![MlirBlock { + id: session.fresh_block(), + arguments: vec![], + operations: vec![arrive, yield_], + }], + }, + MlirRegion { blocks: vec![] }, + ]; + operations.push(conditional); + Ok(operations) + } +} + +/// Drain every remaining stage using its matching slot and parity. +pub(in crate::cute) struct PipelineTailTranslation; + +impl OperationTranslation for PipelineTailTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CutePipelineProducerTailOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.pipeline_producer_tail", &input, 0, 3)?; + let pipeline = load_pipeline_from_operand(ctx, source, 0)?; + let mut slot = input.operands[1].clone(); + let mut phase = input.operands[2].clone(); + let mut operations = vec![]; + for stage in 0..pipeline.stages { + let barrier = barrier_pointer_dynamic( + session, + &mut operations, + input.operands[0].clone(), + slot.clone(), + Some(pipeline.stages), + &input.location, + )?; + let ticks = emit_i32_constant( + session, + &mut operations, + PIPELINE_WAIT_TICKS, + &input.location, + )?; + operations.push(operation( + "nvvm.mbarrier.try_wait.parity", + vec![], + vec![barrier, phase.clone(), ticks], + &input.location, + )?); + if stage + 1 != pipeline.stages { + let (advance, next) = emit_pipeline_state_advance( + session, + slot, + phase, + pipeline.stages, + None, + &input.location, + )?; + operations.extend(advance); + [slot, phase] = next; + } + } + Ok(operations) + } +} + +/// Recover the verified pipeline type from a source operand before type erasure. +fn load_pipeline_from_operand( + ctx: &Context, + source: Ptr, + index: usize, +) -> Result { + let ty = source.deref(ctx).get_operand(index).get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected a cute.tma_load_pipeline operand".to_owned()) +} + +/// Form a shared barrier pointer from a compile-time slot counted in i64 elements. +fn barrier_pointer_constant( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base: MlirValueUse, + slot: u64, + location: &MlirLocation, +) -> Result { + let slot = emit_i64_constant(session, operations, slot, location)?; + gep_value( + session, + operations, + base, + slot, + MlirType::Integer(64), + location, + ) +} + +/// Form a shared barrier pointer and optionally offset into the empty ring. +fn barrier_pointer_dynamic( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base: MlirValueUse, + slot: MlirValueUse, + add_stages: Option, + location: &MlirLocation, +) -> Result { + let (wide_result, mut index) = fresh(session, MlirType::Integer(64)); + operations.push(operation( + "arith.extui", + vec![wide_result], + vec![slot], + location, + )?); + if let Some(stages) = add_stages { + let stages = emit_i64_constant(session, operations, stages, location)?; + index = emit_binary(session, operations, "arith.addi", index, stages, location)?; + } + gep_value( + session, + operations, + base, + index, + MlirType::Integer(64), + location, + ) +} + +/// Emit slot increment, wrap-to-zero, and parity toggle, optionally reusing caller result IDs. +fn emit_pipeline_state_advance( + session: &mut TranslationSession<'_>, + slot: MlirValueUse, + phase: MlirValueUse, + stages: u64, + final_results: Option<[MlirResult; 2]>, + location: &MlirLocation, +) -> Result<(Vec, [MlirValueUse; 2]), String> { + let mut operations = vec![]; + let one = emit_i32_constant(session, &mut operations, 1, location)?; + let incremented = emit_binary(session, &mut operations, "arith.addi", slot, one, location)?; + let width = emit_i32_constant(session, &mut operations, i128::from(stages), location)?; + let (wrap_result, wraps) = fresh(session, MlirType::Integer(1)); + let mut compare = operation( + "arith.cmpi", + vec![wrap_result], + vec![incremented.clone(), width], + location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(64), + }, + ); + operations.push(compare); + let zero = emit_i32_constant(session, &mut operations, 0, location)?; + + let (slot_result, next_slot) = if let Some(results) = &final_results { + ( + results[0].clone(), + MlirValueUse { + id: results[0].id, + ty: results[0].ty.clone(), + }, + ) + } else { + fresh(session, MlirType::Integer(32)) + }; + operations.push(operation( + "arith.select", + vec![slot_result], + vec![wraps.clone(), zero, incremented], + location, + )?); + + let (wrap_u32_result, wrap_u32) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + "arith.extui", + vec![wrap_u32_result], + vec![wraps], + location, + )?); + let (phase_result, next_phase) = if let Some(results) = final_results { + ( + results[1].clone(), + MlirValueUse { + id: results[1].id, + ty: results[1].ty.clone(), + }, + ) + } else { + fresh(session, MlirType::Integer(32)) + }; + operations.push(operation( + "arith.xori", + vec![phase_result], + vec![phase, wrap_u32], + location, + )?); + Ok((operations, [next_slot, next_phase])) +} + +/// Build an unpredicated barrier initialization with the given arrival count. +pub(in crate::cute) fn mbarrier_init( + barrier: MlirValueUse, + arrival_count: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation( + "nvvm.mbarrier.init", + vec![], + vec![barrier, arrival_count], + location, + ) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/registry.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/registry.rs new file mode 100644 index 0000000000..27a5f457e3 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/registry.rs @@ -0,0 +1,174 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared GEMM foundation registration, including registrations used by SM100. +//! This single owner preserves registration order across capability modules. + +use crate::cute::{ + common::{ + ExternSharedTranslation, IdentityFirstOperandTranslation, IdentityTranslation, + InsertFieldTranslation, StaticContractTypeTranslation, StaticMultiIdentityTranslation, + }, + copy::tma::{TmaLoadTranslation, TmaStoreTranslation, TmaViewTypeTranslation}, + epilogue::sm120::{ + EpilogueHalfTranslation, EpilogueStoreTranslation, EpilogueSyncTranslation, + NoOperandRenameTranslation, TmaStoreWaitTranslation, + }, + mma::sm120::{ + FragmentFillTranslation, FragmentSliceKTranslation, MmaLoadATranslation, + MmaLoadScalesTranslation, TiledGemmTranslation, + }, + pipeline::sm120::{ + LoadPipelineTypeTranslation, PipelineExpectTxTranslation, PipelineInitTranslation, + PipelineReleaseTranslation, PipelineStateAdvanceTranslation, PipelineStateNewTranslation, + PipelineStateSlotTranslation, PipelineTailTranslation, PipelineWaitTranslation, + }, + scheduler::sm120::{ + SchedulerAdvanceTranslation, SchedulerHasWorkTranslation, SchedulerNewTranslation, + WorkTileCoordinatesTranslation, WorkTileTypeTranslation, + }, +}; +use crate::{DropAttribute, TranslationError, TranslationRegistry}; +use dialect_cute::{ + attributes::{ + CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpilogueSyncPhaseAttr, + CuteMmaCarrierKindAttr, CutePipelineStateAttr, CuteTileGridAttr, CuteTiledMmaPlanAttr, + CuteTmaStorePipelineAttr, + }, + epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, + }, + scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, + }, + types::{ + CuteEpilogueTileType, CuteSmemTensorType, CuteTmaLoadPipelineType, CuteTmaViewType, + CuteWorkTileType, + }, +}; +use dialect_mir::ops::{MirExternSharedOp, MirInsertFieldOp}; + +/// Register the runtime foundation shared by TMA, the mainloop, and epilogue. +pub(crate) fn register_cute_gemm_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_type::(WorkTileTypeTranslation)?; + registry.register_type::(LoadPipelineTypeTranslation)?; + registry.register_type::(TmaViewTypeTranslation)?; + // These types occur only inside source TypeAttr values. Their recipes read + // the verified source objects directly, so the translated attribute value + // is deliberately an empty sentinel and is removed before emission. + registry.register_type::(StaticContractTypeTranslation)?; + registry.register_type::(StaticContractTypeTranslation)?; + + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + + registry.register_operation::(InsertFieldTranslation)?; + registry.register_operation::(ExternSharedTranslation)?; + + registry.register_operation::(SchedulerNewTranslation)?; + registry.register_operation::(SchedulerHasWorkTranslation)?; + registry.register_operation::(IdentityTranslation( + "cute.scheduler_current", + ))?; + registry.register_operation::(WorkTileCoordinatesTranslation)?; + registry.register_operation::(SchedulerAdvanceTranslation)?; + + registry.register_operation::(IdentityFirstOperandTranslation( + "cute.tma_load_pipeline_make", + ))?; + registry.register_operation::(PipelineInitTranslation)?; + registry.register_operation::(PipelineStateNewTranslation)?; + registry.register_operation::(PipelineStateSlotTranslation)?; + registry.register_operation::(PipelineStateAdvanceTranslation)?; + registry.register_operation::(PipelineWaitTranslation { + empty_ring: true, + name: "cute.pipeline_producer_acquire", + })?; + registry.register_operation::(PipelineWaitTranslation { + empty_ring: false, + name: "cute.pipeline_consumer_wait", + })?; + registry.register_operation::(PipelineExpectTxTranslation)?; + registry.register_operation::(PipelineReleaseTranslation)?; + registry.register_operation::(PipelineTailTranslation)?; + + registry.register_operation::(IdentityFirstOperandTranslation( + "cute.tma_gmem_view", + ))?; + registry.register_operation::(IdentityFirstOperandTranslation( + "cute.tma_smem_view", + ))?; + registry.register_operation::(TmaLoadTranslation)?; + registry.register_operation::(TmaStoreTranslation)?; + + registry.register_operation::(TmaStoreWaitTranslation { + name: "cute.tma_store_acquire", + })?; + registry.register_operation::(NoOperandRenameTranslation( + "nvvm.cp.async.bulk.commit.group", + ))?; + registry.register_operation::(TmaStoreWaitTranslation { + name: "cute.tma_store_tail", + })?; + // Pointer/coordinate-only epilogue operations are included here. The + // accumulator-to-stmatrix leaf is registered by the MMA rung below. + registry.register_operation::(StaticMultiIdentityTranslation { + name: "cute.epilogue_smem_overlay", + attributes: &["epilogue_overlay_tile"], + })?; + registry.register_operation::(StaticMultiIdentityTranslation { + name: "cute.epilogue_warp_slice", + attributes: &["epilogue_warp_slice_tile"], + })?; + registry.register_operation::(EpilogueHalfTranslation)?; + registry.register_operation::(EpilogueSyncTranslation)?; + + register_mma_and_epilogue_mappings(registry) +} + +/// Register SM120 MMA carrier operations and accumulator-to-shared epilogue stores. +fn register_mma_and_epilogue_mappings( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_operation::(StaticMultiIdentityTranslation { + name: "cute.smem_tensor_overlay", + attributes: &["smem_overlay_view"], + })?; + registry + .register_operation::(IdentityTranslation("cute.tiled_mma_slice"))?; + registry.register_operation::(FragmentFillTranslation)?; + registry.register_operation::(MmaLoadScalesTranslation)?; + registry.register_operation::(FragmentSliceKTranslation)?; + registry.register_operation::(MmaLoadATranslation)?; + registry.register_operation::(StaticMultiIdentityTranslation { + name: "cute.mma_partition_b", + attributes: &["mma_partition_b_view"], + })?; + registry.register_operation::(TiledGemmTranslation)?; + registry.register_operation::(EpilogueStoreTranslation)?; + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/mod.rs new file mode 100644 index 0000000000..27bcbd9765 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/mod.rs @@ -0,0 +1,8 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent work assignment and runtime tile coordinate decoding. + +pub(in crate::cute) mod sm120; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/sm120.rs new file mode 100644 index 0000000000..2b1a08fa45 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/scheduler/sm120.rs @@ -0,0 +1,271 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent one-dimensional tile scheduler and runtime work coordinates. + +use crate::cute::common::{ + binary, constant_value, emit_binary, emit_i64_constant, expect_shape, fresh, operation, + stringify, u64_bits, +}; +use crate::{ + MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationRegistry, TranslationSession, TypeTranslation, +}; +use dialect_cute::{ + attributes::CuteTileGridAttr, + scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerHasWorkOp, CuteSchedulerNew1dOp, + CuteWorkTileCoordinatesOp, + }, + types::CuteWorkTileType, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Represent a verified scheduler work tile as its runtime integer cursor. +pub(in crate::cute) struct WorkTileTypeTranslation; + +impl TypeTranslation for WorkTileTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + source + .deref(ctx) + .downcast_ref::() + .ok_or_else(|| "expected cute.work_tile".to_owned())? + .verify(ctx) + .map_err(stringify)?; + Ok(MlirType::Integer(64)) + } +} + +/// Initialize the persistent scheduler cursor from the source launch coordinate. +pub(in crate::cute) struct SchedulerNewTranslation; + +impl OperationTranslation for SchedulerNewTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteSchedulerNew1dOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.scheduler_new_1d", &input, 2, 0)?; + let mut operations = vec![]; + for (name, result) in [ + ("nvvm.read.ptx.sreg.ctaid.x", input.results[0].clone()), + ("nvvm.read.ptx.sreg.nctaid.x", input.results[1].clone()), + ] { + let (narrow_result, narrow) = fresh(session, MlirType::Integer(32)); + operations.push(operation( + name, + vec![narrow_result], + vec![], + &input.location, + )?); + operations.push(operation( + "arith.extui", + vec![result], + vec![narrow], + &input.location, + )?); + } + Ok(operations) + } +} + +/// Compare the cursor with the verified total tile count. +pub(in crate::cute) struct SchedulerHasWorkTranslation; + +impl OperationTranslation for SchedulerHasWorkTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let semantic = CuteSchedulerHasWorkOp::wrap(source); + semantic.verify(ctx).map_err(stringify)?; + expect_shape("cute.scheduler_has_work", &input, 1, 1)?; + let total = semantic + .tile_grid(ctx) + .and_then(CuteTileGridAttr::total_tiles) + .ok_or_else(|| "scheduler tile count overflowed".to_owned())?; + let (constant, total_use) = constant_value( + session, + MlirType::Integer(64), + u64_bits(total), + &input.location, + )?; + let mut compare = operation( + "arith.cmpi", + input.results, + vec![input.operands[0].clone(), total_use], + &input.location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 6, // unsigned less-than + ty: MlirType::Integer(64), + }, + ); + Ok(vec![constant, compare]) + } +} + +/// Decode the cursor into M/N tile coordinates using the verified traversal order. +pub(in crate::cute) struct WorkTileCoordinatesTranslation; + +impl OperationTranslation for WorkTileCoordinatesTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteWorkTileCoordinatesOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.work_tile_coordinates", &input, 4, 1)?; + let tile_ty = source.deref(ctx).get_operand(0).get_type(ctx); + let tile_ref = tile_ty.deref(ctx); + let grid = tile_ref + .downcast_ref::() + .ok_or_else(|| "work_tile_coordinates operand lost its grid".to_owned())? + .grid; + let linear = input.operands[0].clone(); + let mut operations = vec![]; + let zero = emit_i64_constant(session, &mut operations, 0, &input.location)?; + operations.push(binary( + "arith.addi", + input.results[0].clone(), + linear.clone(), + zero, + &input.location, + )?); + + let m_width = emit_i64_constant(session, &mut operations, grid.m_tiles, &input.location)?; + let quotient_m = emit_binary( + session, + &mut operations, + "arith.divui", + linear.clone(), + m_width.clone(), + &input.location, + )?; + let consumed_m = emit_binary( + session, + &mut operations, + "arith.muli", + quotient_m.clone(), + m_width, + &input.location, + )?; + operations.push(binary( + "arith.subi", + input.results[1].clone(), + linear, + consumed_m, + &input.location, + )?); + + let n_width = emit_i64_constant(session, &mut operations, grid.n_tiles, &input.location)?; + let batch = emit_binary( + session, + &mut operations, + "arith.divui", + quotient_m.clone(), + n_width.clone(), + &input.location, + )?; + let consumed_n = emit_binary( + session, + &mut operations, + "arith.muli", + batch.clone(), + n_width, + &input.location, + )?; + operations.push(binary( + "arith.subi", + input.results[2].clone(), + quotient_m, + consumed_n, + &input.location, + )?); + let zero = emit_i64_constant(session, &mut operations, 0, &input.location)?; + operations.push(binary( + "arith.addi", + input.results[3].clone(), + batch, + zero, + &input.location, + )?); + Ok(operations) + } +} + +/// Advance by the persistent grid stride, saturating instead of wrapping on overflow. +pub(in crate::cute) struct SchedulerAdvanceTranslation; + +impl OperationTranslation for SchedulerAdvanceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CuteSchedulerAdvanceOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + expect_shape("cute.scheduler_advance", &input, 1, 2)?; + let mut operations = vec![]; + let wrapped = emit_binary( + session, + &mut operations, + "arith.addi", + input.operands[0].clone(), + input.operands[1].clone(), + &input.location, + )?; + let (overflow_result, overflow) = fresh(session, MlirType::Integer(1)); + let mut compare = operation( + "arith.cmpi", + vec![overflow_result], + vec![wrapped.clone(), input.operands[0].clone()], + &input.location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 6, // unsigned less-than + ty: MlirType::Integer(64), + }, + ); + operations.push(compare); + let maximum = emit_i64_constant(session, &mut operations, u64::MAX, &input.location)?; + operations.push(operation( + "arith.select", + input.results, + vec![overflow, maximum, wrapped], + &input.location, + )?); + Ok(operations) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/sm100.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/sm100.rs new file mode 100644 index 0000000000..3a48a87d36 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/sm100.rs @@ -0,0 +1,527 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Map SM100 CuTe plans to operations accepted by the pinned CUTLASS 4.7 compiler. +//! +//! Allocation, tensor-memory copies, cluster TMA, and shared-memory descriptors +//! use native CuTe operations. CUTLASS 4.7's `sm100.mma` atom cannot select when +//! to fill or release the A collector, so MMA uses an NVVM TCGen05 instruction. +//! CuTe layouts supply its descriptors; the verified tiled-MMA plan supplies +//! its shape and collector sequence: +//! +//! ```text +//! each K16 step: first N partition -> fill A collector +//! second partition -> last use of A collector +//! ``` + +#[cfg(test)] +use crate::cute::mma::sm100::instruction_descriptor; +use crate::cute::{ + emit::{Emit, dialect, integer_property, llvm_ptr_type, property, segments}, + epilogue::sm100::emit_epilogue, + mma::sm100::emit_mma, + pipeline::sm100::emit_pipeline, +}; +use crate::{ + DropAttribute, MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationError, TranslationRegistry, TranslationSession, +}; +#[cfg(test)] +use dialect_cute::attributes::CuteComposedLayoutAttr; +use dialect_cute::{attributes::CuteTmaStorePipelineAttr, sm100_ops::*}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Register paired SM100 GEMM plans and their tensor-memory, copy, and pipeline recipes. +pub(crate) fn register_cute_sm100_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + /// Bind a typed SM100 source operation to its verified plan reader and runtime action. + macro_rules! register { + ($op:ty, $kind:ident, $plan:ident) => { + registry.register_operation::<$op>(Recipe { + kind: Kind::$kind, + read: |ctx, source| { + let op = <$op>::wrap(source); + op.verify(ctx).map_err(|e| e.to_string())?; + op.plan(ctx) + .map(Plan::$plan) + .ok_or_else(|| "SM100 operation lost its verified plan".into()) + }, + })?; + }; + } + register!(CuteSm100TmemAllocOp, Alloc, Tmem); + register!(CuteSm100TmemDeallocOp, Dealloc, Tmem); + register!(CuteSm100TiledMmaOp, Mma, Mma); + register!(CuteSm100TmemEpilogueOp, Epilogue, Epilogue); + register!(CuteSm100ClusterTmaLoadOp, TmaLoad, Copy); + register!(CuteSm100PipelineInitOp, PipelineInit, Pipeline); + register!(CuteSm100PipelineAcquireOp, PipelineAcquire, Pipeline); + register!(CuteSm100PipelineExpectOp, PipelineExpect, Pipeline); + register!(CuteSm100PipelineWaitOp, PipelineWait, Pipeline); + register!(CuteSm100PipelineReleaseOp, PipelineRelease, Pipeline); + register!(CuteSm100AccumulatorInitOp, AccInit, Pipeline); + register!(CuteSm100AccumulatorAcquireOp, AccAcquire, Pipeline); + register!(CuteSm100AccumulatorCommitOp, AccCommit, Pipeline); + register!(CuteSm100AccumulatorWaitOp, AccWait, Pipeline); + register!(CuteSm100AccumulatorReleaseOp, AccRelease, Pipeline); + register!(CuteSm100TmaStoreOp, TmaStore, Epilogue); + register!(CuteSm100StoreCommitOp, StoreCommit, Store); + register!(CuteSm100StoreAcquireOp, StoreAcquire, Store); + register!(CuteSm100StoreTailOp, StoreTail, Store); + Ok(()) +} + +/// Select the runtime action after the source plan has been verified. +#[derive(Clone, Copy)] +pub(in crate::cute) enum Kind { + /// Allocate the plan's tensor-memory columns. + Alloc, + /// Release the previously allocated tensor-memory columns. + Dealloc, + /// Issue paired FP16 MMA with the plan's A-collector sequence. + Mma, + /// Move the accumulator through bias/conversion into shared output storage. + Epilogue, + /// Issue the plan's cluster TMA load and completion routing. + TmaLoad, + /// Initialize the staged load pipeline's full and empty barriers. + PipelineInit, + /// Wait for an empty load stage before issuing its copy. + PipelineAcquire, + /// Announce transaction bytes on the load stage's full barrier. + PipelineExpect, + /// Wait until the load stage's data is available to its consumers. + PipelineWait, + /// Commit consumption so the load stage can be reused. + PipelineRelease, + /// Initialize the accumulator pipeline's readiness and reuse barriers. + AccInit, + /// Wait until the accumulator storage is available to the MMA producer. + AccAcquire, + /// Commit MMA completion to the accumulator's ready barrier. + AccCommit, + /// Wait for the accumulator before reading tensor memory. + AccWait, + /// Order every lane's reads before one elected arrival permits accumulator reuse. + AccRelease, + /// Issue the plan's shared-to-global TMA output copy. + TmaStore, + /// Commit issued TMA stores as one bulk group. + StoreCommit, + /// Wait until the store queue has capacity for the next group. + StoreAcquire, + /// Drain all outstanding TMA store groups. + StoreTail, +} +/// Verified source plan retained until its runtime action has been emitted. +enum Plan { + /// Tensor-memory allocation extent and ownership contract. + Tmem(#[doc = "Allocation extent and tensor-memory ownership plan."] CuteSm100TmemPlanAttr), + /// Verified paired-MMA shape, shared layouts, and collector-lifetime plan. + Mma(#[doc = "Paired-MMA shape, layouts, and collector-lifetime plan."] CuteSm100MmaPlanAttr), + /// Verified accumulator shape, conversion, and shared-output placement plan. + Epilogue( + #[doc = "Accumulator conversion and shared-output placement plan."] + CuteSm100EpiloguePlanAttr, + ), + /// Cluster-copy layout, destination, and completion-routing contract. + Copy( + #[doc = "Cluster-copy shape, layout, and completion-routing plan."] + CuteSm100ClusterCopyPlanAttr, + ), + /// Load or accumulator barrier counts, stages, and completion routing. + Pipeline( + #[doc = "Barrier stages, arrivals, byte counts, and cluster routing."] + CuteSm100PipelinePlanAttr, + ), + /// Maximum permitted number of pending TMA store groups. + Store(#[doc = "Bound on outstanding hardware TMA store groups."] CuteTmaStorePipelineAttr), +} +/// Pair a runtime action with the source-plan reader that validates its operands. +struct Recipe { + /// Runtime action implemented by this recipe. + kind: Kind, + /// Read and verify the operation's typed plan before target emission. + read: fn(&Context, Ptr) -> Result, +} + +impl OperationTranslation for Recipe { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let plan = (self.read)(ctx, source)?; + if !input.results.is_empty() + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err( + "SM100 semantic operation retained unexpected results, attributes or regions" + .into(), + ); + } + let mut e = Emit { + session, + location: input.location, + ops: vec![], + }; + let a = input.operands; + use Kind::*; + match (self.kind, plan) { + (Alloc, Plan::Tmem(p)) => { + let columns = e.constant(32, p.columns as i128)?; + let token = e.cute_ptr(a[0].clone(), "i32", "smem", 4, None)?; + let mut alloc = + e.operation("cute_nvgpu.arch.sm100.alloc_tmem", vec![columns, token])?; + alloc + .properties + .insert("is_two_cta".into(), MlirAttribute::Unit); + e.ops.push(alloc); + let mut relinquish = + e.operation("cute_nvgpu.arch.sm100.relinquish_tmem_alloc_permit", vec![])?; + relinquish + .properties + .insert("is_two_cta".into(), MlirAttribute::Unit); + e.ops.push(relinquish); + } + (Dealloc, Plan::Tmem(p)) => { + let ptr = e.cast( + "builtin.unrealized_conversion_cast", + a[0].clone(), + dialect("!cute.ptr>")?, + )?; + let columns = e.constant(32, p.columns as i128)?; + let mut op = + e.operation("cute_nvgpu.arch.sm100.dealloc_tmem", vec![ptr, columns])?; + op.properties + .insert("is_two_cta".into(), MlirAttribute::Unit); + e.ops.push(op); + } + (Mma, Plan::Mma(p)) => emit_mma(&mut e, &a, &p)?, + (Epilogue, Plan::Epilogue(p)) => emit_epilogue(&mut e, &a, &p)?, + (TmaLoad, Plan::Copy(p)) => { + let desc = e.cute_ptr( + a[0].clone(), + "!cute_nvgpu.tma_descriptor_tiled", + "generic", + 64, + None, + )?; + let dst = e.cute_ptr(a[1].clone(), "f16", "smem", 128, None)?; + let rank = e.constant(32, p.leader_rank as i128)?; + let local_bar = e.llvm_ptr(a[4].clone(), 3)?; + let remote = e.value("nvvm.mapa", vec![local_bar, rank], llvm_ptr_type(7)?)?; + // The two-CTA TMA intrinsic encodes the leader's barrier as a + // shared address. Retype its bits, retaining the mapa rank. + let bits = e.cast("llvm.ptrtoint", remote, MlirType::Integer(32))?; + let bar_ptr = e.cast("llvm.inttoptr", bits, llvm_ptr_type(3)?)?; + let barrier = e.cute_ptr(bar_ptr, "i64", "smem", 8, None)?; + let one = e.constant(32, 1)?; + let mask32 = e.binary("arith.shli", one, a[5].clone())?; + let mask = e.cast("arith.trunci", mask32, MlirType::Integer(16))?; + let mut op = e.operation( + "cute_nvgpu.arch.copy.SM100.tma_load", + vec![desc, dst, barrier, a[3].clone(), a[2].clone(), mask], + )?; + property(&mut op, "mode", "#cute_nvgpu.tma_load_mode")?; + integer_property(&mut op, "num_cta", 32, p.cta_group as i128); + segments(&mut op, &[1, 1, 1, 2, 1, 0, 0]); + e.ops.push(op); + } + (TmaStore, Plan::Epilogue(_)) => { + let desc = e.cute_ptr( + a[0].clone(), + "!cute_nvgpu.tma_descriptor_tiled", + "generic", + 64, + None, + )?; + let src = e.cute_ptr(a[1].clone(), "f16", "smem", 128, None)?; + let mut op = e.operation( + "cute_nvgpu.arch.copy.SM100.tma_store", + vec![desc, src, a[3].clone(), a[2].clone()], + )?; + property(&mut op, "mode", "#cute_nvgpu.tma_store_mode")?; + segments(&mut op, &[1, 1, 2, 0]); + e.ops.push(op); + } + (StoreCommit, Plan::Store(_)) => e.emit("nvvm.cp.async.bulk.commit.group", vec![])?, + (StoreAcquire | StoreTail, Plan::Store(p)) => { + let groups = if matches!(self.kind, StoreTail) { + 0 + } else { + p.stages - 1 + }; + let mut op = e.operation("nvvm.cp.async.bulk.wait_group", vec![])?; + integer_property(&mut op, "group", 32, groups as i128); + op.properties.insert("read".into(), MlirAttribute::Unit); + e.ops.push(op); + } + (kind, Plan::Pipeline(p)) => emit_pipeline(&mut e, kind, &a, &p)?, + _ => return Err("SM100 operation/plan mismatch".into()), + } + Ok(e.ops) + } +} + +#[cfg(test)] +mod tests { + //! Paired-GEMM mapping tests for collector lifetime, vector epilogues, and accumulator release. + + use super::*; + use crate::{CutlassFullCuteMlir22, MlirConsumerProfile}; + use dialect_cute::layout::{ComposedLayout, IntTuple, OffsetUnit, Swizzle}; + use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType, address_space}, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + op::Op, + r#type::TypeHandle, + value::Value, + }; + + /// Build the canonical row-major FP16 composed layout with a selected swizzle width. + fn layout(rows: u32, cols: u32, bits: u32) -> CuteComposedLayoutAttr { + CuteComposedLayoutAttr( + ComposedLayout::new( + Swizzle::new(bits, 3, 3), + 0, + format!("({rows},{cols}):({cols},1)").parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ) + } + /// Build the paired-MMA fixture plan for two independently selected N partitions. + fn mma_plan(n0: u32, n1: u32) -> CuteSm100MmaPlanAttr { + CuteSm100MmaPlanAttr { + m: 256, + n0, + n1, + k: 64, + tmem_columns: 512, + cta_group: 2, + collector: CuteSm100CollectorAttr::AFillLastUse, + a_layout: layout(128, 64, 3), + b0_layout: layout(n0 / 2, 64, 3), + b1_layout: layout(n1 / 2, 64, 3), + } + } + /// Create a typed SM100 function with the runtime operands used by individual recipe tests. + fn fixture() -> (Context, ModuleOp, Ptr, Vec) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let u32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let f32: TypeHandle = FP32Type::get(&ctx).into(); + let f16: TypeHandle = MirFP16Type::get(&ctx).into(); + let halfptr: TypeHandle = + MirPtrType::get(&mut ctx, f16, true, address_space::SHARED).into(); + let barptr: TypeHandle = MirPtrType::get(&mut ctx, u64, true, address_space::SHARED).into(); + let tokenptr: TypeHandle = + MirPtrType::get(&mut ctx, u32, true, address_space::SHARED).into(); + let args = vec![ + u32, halfptr, halfptr, halfptr, i1, barptr, u32, u32, f32, tokenptr, + ]; + let module = ModuleOp::new(&mut ctx, "sm100".try_into().unwrap()); + let ft = FunctionType::get(&ctx, args.clone(), vec![]); + let func = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let f = MirFuncOp::new(&mut ctx, func, TypeAttr::new(ft.into())); + f.set_symbol_name(&mut ctx, "sm100_probe".try_into().unwrap()); + module.append_operation(&mut ctx, func, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(func.deref(&ctx).get_region(0), &ctx); + let values = block.deref(&ctx).arguments().collect(); + (ctx, module, block, values) + } + /// Terminate and render the fixture while intentionally bypassing whole-kernel lifecycle checks. + fn render(ctx: &mut Context, module: &ModuleOp, block: Ptr) -> String { + let ret = Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, ctx); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let registry = profile.build_registry().unwrap(); + // Mapping-unit tests deliberately bypass whole-kernel lifecycle checks; + // each source operation is still verified by its own translation. + let target = + crate::translate_module(ctx, module, ®istry, &profile.translation_config()).unwrap(); + crate::render_module(&target).unwrap() + } + /// Check each paired-MMA partition uses plan-derived descriptors and the intended collector transition. + #[test] + fn paired_mma_derives_native_layout_descriptors_and_collector_lifetimes() { + for (n0, n1) in [(192, 160), (128, 224)] { + let (mut ctx, module, block, a) = fixture(); + let op = CuteSm100TiledMmaOp::new(&mut ctx, a[..5].to_vec(), mma_plan(n0, n1)); + op.get_operation().insert_at_back(block, &ctx); + let text = render(&mut ctx, &module, block); + assert_eq!( + text.matches("\"cute_nvgpu.make_umma_smem_desc\"").count(), + 3, + "{text}" + ); + assert!(text.contains("S<3,4,3>"), "{text}"); + assert!( + text.contains(&format!("(({},16),1):((64,1),0)", n0 / 2)), + "{text}" + ); + assert!( + text.contains(&format!("(({},16),1):((64,1),0)", n1 / 2)), + "{text}" + ); + assert_eq!(text.matches("\"nvvm.tcgen05.mma\"").count(), 8, "{text}"); + assert_eq!( + text.matches("#nvvm.tcgen05_mma_collectorop").count(), + 4, + "{text}" + ); + assert_eq!( + text.matches("#nvvm.tcgen05_mma_collectorop") + .count(), + 4, + "{text}" + ); + assert!( + text.contains(&format!("{} : i32", instruction_descriptor(256, n0))), + "{text}" + ); + assert!(!text.contains("llvm.inline_asm"), "{text}"); + } + } + /// Check TMEM results remain vector-valued through bias, FP16 conversion, and shared chunk stores. + #[test] + fn tensor_memory_epilogue_preserves_vector_bias_conversion_and_chunk_stores() { + let (mut ctx, module, block, a) = fixture(); + let plan = CuteSm100EpiloguePlanAttr { + layout: layout(128, 32, 2), + rows: 128, + columns: 32, + }; + let op = CuteSm100TmemEpilogueOp::new(&mut ctx, vec![a[0], a[1], a[6], a[8]], plan); + op.get_operation().insert_at_back(block, &ctx); + let text = render(&mut ctx, &module, block); + assert_eq!( + text.matches("\"cute_nvgpu.arch.copy.SM100.tmem_load\"") + .count(), + 1, + "{text}" + ); + assert_eq!(text.matches("\"vector.broadcast\"").count(), 1, "{text}"); + assert_eq!(text.matches("\"arith.addf\"").count(), 1, "{text}"); + assert_eq!(text.matches("\"arith.truncf\"").count(), 1, "{text}"); + assert_eq!(text.matches("\"llvm.store\"").count(), 4, "{text}"); + assert_eq!(text.matches("alignment = 2 : i64").count(), 4, "{text}"); + assert_eq!(text.matches("\"vector.shuffle\"").count(), 4, "{text}"); + assert!(!text.contains("\"vector.extract\""), "{text}"); + assert!(text.contains("vector<32xf32>"), "{text}"); + assert!(text.contains("vector<32xf16>"), "{text}"); + assert!(text.contains("vector<8xf16>"), "{text}"); + assert!(text.contains("vector<32xi32>"), "{text}"); + assert!(!text.contains("\"nvvm.tcgen05.ld\""), "{text}"); + } + /// Enumerate every chunk to check element coverage, swizzle mapping, and relative store alignment. + #[test] + fn epilogue_vector_chunks_preserve_every_swizzled_element_and_relative_alignment() { + let ctx = Context::new(); + let plan = CuteSm100EpiloguePlanAttr { + layout: layout(128, 32, 2), + rows: 128, + columns: 32, + }; + plan.verify(&ctx).unwrap(); + let coordinate = |row, col| IntTuple::Tuple(vec![IntTuple::Leaf(row), IntTuple::Leaf(col)]); + let mut covered = std::collections::BTreeSet::new(); + for row in 0..i64::from(plan.rows) { + for col in (0..i64::from(plan.columns)).step_by(8) { + let base = plan.layout.0.checked_call(&coordinate(row, col)).unwrap(); + // A 16-byte-aligned destination retains that alignment for + // every chunk. This proves relative alignment only: the raw + // pointer API still permits merely FP16-aligned destinations. + assert_eq!((base * 2) % 16, 0, "row {row}, column {col}"); + for lane in 0..8 { + let physical = plan + .layout + .0 + .checked_call(&coordinate(row, col + lane)) + .unwrap(); + assert_eq!(physical, base + lane, "row {row}, column {}", col + lane); + assert!(covered.insert(physical), "duplicate offset {physical}"); + } + } + } + assert!( + covered + .into_iter() + .eq(0..i64::from(plan.rows * plan.columns)) + ); + } + + /// Check all lanes fence TMEM reads before a single elected arrival releases the accumulator. + #[test] + fn accumulator_release_fences_every_lane_before_electing_one_arrival() { + let (mut ctx, module, block, a) = fixture(); + let plan = CuteSm100PipelinePlanAttr { + stages: 1, + transaction_bytes: 0, + consumer_arrivals: 8, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + }; + let wait = CuteSm100AccumulatorWaitOp::new(&mut ctx, vec![a[5], a[7]], plan); + wait.get_operation().insert_at_back(block, &ctx); + let release = CuteSm100AccumulatorReleaseOp::new(&mut ctx, vec![a[5]], plan); + release.get_operation().insert_at_back(block, &ctx); + let text = render(&mut ctx, &module, block); + assert!(text.contains("\"scf.while\""), "{text}"); + assert!(text.contains("#nvvm.mbar_scope"), "{text}"); + let fence = text.find("#nvvm.tcgen05_fence").unwrap(); + let elected = text.find("\"nvvm.elect.sync\"").unwrap(); + let conditional = text.find("\"scf.if\"").unwrap(); + let arrive = text.find("\"nvvm.mbarrier.arrive\"").unwrap(); + assert!( + fence < elected && elected < conditional && conditional < arrive, + "{text}" + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs new file mode 100644 index 0000000000..c00496f04a --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/elementwise.rs @@ -0,0 +1,1125 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Map tensor views and copies to the official CUTLASS `cute` dialect. +//! +//! Full tiles use CuTe operations; edge tiles use checked scalar accesses: +//! +//! ```text +//! tensor_make -> zipped_divide -> slice -> full-tile copy +//! | CUTLASS CuTe views and algorithms +//! | +//! +-> is_full / base / scalar tail store +//! ordinary arith + LLVM, using verified producer provenance +//! ``` +//! +//! CUTLASS handles view transformations and bulk copies. Edge stores still +//! need the original logical length and absolute index, so they map to scalar +//! arithmetic and LLVM memory operations. + +use crate::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use dialect_cute::{ + attributes::{ + CuteAlignmentAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, + CuteTensorRoleAttr, + }, + tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, + }, + types::CuteTensorViewType, +}; +use pliron::{ + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +/// Register the elementwise tensor-view mapping for the pinned CUTLASS profile. +/// +/// The caller must register builtin and MIR-core packs as well: this pack +/// intentionally reuses their scalar, pointer, function, and CFG mappings. +pub fn register_cute_elementwise_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_type::(TensorViewTypeTranslation)?; + register_cute_elementwise_mappings(registry) +} + +/// Register the elementwise operations and attributes while leaving the +/// shared `cute.tensor_view` type slot to a composed profile. +pub(crate) fn register_cute_elementwise_mappings( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + // These source attributes have already been consumed by the recipes or + // are encoded in the translated CUTLASS type. They must still be registered + // so the exporter's missing-mapping preflight can distinguish an explicit + // drop from an accidental omission. + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::( + GlobalStoreNoAllocateTranslation, + )?; + + registry.register_operation::(TensorMakeTranslation)?; + registry.register_operation::(TensorZippedDivideTranslation)?; + registry.register_operation::(TensorSliceTranslation)?; + registry.register_operation::(TensorIsFullTranslation)?; + registry.register_operation::(TensorBaseTranslation)?; + registry.register_operation::(TensorLoadIntoTranslation)?; + registry.register_operation::(TensorStoreFromTranslation)?; + registry.register_operation::(TensorStoreElementAbsTranslation)?; + Ok(()) +} + +/// Preserve reviewed elementwise view layouts as CuTe memref types. +pub(crate) struct TensorViewTypeTranslation; + +impl TypeTranslation for TensorViewTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let view = source_ref + .downcast_ref::() + .ok_or_else(|| "expected cute.tensor_view".to_owned())?; + require_elementwise_view(view)?; + + let scalar = transport_scalar(ctx, view)?; + let element = scalar_spelling(®istry.translate_type(ctx, scalar)?)?; + let space = address_space_spelling(view.space)?; + let layout = layout_spelling(view.layout)?; + MlirType::dialect(format!( + "!cute.memref<{element}, {space}, align<{}>, \"{layout}\">", + view.alignment.0 + )) + } +} + +/// Build a contiguous global CuTe view with an explicit LLVM-to-CuTe pointer bridge. +struct TensorMakeTranslation; + +impl OperationTranslation for TensorMakeTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_make", &input, 1, 2)?; + let source_result = source.deref(ctx).get_result(0); + let view = source_view(ctx, source_result)?; + require_elementwise_view(&view)?; + if view.layout != CuteTensorLayoutAttr::Contiguous1D { + return Err("cute.tensor_make requires a contiguous 1D result".into()); + } + + let scalar = transport_scalar(ctx, &view)?; + let element = scalar_spelling(&session.translate_type(scalar).map_err(stringify)?)?; + let space = address_space_spelling(view.space)?; + let pointer_type = cute_pointer_type(&element, space, view.alignment.0)?; + let shape_type = MlirType::dialect("!cute.shape<\"?{i64}\">")?; + let layout_type = MlirType::dialect("!cute.layout<\"?{i64}:1\">")?; + let location = input.location.clone(); + + // Rust slice pointers use the generic LLVM address space in their + // stable two-word ABI. CUTLASS's gmem pointer lowers to NVVM address space + // 1, so make that transition explicit before crossing the temporary + // LLVM/CuTe type bridge. CUTLASS can then reconcile the two opposite + // unrealized casts around the CuTe pointer without trying to use one + // as an address-space cast. + let (mut operations, pointer_use) = bridge_llvm_pointer_to_cute( + session, + input.operands[0].clone(), + pointer_type, + 1, + &location, + )?; + + let (shape_result, shape_use) = fresh(session, shape_type); + let shape = operation( + "cute.make_shape", + vec![shape_result], + vec![input.operands[1].clone()], + &location, + )?; + + let (layout_result, layout_use) = fresh(session, layout_type); + let mut layout = operation( + "cute.make_layout", + vec![layout_result], + vec![shape_use], + &location, + )?; + layout.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 0]), + ); + + let view = operation( + "cute.make_view", + input.results, + vec![pointer_use, layout_use], + &location, + )?; + operations.extend([shape, layout, view]); + Ok(operations) + } +} + +/// Divide the verified one-dimensional view into statically sized CuTe tiles. +struct TensorZippedDivideTranslation; + +impl OperationTranslation for TensorZippedDivideTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_zipped_divide", &input, 1, 1)?; + let result = source.deref(ctx).get_result(0); + let view = source_view(ctx, result)?; + require_elementwise_view(&view)?; + let CuteTensorLayoutAttr::Zipped1D(tile_size) = view.layout else { + return Err("cute.tensor_zipped_divide requires a Zipped1D result".into()); + }; + + let tiler_type = MlirType::dialect(format!("!cute.shape<\"{tile_size}\">"))?; + let (tiler_result, tiler_use) = fresh(session, tiler_type); + let tiler = operation("cute.static", vec![tiler_result], vec![], &input.location)?; + let divide = operation( + "cute.zipped_divide", + input.results, + vec![input.operands[0].clone(), tiler_use], + &input.location, + )?; + Ok(vec![tiler, divide]) + } +} + +/// Select one tile using its runtime coordinate and verified layout. +struct TensorSliceTranslation; + +impl OperationTranslation for TensorSliceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_slice", &input, 1, 2)?; + let result = source.deref(ctx).get_result(0); + let view = source_view(ctx, result)?; + require_elementwise_view(&view)?; + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) { + return Err("cute.tensor_slice requires a Tile1D result".into()); + } + + let coord_type = MlirType::dialect("!cute.coord<\"(_,?{i64})\">")?; + let (coord_result, coord_use) = fresh(session, coord_type); + let coord = operation( + "cute.make_coord", + vec![coord_result], + vec![input.operands[1].clone()], + &input.location, + )?; + let slice = operation( + "cute.slice", + input.results, + vec![input.operands[0].clone(), coord_use], + &input.location, + )?; + Ok(vec![coord, slice]) + } +} + +/// Recover an absolute element index from the verified view's producer chain. +struct TensorBaseTranslation; + +impl OperationTranslation for TensorBaseTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_base", &input, 1, 1)?; + let tensor = source.deref(ctx).get_operand(0); + let state = tensor_tile_state(ctx, tensor)?; + let final_result = input.results[0].clone(); + let (operations, _) = + emit_saturating_tile_base(session, &state, final_result, &input.location)?; + Ok(operations) + } +} + +/// Check that the selected tile lies wholly within the original logical length. +struct TensorIsFullTranslation; + +impl OperationTranslation for TensorIsFullTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_is_full", &input, 1, 1)?; + let tensor = source.deref(ctx).get_operand(0); + let state = tensor_tile_state(ctx, tensor)?; + let location = input.location.clone(); + let i64_type = MlirType::Integer(64); + let i1_type = MlirType::Integer(1); + + let (base_result, base_use) = fresh(session, i64_type.clone()); + let (mut operations, _) = + emit_saturating_tile_base(session, &state, base_result, &location)?; + + let (width_result, width_use) = fresh(session, i64_type.clone()); + operations.push(integer_constant( + width_result, + u64_bits(state.tile_size), + &location, + )?); + + let len = session.target_value_use(state.len).map_err(stringify)?; + let (in_range_result, in_range_use) = fresh(session, i1_type.clone()); + operations.push(integer_compare( + in_range_result, + base_use.clone(), + len.clone(), + 6, + &location, + )?); + + let (remaining_result, remaining_use) = fresh(session, i64_type); + operations.push(binary( + "arith.subi", + remaining_result, + len, + base_use, + &location, + )?); + + let (enough_result, enough_use) = fresh(session, i1_type); + operations.push(integer_compare( + enough_result, + remaining_use, + width_use, + 9, + &location, + )?); + + operations.push(binary( + "arith.andi", + input.results[0].clone(), + in_range_use, + enough_use, + &location, + )?); + Ok(operations) + } +} + +/// Copy a full global tile into its register fragment. +struct TensorLoadIntoTranslation; + +impl OperationTranslation for TensorLoadIntoTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_load_into", &input, 0, 2)?; + let source_op = CuteTensorLoadIntoOp::wrap(source); + let alignment = source_op + .get_attr_load_alignment_bytes(ctx) + .map(|attribute| attribute.0) + .ok_or_else(|| "cute.tensor_load_into has no alignment promise".to_owned())?; + let tensor_value = source.deref(ctx).get_operand(0); + emit_full_tile_copy( + ctx, + session, + input.operands[0].clone(), + input.operands[1].clone(), + tensor_value, + alignment, + CopyDirection::Load, + &input.location, + ) + } +} + +/// Copy a register fragment into a full global tile. +struct TensorStoreFromTranslation; + +impl OperationTranslation for TensorStoreFromTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_store_from", &input, 0, 2)?; + let source_op = CuteTensorStoreFromOp::wrap(source); + let alignment = source_op + .get_attr_store_alignment_bytes(ctx) + .map(|attribute| attribute.0) + .ok_or_else(|| "cute.tensor_store_from has no alignment promise".to_owned())?; + let tensor_value = source.deref(ctx).get_operand(1); + emit_full_tile_copy( + ctx, + session, + input.operands[1].clone(), + input.operands[0].clone(), + tensor_value, + alignment, + CopyDirection::Store, + &input.location, + ) + } +} + +/// Store one element only when its absolute index is within the original view. +struct TensorStoreElementAbsTranslation; + +impl OperationTranslation for TensorStoreElementAbsTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_store_element_abs", &input, 0, 3)?; + let source_ref = source.deref(ctx); + let tensor = source_ref.get_operand(0); + let absolute_index = source_ref.get_operand(1); + let value = source_ref.get_operand(2); + drop(source_ref); + + let state = tensor_tile_state(ctx, tensor)?; + let data = session.target_value_use(state.data).map_err(stringify)?; + let index = session + .target_value_use(absolute_index) + .map_err(stringify)?; + let scalar = session.target_value_use(value).map_err(stringify)?; + let element_type = scalar.ty.clone(); + let location = input.location.clone(); + + let (pointer_result, pointer_use) = fresh(session, data.ty.clone()); + let mut gep = operation( + "llvm.getelementptr", + vec![pointer_result], + vec![data, index], + &location, + )?; + gep.properties + .insert("elem_type".into(), MlirAttribute::Type(element_type)); + gep.properties.insert( + "noWrapFlags".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(32), + }, + ); + gep.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![i32::MIN]), + ); + + let mut store = operation("llvm.store", vec![], vec![scalar, pointer_use], &location)?; + store.properties.insert( + "ordering".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(64), + }, + ); + Ok(vec![gep, store]) + } +} + +/// Runtime origin and tile metadata recovered from verified tensor producers. +#[derive(Clone, Copy)] +struct TensorViewState { + /// Original global data pointer before tile transformations. + data: Value, + /// Original logical element count used for bounds checks. + len: Value, + /// Number of elements in each statically sized tile. + tile_size: u64, + /// Runtime tile coordinate relative to the original view. + tile_index: Value, +} + +/// Follow verified producer operations to recover the original pointer and tile metadata. +fn resolve_tensor_view( + ctx: &Context, + value: Value, + depth: usize, +) -> Result<(Value, Value, Option, Option), String> { + if depth > 16 { + return Err("cute tensor-view producer chain is unexpectedly deep".into()); + } + let defining_op = value.defining_op().ok_or_else(|| { + "cute tensor view reached a block argument; the elementwise mapping needs a direct make/zipped/slice chain".to_owned() + })?; + let opid = Operation::get_opid(defining_op, ctx); + if opid == CuteTensorMakeOp::get_opid_static() { + let op = defining_op.deref(ctx); + return Ok((op.get_operand(0), op.get_operand(1), None, None)); + } + if opid == CuteTensorZippedDivideOp::get_opid_static() { + let operand = defining_op.deref(ctx).get_operand(0); + let (data, len, _, tile_index) = resolve_tensor_view(ctx, operand, depth + 1)?; + let view = source_view(ctx, value)?; + let tile_size = view.layout.tile_size().ok_or_else(|| { + "cute.tensor_zipped_divide result has no static tile width".to_owned() + })?; + return Ok((data, len, Some(tile_size), tile_index)); + } + if opid == CuteTensorSliceOp::get_opid_static() { + let op = defining_op.deref(ctx); + let operand = op.get_operand(0); + let selected_index = op.get_operand(1); + drop(op); + let (data, len, _, _) = resolve_tensor_view(ctx, operand, depth + 1)?; + let view = source_view(ctx, value)?; + let tile_size = view + .selected_tile_size() + .ok_or_else(|| "cute.tensor_slice result has no selected tile width".to_owned())?; + return Ok((data, len, Some(tile_size), Some(selected_index))); + } + Err(format!( + "tensor view is produced by unsupported operation `{opid}`" + )) +} + +/// Require a sliced tile and recover its original logical bounds. +fn tensor_tile_state(ctx: &Context, value: Value) -> Result { + let (data, len, tile_size, tile_index) = resolve_tensor_view(ctx, value, 0)?; + Ok(TensorViewState { + data, + len, + tile_size: tile_size + .ok_or_else(|| "tensor operation needs a selected tile width".to_owned())?, + tile_index: tile_index + .ok_or_else(|| "tensor operation needs a selected tile index".to_owned())?, + }) +} + +/// Compute tile_index times tile_size without wrapping past the logical end. +fn emit_saturating_tile_base( + session: &mut TranslationSession<'_>, + state: &TensorViewState, + final_result: MlirResult, + location: &MlirLocation, +) -> Result<(Vec, MlirValueUse), String> { + let i64_type = MlirType::Integer(64); + let i1_type = MlirType::Integer(1); + let tile_index = session + .target_value_use(state.tile_index) + .map_err(stringify)?; + let mut operations = vec![]; + + let (width_result, width_use) = fresh(session, i64_type.clone()); + operations.push(integer_constant( + width_result, + u64_bits(state.tile_size), + location, + )?); + + let (product_result, product_use) = fresh(session, i64_type.clone()); + operations.push(binary( + "arith.muli", + product_result, + tile_index.clone(), + width_use, + location, + )?); + + let (limit_result, limit_use) = fresh(session, i64_type.clone()); + operations.push(integer_constant( + limit_result, + u64_bits(u64::MAX / state.tile_size), + location, + )?); + + let (overflow_result, overflow_use) = fresh(session, i1_type); + operations.push(integer_compare( + overflow_result, + limit_use, + tile_index, + 6, + location, + )?); + + let (extended_result, extended_use) = fresh(session, i64_type.clone()); + operations.push(unary( + "arith.extui", + extended_result, + overflow_use, + location, + )?); + + let (zero_result, zero_use) = fresh(session, i64_type); + operations.push(integer_constant(zero_result, 0, location)?); + + let (mask_result, mask_use) = fresh(session, MlirType::Integer(64)); + operations.push(binary( + "arith.subi", + mask_result, + zero_use, + extended_use, + location, + )?); + + let final_use = MlirValueUse { + id: final_result.id, + ty: final_result.ty.clone(), + }; + operations.push(binary( + "arith.ori", + final_result, + product_use, + mask_use, + location, + )?); + Ok((operations, final_use)) +} + +/// Select the source and destination spaces for a full-tile register copy. +#[derive(Clone, Copy)] +enum CopyDirection { + /// Copy global-memory elements into a register fragment. + Load, + /// Copy a register fragment into global-memory elements. + Store, +} + +/// Build the CuTe layouts and copy operation for a verified complete tile. +#[allow(clippy::too_many_arguments)] +fn emit_full_tile_copy( + ctx: &Context, + session: &mut TranslationSession<'_>, + tensor: MlirValueUse, + carrier: MlirValueUse, + source_tensor: Value, + promised_alignment: u64, + direction: CopyDirection, + location: &MlirLocation, +) -> Result, String> { + let view = source_view(ctx, source_tensor)?; + require_elementwise_view(&view)?; + let tile_size = view + .selected_tile_size() + .ok_or_else(|| "full-tile copy requires a selected Tile1D view".to_owned())?; + let storage_bytes = view + .storage_bytes(ctx) + .ok_or_else(|| "full-tile copy storage has no whole-byte width".to_owned())?; + let transfer_bytes = tile_size + .checked_mul(storage_bytes) + .ok_or_else(|| "full-tile copy byte width overflows u64".to_owned())?; + if !matches!(transfer_bytes, 4 | 8 | 16) { + return Err(format!( + "CUTLASS universal_copy mapping supports 4, 8, or 16 bytes, got {transfer_bytes}" + )); + } + if promised_alignment < transfer_bytes || !promised_alignment.is_power_of_two() { + return Err(format!( + "full-tile copy alignment {promised_alignment} does not cover {transfer_bytes} bytes" + )); + } + + let scalar = transport_scalar(ctx, &view)?; + let element = scalar_spelling(&session.translate_type(scalar).map_err(stringify)?)?; + let space = address_space_spelling(view.space)?; + let natural_pointer = cute_pointer_type(&element, space, view.alignment.0)?; + let aligned_pointer = cute_pointer_type(&element, space, promised_alignment)?; + let aligned_tensor = cute_memref_type( + &element, + space, + promised_alignment, + &format!("({tile_size}):(1)"), + )?; + let carrier_pointer = cute_pointer_type(&element, "rmem", transfer_bytes)?; + let carrier_view = cute_memref_type( + &element, + "rmem", + transfer_bytes, + &format!("({tile_size}):(1)"), + )?; + let constrained = MlirType::dialect(format!("!cute.i64"))?; + let atom_type = MlirType::dialect(format!( + "!cute_nvgpu.atom.universal_copy<{element}, {}b>", + transfer_bytes * 8 + ))?; + + let mut operations = vec![]; + let (iter_result, iter_use) = fresh(session, natural_pointer); + operations.push(operation( + "cute.get_iter", + vec![iter_result], + vec![tensor], + location, + )?); + + let (address_result, address_use) = fresh(session, MlirType::Integer(64)); + operations.push(operation( + "cute.ptrtoint", + vec![address_result], + vec![iter_use], + location, + )?); + + let (assumed_result, assumed_use) = fresh(session, constrained); + operations.push(operation( + "cute.assume", + vec![assumed_result], + vec![address_use], + location, + )?); + + let (aligned_pointer_result, aligned_pointer_use) = fresh(session, aligned_pointer); + operations.push(operation( + "cute.inttoptr", + vec![aligned_pointer_result], + vec![assumed_use], + location, + )?); + + let (aligned_view_result, aligned_view_use) = fresh(session, aligned_tensor); + operations.push(operation( + "cute.make_view", + vec![aligned_view_result], + vec![aligned_pointer_use], + location, + )?); + + // CuTe rmem pointers deliberately lower to generic LLVM pointers because + // NVVM allocas cannot use the local-memory address space. Synthetic and + // hand-authored MIR may still present an explicit local pointer, so + // normalize it to AS0 before the temporary CuTe bridge as well. + let (carrier_bridge, carrier_pointer_use) = + bridge_llvm_pointer_to_cute(session, carrier, carrier_pointer, 0, location)?; + operations.extend(carrier_bridge); + + let (carrier_view_result, carrier_view_use) = fresh(session, carrier_view); + operations.push(operation( + "cute.make_view", + vec![carrier_view_result], + vec![carrier_pointer_use], + location, + )?); + + let (atom_result, atom_use) = fresh(session, atom_type); + operations.push(operation( + "cute.make_atom", + vec![atom_result], + vec![], + location, + )?); + + let (source, destination) = match direction { + CopyDirection::Load => (aligned_view_use, carrier_view_use), + CopyDirection::Store => (carrier_view_use, aligned_view_use), + }; + let mut copy = operation( + "cute.copy", + vec![], + vec![atom_use, source, destination], + location, + )?; + // CUTLASS 4.7's AlgorithmCopyOp uses AttrSizedOperandSegments. Generic + // syntax does not infer these four groups, and a missing property reparses + // as all-zero segments. Canonicalization may then erase the copy as an + // effect-free operation, so preserve the exact atom/src/dst/predicate + // segmentation explicitly. + copy.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 1, 0]), + ); + operations.push(copy); + Ok(operations) +} + +/// Reject elementwise view formats, roles, spaces, or layouts outside the reviewed subset. +fn require_elementwise_view(view: &CuteTensorViewType) -> Result<(), String> { + if view.format != CuteTensorFormatAttr::Plain { + return Err(format!( + "the elementwise CUTLASS pack supports Plain tensor views, got {:?}", + view.format + )); + } + if view.role != CuteTensorRoleAttr::Generic { + return Err(format!( + "the elementwise CUTLASS pack supports Generic tensor roles, got {:?}", + view.role + )); + } + if view.space != CuteTensorAddressSpaceAttr::Gmem { + return Err(format!( + "the elementwise CUTLASS pack supports global-memory tensor views, got {:?}", + view.space + )); + } + if !matches!( + view.layout, + CuteTensorLayoutAttr::Contiguous1D + | CuteTensorLayoutAttr::Zipped1D(_) + | CuteTensorLayoutAttr::Tile1D(_) + ) { + return Err(format!( + "the elementwise CUTLASS pack does not map layout {:?}", + view.layout + )); + } + Ok(()) +} + +/// Recover a CuTe tensor-view type from a source SSA value or report a type error. +fn source_view(ctx: &Context, value: Value) -> Result { + let ty = value.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected a cute.tensor_view value".to_owned()) +} + +/// Map the source memory-space contract to CuTe's gmem or smem spelling. +fn address_space_spelling(space: CuteTensorAddressSpaceAttr) -> Result<&'static str, String> { + match space { + CuteTensorAddressSpaceAttr::Gmem => Ok("gmem"), + other => Err(format!( + "the elementwise CUTLASS pack does not map address space {other:?}" + )), + } +} + +/// Spell a reviewed elementwise layout in CUTLASS syntax. +fn layout_spelling(layout: CuteTensorLayoutAttr) -> Result { + match layout { + CuteTensorLayoutAttr::Contiguous1D => Ok("?{i64}:1".into()), + CuteTensorLayoutAttr::Zipped1D(tile) if tile > 0 => { + Ok(format!("({tile},?{{i64}}):(1,{tile})")) + } + CuteTensorLayoutAttr::Tile1D(tile) if tile > 0 => Ok(format!("({tile}):(1)")), + other => Err(format!( + "the elementwise CUTLASS pack does not map layout {other:?}" + )), + } +} + +/// Resolve the validated scalar bit representation without changing the source element type. +fn transport_scalar(ctx: &Context, view: &CuteTensorViewType) -> Result { + view.transport_scalar(ctx) + .ok_or_else(|| "elementwise tensor storage has no proven scalar transport layout".into()) +} + +/// Spell supported storage scalars for a CuTe pointer or memref type. +fn scalar_spelling(ty: &MlirType) -> Result { + match ty { + MlirType::Integer(width) => Ok(format!("i{width}")), + MlirType::Float(crate::MlirFloatType::F16) => Ok("f16".into()), + MlirType::Float(crate::MlirFloatType::F32) => Ok("f32".into()), + other => Err(format!( + "the elementwise CUTLASS pack needs an i8, i16, f16, or f32 scalar, got {other:?}" + )), + } +} + +/// Build a CuTe pointer spelling with element type, memory space, and byte alignment. +fn cute_pointer_type(element: &str, space: &str, alignment: u64) -> Result { + MlirType::dialect(format!("!cute.ptr<{element}, {space}, align<{alignment}>>")) +} + +/// Cross from an LLVM pointer to a CuTe pointer without asking an unrealized +/// cast to change address spaces. +/// +/// CUTLASS lowers the CuTe pointer back to its canonical LLVM address space. When +/// the incoming pointer already has that address space, the two temporary +/// unrealized casts reconcile away. When it does not, the explicit +/// `llvm.addrspacecast` remains and the temporary casts still reconcile away. +fn bridge_llvm_pointer_to_cute( + session: &mut TranslationSession<'_>, + pointer: MlirValueUse, + cute_pointer: MlirType, + target_address_space: u32, + location: &MlirLocation, +) -> Result<(Vec, MlirValueUse), String> { + let source_address_space = llvm_pointer_address_space(&pointer.ty)?; + let mut operations = vec![]; + let pointer = if source_address_space == target_address_space { + pointer + } else { + let target_type = llvm_pointer_type(target_address_space)?; + let (result, use_) = fresh(session, target_type); + operations.push(operation( + "llvm.addrspacecast", + vec![result], + vec![pointer], + location, + )?); + use_ + }; + + let (result, use_) = fresh(session, cute_pointer); + operations.push(operation( + "builtin.unrealized_conversion_cast", + vec![result], + vec![pointer], + location, + )?); + Ok((operations, use_)) +} + +/// Read the numeric address space from an opaque LLVM pointer type. +fn llvm_pointer_address_space(ty: &MlirType) -> Result { + let MlirType::Dialect(spelling) = ty else { + return Err(format!( + "CuTe pointer bridge expected an LLVM pointer, got {ty:?}" + )); + }; + if spelling == "!llvm.ptr" { + return Ok(0); + } + let Some(address_space) = spelling + .strip_prefix("!llvm.ptr<") + .and_then(|suffix| suffix.strip_suffix('>')) + else { + return Err(format!( + "CuTe pointer bridge expected an LLVM pointer, got {spelling}" + )); + }; + address_space.parse::().map_err(|_| { + format!("CuTe pointer bridge could not parse LLVM address space in {spelling}") + }) +} + +/// Spell an opaque LLVM pointer in the requested numeric address space. +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +/// Build the CuTe memref type for a register-copy layout and its storage element. +fn cute_memref_type( + element: &str, + space: &str, + alignment: u64, + layout: &str, +) -> Result { + MlirType::dialect(format!( + "!cute.memref<{element}, {space}, align<{alignment}>, \"{layout}\">" + )) +} + +/// Require exact result/operand counts and no leftover attributes, regions, or successors. +fn expect_shape( + name: &str, + input: &OperationInput, + results: usize, + operands: usize, +) -> Result<(), String> { + if input.results.len() != results || input.operands.len() != operands { + return Err(format!( + "{name} expected {results} results and {operands} operands, got {} and {}", + input.results.len(), + input.operands.len() + )); + } + if !input.successors.is_empty() || !input.regions.is_empty() || !input.attributes.is_empty() { + return Err(format!( + "{name} unexpectedly retained successors, regions, or source-only attributes" + )); + } + Ok(()) +} + +/// Build a target operation with the supplied operands, results, and source location. +fn operation( + name: &str, + results: Vec, + operands: Vec, + location: &MlirLocation, +) -> Result { + let mut operation = MlirOperation::new(name)?; + operation.results = results; + operation.operands = operands; + operation.location = location.clone(); + Ok(operation) +} + +/// Allocate one SSA identifier and return matching result and operand handles. +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +/// Build an arith.constant using the supplied result's integer type. +fn integer_constant( + result: MlirResult, + value: i128, + location: &MlirLocation, +) -> Result { + let mut operation = operation("arith.constant", vec![result.clone()], vec![], location)?; + operation.properties.insert( + "value".into(), + MlirAttribute::Integer { + value, + ty: result.ty, + }, + ); + Ok(operation) +} + +/// Build an arith.cmpi using the explicit target predicate number. +fn integer_compare( + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + predicate: i128, + location: &MlirLocation, +) -> Result { + let mut operation = operation("arith.cmpi", vec![result], vec![left, right], location)?; + operation.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: predicate, + ty: MlirType::Integer(64), + }, + ); + Ok(operation) +} + +/// Build a two-operand target operation using the supplied result. +fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +/// Build a one-operand operation with the supplied result and source location. +fn unary( + name: &str, + result: MlirResult, + operand: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![operand], location) +} + +/// Preserve all 64 bits while spelling an unsigned value as a signed MLIR literal. +fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +/// Preserve a source diagnostic as the string expected by translation recipes. +fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} + +/// Emit a global store preserving its no-allocate cache policy and vector-width contract. +struct GlobalStoreNoAllocateTranslation; +impl OperationTranslation for GlobalStoreNoAllocateTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + use dialect_cute::copy_ops::{CuteStoreGlobalNoAllocateOp, global_store_element_bytes}; + use pliron::common_traits::Verify; + expect_shape("cute.store_global_no_allocate", &input, 0, 2)?; + let op = CuteStoreGlobalNoAllocateOp::wrap(source); + op.verify(ctx).map_err(stringify)?; + let n = op.elements(ctx).ok_or("missing global store width")?; + let dst_ty = source.deref(ctx).get_operand(0).get_type(ctx); + let element_ty = dst_ty + .deref(ctx) + .downcast_ref::() + .ok_or("missing destination pointer type")? + .pointee; + let scalar_bytes = global_store_element_bytes(ctx, element_ty) + .ok_or("unsupported global store element")?; + let bytes = scalar_bytes * n; + let element = scalar_spelling(&session.translate_type(element_ty).map_err(stringify)?)?; + let location = &input.location; + let mut operations = vec![]; + let mut views = vec![]; + for (index, space, align, address_space) in [(0, "gmem", bytes, 1), (1, "rmem", 32, 0)] { + let pointer_ty = cute_pointer_type(&element, space, align)?; + let (bridge, pointer) = bridge_llvm_pointer_to_cute( + session, + input.operands[index].clone(), + pointer_ty, + address_space, + location, + )?; + operations.extend(bridge); + let view_ty = cute_memref_type(&element, space, align, &format!("({n}):(1)"))?; + let (result, value) = fresh(session, view_ty); + operations.push(operation( + "cute.make_view", + vec![result], + vec![pointer], + location, + )?); + views.push(value); + } + let atom_ty = MlirType::dialect(format!( + "!cute_nvgpu.atom.r2g<{element} copy_bits = {} l1_cache_evict_priority = no_allocate>", + bytes * 8 + ))?; + let (atom_result, atom) = fresh(session, atom_ty); + operations.push(operation( + "cute.make_atom", + vec![atom_result], + vec![], + location, + )?); + let mut copy = operation( + "cute.copy", + vec![], + vec![atom, views[1].clone(), views[0].clone()], + location, + )?; + copy.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 1, 0]), + ); + operations.push(copy); + Ok(operations) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs new file mode 100644 index 0000000000..f6f96c0a51 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/gemv.rs @@ -0,0 +1,1316 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! NVFP4 GEMV mappings for the official CUTLASS full-CuTe MLIR profile. +//! +//! Source views become LLVM aggregates holding runtime pointers and shapes. +//! Keeping these fields in SSA preserves the dependencies between operations: +//! +//! ```text +//! tensor_make_2d -> scaled_view_make -> row -> KTile<64> +//! | +//! v +//! f4/UE8M0 CuTe views <- selected byte addresses +//! | +//! v +//! load_vec -> cvt_fpext -> ordered f32 dot +//! ``` +//! +//! The mapping translates CuTe operations directly into operations accepted +//! by CUTLASS. + +use crate::cute::tensor::elementwise::{ + TensorViewTypeTranslation, register_cute_elementwise_mappings, +}; + +use crate::{ + MlirAttribute, MlirFloatType, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use dialect_cute::{ + attributes::{ + CuteScaledLayoutAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, + }, + gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, GEMV_K_TILE_WIDTH, + }, + types::{CuteFragmentType, CuteScaledViewType, CuteTensorViewType}, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Packed FP4 storage bytes consumed by one 64-element K tile. +const VALUE_BYTES_PER_TILE: u64 = 32; +/// Packed value bytes fetched by one vector load. +const VALUE_BYTES_PER_LOAD: u64 = 16; +/// FP4 values represented by one packed vector load. +const VALUE_LANES_PER_LOAD: u64 = 32; +/// Scale storage bytes associated with one K tile. +const SCALE_BYTES_PER_TILE: u64 = 4; +/// Row extent of the verified blocked scale layout. +const SCALE_ATOM_ROWS: u64 = 128; +/// K extent of one scale-layout atom. +const SCALE_ATOM_K: u64 = 64; +/// Storage bytes occupied by one blocked scale atom. +const SCALE_ATOM_BYTES: u64 = 512; + +/// Register the composed elementwise + NVFP4 GEMV mapping pack. +/// +/// `cute.tensor_view` has one registry owner. Its translator dispatches from +/// the verified source type and fails closed for every view outside the two +/// reviewed mapping families. +pub(crate) fn register_cute_cutlass_profile_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_type::(CombinedTensorViewTypeTranslation)?; + register_cute_elementwise_mappings(registry)?; + register_gemv_mappings(registry) +} + +/// Register GEMV-specific view, fragment, load, and ordered-dot mappings. +fn register_gemv_mappings(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_type::(ScaledViewTypeTranslation)?; + registry.register_type::(FragmentTypeTranslation)?; + + registry.register_operation::(TensorMake2DTranslation)?; + registry.register_operation::(ScaledViewMakeTranslation)?; + registry.register_operation::(ScaledViewRowTranslation)?; + registry.register_operation::(ScaledViewKTileTranslation)?; + registry.register_operation::(ScaledViewLoadTranslation)?; + registry.register_operation::(DotTranslation)?; + Ok(()) +} + +/// Dispatch verified tensor views to their GEMV carrier or elementwise CuTe representation. +struct CombinedTensorViewTypeTranslation; + +impl TypeTranslation for CombinedTensorViewTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let view = source_ref + .downcast_ref::() + .ok_or_else(|| "expected cute.tensor_view".to_owned())?; + if require_gemv_tensor_view(ctx, view).is_ok() { + return tensor_carrier_type(); + } + drop(source_ref); + TensorViewTypeTranslation.translate(ctx, source, registry) + } +} + +/// Represent a verified scaled view with its runtime pointer and shape carrier. +struct ScaledViewTypeTranslation; + +impl TypeTranslation for ScaledViewTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let view = source_ref + .downcast_ref::() + .ok_or_else(|| "expected cute.scaled_view".to_owned())?; + view.verify(ctx).map_err(stringify)?; + scaled_carrier_type() + } +} + +/// Represent a verified GEMV fragment as the ordered FP32 lane carrier. +struct FragmentTypeTranslation; + +impl TypeTranslation for FragmentTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let fragment = source_ref + .downcast_ref::() + .ok_or_else(|| "expected cute.fragment".to_owned())?; + fragment.verify(ctx).map_err(stringify)?; + fragment_carrier_type() + } +} + +/// Pack runtime pointer, rows, columns, and stride into the two-dimensional view carrier. +struct TensorMake2DTranslation; + +impl OperationTranslation for TensorMake2DTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.tensor_make_2d", &input, 1, 4)?; + CuteTensorMake2DOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + let result_view = tensor_view(ctx, source.deref(ctx).get_result(0))?; + require_gemv_tensor_view(ctx, &result_view)?; + build_aggregate( + session, + input.results[0].clone(), + input.operands, + &input.location, + ) + } +} + +/// Combine values and scales into the verified full scaled-view carrier. +struct ScaledViewMakeTranslation; + +impl OperationTranslation for ScaledViewMakeTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.scaled_view_make", &input, 1, 2)?; + CuteScaledViewMakeOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + let result = scaled_view(ctx, source.deref(ctx).get_result(0))?; + if result.layout != CuteScaledLayoutAttr::Full { + return Err("cute.scaled_view_make must produce Full layout".into()); + } + + let i64_type = MlirType::Integer(64); + let (zero_result, zero_use) = fresh(session, i64_type); + let zero = integer_constant(zero_result, 0, &input.location)?; + let mut operations = vec![zero]; + operations.extend(build_aggregate( + session, + input.results[0].clone(), + vec![ + input.operands[0].clone(), + input.operands[1].clone(), + zero_use.clone(), + zero_use.clone(), + zero_use, + ], + &input.location, + )?); + Ok(operations) + } +} + +/// Select the runtime row while retaining scale-layout and shape dependencies. +struct ScaledViewRowTranslation; + +impl OperationTranslation for ScaledViewRowTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.scaled_view_row", &input, 1, 3)?; + CuteScaledViewRowOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + let result = scaled_view(ctx, source.deref(ctx).get_result(0))?; + if result.layout != CuteScaledLayoutAttr::Row { + return Err("cute.scaled_view_row must produce Row layout".into()); + } + + let (partial_result, partial_use) = fresh(session, input.results[0].ty.clone()); + let batch = insert_value( + partial_result, + input.operands[0].clone(), + input.operands[1].clone(), + 2, + &input.location, + )?; + let row = insert_value( + input.results[0].clone(), + partial_use, + input.operands[2].clone(), + 3, + &input.location, + )?; + Ok(vec![batch, row]) + } +} + +/// Select the source row's 64-element K tile and corresponding scale coordinates. +struct ScaledViewKTileTranslation; + +impl OperationTranslation for ScaledViewKTileTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.scaled_view_k_tile", &input, 1, 2)?; + CuteScaledViewKTileOp::wrap(source) + .verify(ctx) + .map_err(stringify)?; + let result = scaled_view(ctx, source.deref(ctx).get_result(0))?; + if result.layout != CuteScaledLayoutAttr::KTile(GEMV_K_TILE_WIDTH) { + return Err("cute.scaled_view_k_tile must produce KTile<64>".into()); + } + Ok(vec![insert_value( + input.results[0].clone(), + input.operands[0].clone(), + input.operands[1].clone(), + 4, + &input.location, + )?]) + } +} + +/// Load one verified K64 tile and convert packed values/scales into FP32 fragments. +struct ScaledViewLoadTranslation; + +impl OperationTranslation for ScaledViewLoadTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.scaled_view_load", &input, 1, 1)?; + let source_op = CuteScaledViewLoadOp::wrap(source); + source_op.verify(ctx).map_err(stringify)?; + let value_alignment = source_op + .promised_value_alignment(ctx) + .ok_or_else(|| "cute.scaled_view_load has no value alignment promise".to_owned())?; + let scale_alignment = source_op + .promised_scale_alignment(ctx) + .ok_or_else(|| "cute.scaled_view_load has no scale alignment promise".to_owned())?; + if value_alignment < VALUE_BYTES_PER_LOAD || scale_alignment < SCALE_BYTES_PER_TILE { + return Err("K=64 GEMV load needs value alignment >=16 and scale alignment >=4".into()); + } + + let source_fragment = fragment(ctx, source.deref(ctx).get_result(0))?; + let source_view = source_fragment + .source_view(ctx) + .ok_or_else(|| "GEMV fragment lost its scaled-view provenance".to_owned())?; + if source_view.layout != CuteScaledLayoutAttr::KTile(GEMV_K_TILE_WIDTH) { + return Err("GEMV fragment must come from KTile<64>".into()); + } + let scale_format = source_view + .scales + .deref(ctx) + .downcast_ref::() + .ok_or_else(|| "GEMV scaled view lost scale tensor".to_owned())? + .format; + let scale_kind = match scale_format { + CuteTensorFormatAttr::E4M3 => NarrowLoad::E4M3, + _ => NarrowLoad::UE8M0, + }; + + emit_scaled_k64_load( + session, + input.operands[0].clone(), + input.results[0].clone(), + value_alignment, + scale_alignment, + scale_kind, + &input.location, + ) + } +} + +/// Accumulate the K64 fragment product in the source's explicit FP32 order. +struct DotTranslation; + +impl OperationTranslation for DotTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + expect_shape("cute.dot", &input, 1, 3)?; + CuteDotOp::wrap(source).verify(ctx).map_err(stringify)?; + + let source_ref = source.deref(ctx); + let matrix = fragment(ctx, source_ref.get_operand(0))? + .source_view(ctx) + .ok_or_else(|| "cute.dot matrix fragment lost provenance".to_owned())?; + let vector = fragment(ctx, source_ref.get_operand(1))? + .source_view(ctx) + .ok_or_else(|| "cute.dot vector fragment lost provenance".to_owned())?; + if matrix.role != CuteTensorRoleAttr::Mkl || vector.role != CuteTensorRoleAttr::Nkl { + return Err("cute.dot requires Mkl on the left and Nkl on the right".into()); + } + drop(source_ref); + + emit_ordered_k64_dot( + session, + input.operands[0].clone(), + input.operands[1].clone(), + input.operands[2].clone(), + input.results[0].clone(), + &input.location, + ) + } +} + +/// Load the K64 values and scales into separate FP32 vectors; the dot applies the scales. +fn emit_scaled_k64_load( + session: &mut TranslationSession<'_>, + scaled: MlirValueUse, + final_result: MlirResult, + value_alignment: u64, + scale_alignment: u64, + scale_kind: NarrowLoad, + location: &MlirLocation, +) -> Result, String> { + let tensor_type = tensor_carrier_type()?; + let pointer_type = llvm_pointer_type(0)?; + let i64_type = MlirType::Integer(64); + let mut operations = vec![]; + + let (values_result, values) = fresh(session, tensor_type.clone()); + operations.push(extract_value(values_result, scaled.clone(), 0, location)?); + let (scales_result, scales) = fresh(session, tensor_type); + operations.push(extract_value(scales_result, scaled.clone(), 1, location)?); + let (batch_result, batch) = fresh(session, i64_type.clone()); + operations.push(extract_value(batch_result, scaled.clone(), 2, location)?); + let (row_result, row) = fresh(session, i64_type.clone()); + operations.push(extract_value(row_result, scaled.clone(), 3, location)?); + let (tile_result, tile) = fresh(session, i64_type.clone()); + operations.push(extract_value(tile_result, scaled, 4, location)?); + + let (value_pointer_result, value_pointer) = fresh(session, pointer_type.clone()); + operations.push(extract_value( + value_pointer_result, + values.clone(), + 0, + location, + )?); + let (value_rows_result, value_rows) = fresh(session, i64_type.clone()); + operations.push(extract_value( + value_rows_result, + values.clone(), + 2, + location, + )?); + let (value_k_result, value_k) = fresh(session, i64_type.clone()); + operations.push(extract_value(value_k_result, values, 3, location)?); + + let (scale_pointer_result, scale_pointer) = fresh(session, pointer_type); + operations.push(extract_value( + scale_pointer_result, + scales.clone(), + 0, + location, + )?); + let (scale_rows_result, scale_rows) = fresh(session, i64_type.clone()); + operations.push(extract_value( + scale_rows_result, + scales.clone(), + 2, + location, + )?); + let (scale_k_result, scale_k) = fresh(session, i64_type.clone()); + operations.push(extract_value(scale_k_result, scales, 3, location)?); + + let two = emit_i64_constant(session, &mut operations, 2, location)?; + let packed_k = emit_binary( + session, + &mut operations, + "arith.divui", + value_k, + two, + location, + )?; + let batch_rows = emit_binary( + session, + &mut operations, + "arith.muli", + batch.clone(), + value_rows, + location, + )?; + let linear_row = emit_binary( + session, + &mut operations, + "arith.addi", + batch_rows, + row.clone(), + location, + )?; + let value_row_base = emit_binary( + session, + &mut operations, + "arith.muli", + linear_row, + packed_k, + location, + )?; + let value_stride = emit_i64_constant(session, &mut operations, VALUE_BYTES_PER_TILE, location)?; + let value_tile_offset = emit_binary( + session, + &mut operations, + "arith.muli", + tile.clone(), + value_stride, + location, + )?; + let value_base = emit_binary( + session, + &mut operations, + "arith.addi", + value_row_base, + value_tile_offset, + location, + )?; + + let rest_m = emit_ceil_div( + session, + &mut operations, + scale_rows, + SCALE_ATOM_ROWS, + location, + )?; + let rest_k = emit_ceil_div(session, &mut operations, scale_k, SCALE_ATOM_K, location)?; + let block_bytes = emit_i64_constant(session, &mut operations, SCALE_ATOM_BYTES, location)?; + let batch_blocks = emit_binary( + session, + &mut operations, + "arith.muli", + batch, + rest_m, + location, + )?; + let batch_blocks = emit_binary( + session, + &mut operations, + "arith.muli", + batch_blocks, + rest_k.clone(), + location, + )?; + let batch_bytes = emit_binary( + session, + &mut operations, + "arith.muli", + batch_blocks, + block_bytes.clone(), + location, + )?; + + let shift_seven = emit_i64_constant(session, &mut operations, 7, location)?; + let row_block = emit_binary( + session, + &mut operations, + "arith.shrui", + row.clone(), + shift_seven, + location, + )?; + let row_blocks = emit_binary( + session, + &mut operations, + "arith.muli", + row_block, + rest_k, + location, + )?; + let row_block_bytes = emit_binary( + session, + &mut operations, + "arith.muli", + row_blocks, + block_bytes, + location, + )?; + + let mask_31 = emit_i64_constant(session, &mut operations, 31, location)?; + let row_in_quadrant = emit_binary( + session, + &mut operations, + "arith.andi", + row.clone(), + mask_31, + location, + )?; + let shift_four = emit_i64_constant(session, &mut operations, 4, location)?; + let row_in_quadrant_bytes = emit_binary( + session, + &mut operations, + "arith.shli", + row_in_quadrant, + shift_four, + location, + )?; + + let shift_five = emit_i64_constant(session, &mut operations, 5, location)?; + let quadrant = emit_binary( + session, + &mut operations, + "arith.shrui", + row, + shift_five, + location, + )?; + let mask_three = emit_i64_constant(session, &mut operations, 3, location)?; + let quadrant = emit_binary( + session, + &mut operations, + "arith.andi", + quadrant, + mask_three, + location, + )?; + let shift_two = emit_i64_constant(session, &mut operations, 2, location)?; + let quadrant_bytes = emit_binary( + session, + &mut operations, + "arith.shli", + quadrant, + shift_two, + location, + )?; + + let scale_row_base = emit_binary( + session, + &mut operations, + "arith.addi", + batch_bytes, + row_block_bytes, + location, + )?; + let scale_row_base = emit_binary( + session, + &mut operations, + "arith.addi", + scale_row_base, + row_in_quadrant_bytes, + location, + )?; + let scale_row_base = emit_binary( + session, + &mut operations, + "arith.addi", + scale_row_base, + quadrant_bytes, + location, + )?; + let scale_stride = emit_i64_constant(session, &mut operations, SCALE_ATOM_BYTES, location)?; + let scale_tile_offset = emit_binary( + session, + &mut operations, + "arith.muli", + tile, + scale_stride, + location, + )?; + let scale_base = emit_binary( + session, + &mut operations, + "arith.addi", + scale_row_base, + scale_tile_offset, + location, + )?; + + let value_lo = emit_narrow_load( + session, + &mut operations, + value_pointer.clone(), + value_base.clone(), + NarrowLoad::E2M1, + value_alignment, + location, + )?; + let sixteen = emit_i64_constant(session, &mut operations, VALUE_BYTES_PER_LOAD, location)?; + let value_high_base = emit_binary( + session, + &mut operations, + "arith.addi", + value_base, + sixteen, + location, + )?; + let value_hi = emit_narrow_load( + session, + &mut operations, + value_pointer, + value_high_base, + NarrowLoad::E2M1, + value_alignment, + location, + )?; + let scale_values = emit_narrow_load( + session, + &mut operations, + scale_pointer, + scale_base, + scale_kind, + scale_alignment, + location, + )?; + + operations.extend(build_aggregate( + session, + final_result, + vec![value_lo, value_hi, scale_values], + location, + )?); + Ok(operations) +} + +/// Select the narrow storage format and vector shape for a CuTe load. +#[derive(Clone, Copy)] +enum NarrowLoad { + /// Packed FP4 values with two elements per byte. + E2M1, + /// Unsigned exponent-only byte scales. + UE8M0, + /// Eight-bit E4M3FN byte scales. + E4M3, +} + +/// Build a narrow-format CuTe memref and load its verified contiguous register vector. +fn emit_narrow_load( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + base_pointer: MlirValueUse, + byte_offset: MlirValueUse, + kind: NarrowLoad, + alignment: u64, + location: &MlirLocation, +) -> Result { + let pointer_type = llvm_pointer_type(0)?; + let (offset_result, offset_pointer) = fresh(session, pointer_type); + let mut gep = operation( + "llvm.getelementptr", + vec![offset_result], + vec![base_pointer, byte_offset], + location, + )?; + gep.properties.insert( + "elem_type".into(), + MlirAttribute::Type(MlirType::Integer(8)), + ); + gep.properties.insert( + "noWrapFlags".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(32), + }, + ); + gep.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![i32::MIN]), + ); + operations.push(gep); + + let global_pointer_type = llvm_pointer_type(1)?; + let (global_result, global_pointer) = fresh(session, global_pointer_type); + operations.push(operation( + "llvm.addrspacecast", + vec![global_result], + vec![offset_pointer], + location, + )?); + + let (element, lanes, intermediate) = match kind { + NarrowLoad::E2M1 => ("f4E2M1FN", VALUE_LANES_PER_LOAD, "f16"), + NarrowLoad::UE8M0 => ("f8E8M0FNU", SCALE_BYTES_PER_TILE, "bf16"), + NarrowLoad::E4M3 => ("f8E4M3FN", SCALE_BYTES_PER_TILE, "f16"), + }; + // ptrtoint intentionally accepts only the canonical byte-aligned pointer + // form. State align<1> before strengthening the selected address below. + let loose_pointer_type = MlirType::dialect(format!("!cute.ptr<{element}, gmem, align<1>>"))?; + let (loose_result, loose_pointer) = fresh(session, loose_pointer_type); + operations.push(operation( + "builtin.unrealized_conversion_cast", + vec![loose_result], + vec![global_pointer], + location, + )?); + + let (address_result, address) = fresh(session, MlirType::Integer(64)); + operations.push(operation( + "cute.ptrtoint", + vec![address_result], + vec![loose_pointer], + location, + )?); + let constrained_type = MlirType::dialect(format!("!cute.i64"))?; + let (constrained_result, constrained) = fresh(session, constrained_type); + operations.push(operation( + "cute.assume", + vec![constrained_result], + vec![address], + location, + )?); + let aligned_pointer_type = + MlirType::dialect(format!("!cute.ptr<{element}, gmem, align<{alignment}>>"))?; + let (aligned_result, aligned_pointer) = fresh(session, aligned_pointer_type); + operations.push(operation( + "cute.inttoptr", + vec![aligned_result], + vec![constrained], + location, + )?); + + let memref_type = MlirType::dialect(format!( + "!cute.memref<{element}, gmem, align<{alignment}>, \"({lanes}):(1)\">" + ))?; + let (view_result, view) = fresh(session, memref_type); + operations.push(operation( + "cute.make_view", + vec![view_result], + vec![aligned_pointer], + location, + )?); + + let narrow_type = narrow_vector_type(kind); + let (load_result, loaded) = fresh(session, narrow_type); + let mut load = operation( + "cute.memref.load_vec", + vec![load_result], + vec![view], + location, + )?; + load.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 0, 0]), + ); + operations.push(load); + + let intermediate_type = vector_type(lanes, narrow_float(intermediate)); + let (convert_result, converted) = fresh(session, intermediate_type); + operations.push(operation( + "nvgpu.cvt_fpext", + vec![convert_result], + vec![loaded], + location, + )?); + + let f32_type = vector_type(lanes, MlirType::Float(MlirFloatType::F32)); + let (extended_result, extended) = fresh(session, f32_type); + operations.push(operation( + "arith.extf", + vec![extended_result], + vec![converted], + location, + )?); + Ok(extended) +} + +/// Emit the 64 FP32 products and additions in source order without reassociation. +fn emit_ordered_k64_dot( + session: &mut TranslationSession<'_>, + matrix: MlirValueUse, + vector: MlirValueUse, + mut accumulator: MlirValueUse, + final_result: MlirResult, + location: &MlirLocation, +) -> Result, String> { + let value_vector_type = vector_type(VALUE_LANES_PER_LOAD, MlirType::Float(MlirFloatType::F32)); + let scale_vector_type = vector_type(SCALE_BYTES_PER_TILE, MlirType::Float(MlirFloatType::F32)); + let mut operations = vec![]; + + let mut matrix_values = Vec::with_capacity(2); + let mut vector_values = Vec::with_capacity(2); + for position in 0..2 { + let (result, value) = fresh(session, value_vector_type.clone()); + operations.push(extract_value(result, matrix.clone(), position, location)?); + matrix_values.push(value); + let (result, value) = fresh(session, value_vector_type.clone()); + operations.push(extract_value(result, vector.clone(), position, location)?); + vector_values.push(value); + } + let (matrix_scales_result, matrix_scales) = fresh(session, scale_vector_type.clone()); + operations.push(extract_value(matrix_scales_result, matrix, 2, location)?); + let (vector_scales_result, vector_scales) = fresh(session, scale_vector_type); + operations.push(extract_value(vector_scales_result, vector, 2, location)?); + + for group in 0..SCALE_BYTES_PER_TILE { + let scale_a = emit_vector_extract( + session, + &mut operations, + matrix_scales.clone(), + group, + location, + )?; + let scale_b = emit_vector_extract( + session, + &mut operations, + vector_scales.clone(), + group, + location, + )?; + for lane_in_group in 0..16 { + let lane = group * 16 + lane_in_group; + let chunk = usize::try_from(lane / VALUE_LANES_PER_LOAD) + .map_err(|_| "GEMV vector chunk does not fit usize")?; + let chunk_lane = lane % VALUE_LANES_PER_LOAD; + let a = emit_vector_extract( + session, + &mut operations, + matrix_values[chunk].clone(), + chunk_lane, + location, + )?; + let b = emit_vector_extract( + session, + &mut operations, + vector_values[chunk].clone(), + chunk_lane, + location, + )?; + let scaled_a = emit_binary( + session, + &mut operations, + "arith.mulf", + a, + scale_a.clone(), + location, + )?; + let product = emit_binary( + session, + &mut operations, + "arith.mulf", + scaled_a, + b, + location, + )?; + let scaled_b = emit_binary( + session, + &mut operations, + "arith.mulf", + product, + scale_b.clone(), + location, + )?; + let is_last = group + 1 == SCALE_BYTES_PER_TILE && lane_in_group == 15; + let result = if is_last { + final_result.clone() + } else { + fresh(session, MlirType::Float(MlirFloatType::F32)).0 + }; + let result_use = MlirValueUse { + id: result.id, + ty: result.ty.clone(), + }; + operations.push(binary( + "arith.addf", + result, + accumulator, + scaled_b, + location, + )?); + accumulator = result_use; + } + } + Ok(operations) +} + +/// Compute an unsigned ceiling quotient without overflowing an add-before-divide expression. +fn emit_ceil_div( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: MlirValueUse, + divisor: u64, + location: &MlirLocation, +) -> Result { + let divisor = emit_i64_constant(session, operations, divisor, location)?; + let quotient = emit_binary( + session, + operations, + "arith.divui", + value.clone(), + divisor.clone(), + location, + )?; + let remainder = emit_binary(session, operations, "arith.remui", value, divisor, location)?; + let zero = emit_i64_constant(session, operations, 0, location)?; + let (nonzero_result, nonzero) = fresh(session, MlirType::Integer(1)); + let mut compare = operation( + "arith.cmpi", + vec![nonzero_result], + vec![remainder, zero], + location, + )?; + compare.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(64), + }, + ); + operations.push(compare); + let (round_up_result, round_up) = fresh(session, MlirType::Integer(64)); + operations.push(operation( + "arith.extui", + vec![round_up_result], + vec![nonzero], + location, + )?); + emit_binary( + session, + operations, + "arith.addi", + quotient, + round_up, + location, + ) +} + +/// Append an i64 constant while preserving the unsigned input's bit pattern. +fn emit_i64_constant( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: u64, + location: &MlirLocation, +) -> Result { + let (result, use_) = fresh(session, MlirType::Integer(64)); + operations.push(integer_constant(result, u64_bits(value), location)?); + Ok(use_) +} + +/// Append a two-operand operation whose result has the left operand's type. +fn emit_binary( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + name: &str, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + if left.ty != right.ty { + return Err(format!( + "{name} operands have different types: {:?} and {:?}", + left.ty, right.ty + )); + } + let (result, use_) = fresh(session, left.ty.clone()); + operations.push(binary(name, result, left, right, location)?); + Ok(use_) +} + +/// Extract one compile-time vector lane into a fresh scalar result. +fn emit_vector_extract( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + vector: MlirValueUse, + position: u64, + location: &MlirLocation, +) -> Result { + let MlirType::Vector { shape, element } = &vector.ty else { + return Err(format!( + "vector.extract expected a vector, got {:?}", + vector.ty + )); + }; + if shape.as_slice() != [shape[0]] || position >= shape[0] { + return Err(format!( + "vector.extract position {position} is outside {:?}", + shape + )); + } + let (result, use_) = fresh(session, (**element).clone()); + let mut extract = operation("vector.extract", vec![result], vec![vector], location)?; + extract.properties.insert( + "static_position".into(), + MlirAttribute::DenseI64Array(vec![position as i64]), + ); + operations.push(extract); + Ok(use_) +} + +/// Insert ordered runtime fields into the requested LLVM struct result. +fn build_aggregate( + session: &mut TranslationSession<'_>, + final_result: MlirResult, + fields: Vec, + location: &MlirLocation, +) -> Result, String> { + if fields.is_empty() { + return Err("cannot build an empty GEMV aggregate".into()); + } + let (poison_result, mut aggregate) = fresh(session, final_result.ty.clone()); + let poison = operation("llvm.mlir.poison", vec![poison_result], vec![], location)?; + let mut operations = vec![poison]; + let field_count = fields.len(); + for (position, field) in fields.into_iter().enumerate() { + let result = if position + 1 == field_count { + final_result.clone() + } else { + fresh(session, final_result.ty.clone()).0 + }; + let next = MlirValueUse { + id: result.id, + ty: result.ty.clone(), + }; + operations.push(insert_value( + result, + aggregate, + field, + position as i64, + location, + )?); + aggregate = next; + } + Ok(operations) +} + +/// Extract one LLVM aggregate field into a fresh result of its known type. +fn extract_value( + result: MlirResult, + aggregate: MlirValueUse, + position: i64, + location: &MlirLocation, +) -> Result { + let mut operation = operation("llvm.extractvalue", vec![result], vec![aggregate], location)?; + operation.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![position]), + ); + Ok(operation) +} + +/// Insert a field at its physical LLVM aggregate index. +fn insert_value( + result: MlirResult, + aggregate: MlirValueUse, + field: MlirValueUse, + position: i64, + location: &MlirLocation, +) -> Result { + let mut operation = operation( + "llvm.insertvalue", + vec![result], + vec![aggregate, field], + location, + )?; + operation.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![position]), + ); + Ok(operation) +} + +/// LLVM struct transporting a data pointer, two extents, and a row stride. +fn tensor_carrier_type() -> Result { + MlirType::dialect("!llvm.struct<(!llvm.ptr, i64, i64, i64)>") +} + +/// LLVM struct transporting value/scale pointers and runtime scaled-view coordinates. +fn scaled_carrier_type() -> Result { + MlirType::dialect( + "!llvm.struct<(!llvm.struct<(!llvm.ptr, i64, i64, i64)>, !llvm.struct<(!llvm.ptr, i64, i64, i64)>, i64, i64, i64)>", + ) +} + +/// LLVM carrier holding two 32-lane FP32 value vectors and four FP32 scales. +fn fragment_carrier_type() -> Result { + MlirType::dialect("!llvm.struct<(vector<32xf32>, vector<32xf32>, vector<4xf32>)>") +} + +/// Return the fixed vector shape and element format for the selected narrow load. +fn narrow_vector_type(kind: NarrowLoad) -> MlirType { + match kind { + NarrowLoad::E2M1 => vector_type( + VALUE_LANES_PER_LOAD, + MlirType::Float(MlirFloatType::Other("f4E2M1FN".into())), + ), + NarrowLoad::UE8M0 => vector_type( + SCALE_BYTES_PER_TILE, + MlirType::Float(MlirFloatType::Other("f8E8M0FNU".into())), + ), + NarrowLoad::E4M3 => vector_type( + SCALE_BYTES_PER_TILE, + MlirType::Float(MlirFloatType::Other("f8E4M3FN".into())), + ), + } +} + +/// Wrap CUTLASS's narrow floating-point type spelling without widening it. +fn narrow_float(spelling: &str) -> MlirType { + match spelling { + "f16" => MlirType::Float(MlirFloatType::F16), + "bf16" => MlirType::Float(MlirFloatType::BF16), + "f4E2M1FN" | "f8E8M0FNU" | "f8E4M3FN" => { + MlirType::Float(MlirFloatType::Other(spelling.into())) + } + _ => unreachable!("fixed GEMV conversion intermediate"), + } +} + +/// Build a one-dimensional vector with the requested lane count and element type. +fn vector_type(size: u64, element: MlirType) -> MlirType { + MlirType::Vector { + shape: vec![size], + element: Box::new(element), + } +} + +/// Recover the source tensor-view contract before its runtime carrier is translated. +fn tensor_view(ctx: &Context, value: pliron::value::Value) -> Result { + let ty = value.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected cute.tensor_view value".to_owned()) +} + +/// Recover the source scaled-view contract before its runtime carrier is translated. +fn scaled_view(ctx: &Context, value: pliron::value::Value) -> Result { + let ty = value.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected cute.scaled_view value".to_owned()) +} + +/// Recover the source fragment contract before its runtime carrier is translated. +fn fragment(ctx: &Context, value: pliron::value::Value) -> Result { + let ty = value.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected cute.fragment value".to_owned()) +} + +/// Reject source tensor views outside the reviewed packed-value and scale formats. +fn require_gemv_tensor_view(ctx: &Context, view: &CuteTensorViewType) -> Result<(), String> { + view.verify(ctx).map_err(stringify)?; + if view.space != CuteTensorAddressSpaceAttr::Gmem + || view.access != CuteTensorAccessAttr::ReadOnly + || view.alignment.0 != 1 + || !matches!(view.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) + { + return Err(format!( + "GEMV tensor must be a byte-aligned read-only Mkl/Nkl gmem view, got {:?}", + view + )); + } + match (view.format, view.layout) { + (CuteTensorFormatAttr::E2M1, CuteTensorLayoutAttr::KMajor) + | ( + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3, + CuteTensorLayoutAttr::BlockScaleKMajor(16), + ) => Ok(()), + other => Err(format!( + "GEMV tensor needs E2M1/KMajor or UE8M0/E4M3 BlockScaleKMajor<16>, got {other:?}" + )), + } +} + +/// Require exact result/operand counts and no leftover attributes, regions, or successors. +fn expect_shape( + name: &str, + input: &OperationInput, + results: usize, + operands: usize, +) -> Result<(), String> { + if input.results.len() != results || input.operands.len() != operands { + return Err(format!( + "{name} expected {results} results and {operands} operands, got {} and {}", + input.results.len(), + input.operands.len() + )); + } + if !input.successors.is_empty() || !input.regions.is_empty() || !input.attributes.is_empty() { + return Err(format!( + "{name} unexpectedly retained successors, regions, or source-only attributes" + )); + } + Ok(()) +} + +/// Build a target operation with the supplied operands, results, and source location. +fn operation( + name: &str, + results: Vec, + operands: Vec, + location: &MlirLocation, +) -> Result { + let mut operation = MlirOperation::new(name)?; + operation.results = results; + operation.operands = operands; + operation.location = location.clone(); + Ok(operation) +} + +/// Allocate one SSA identifier and return matching result and operand handles. +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +/// Build an arith.constant using the supplied result's integer type. +fn integer_constant( + result: MlirResult, + value: i128, + location: &MlirLocation, +) -> Result { + let mut operation = operation("arith.constant", vec![result.clone()], vec![], location)?; + operation.properties.insert( + "value".into(), + MlirAttribute::Integer { + value, + ty: result.ty, + }, + ); + Ok(operation) +} + +/// Build a two-operand target operation using the supplied result. +fn binary( + name: &str, + result: MlirResult, + left: MlirValueUse, + right: MlirValueUse, + location: &MlirLocation, +) -> Result { + operation(name, vec![result], vec![left, right], location) +} + +/// Spell an opaque LLVM pointer in the requested numeric address space. +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +/// Preserve all 64 bits while spelling an unsigned value as a signed MLIR literal. +fn u64_bits(value: u64) -> i128 { + i64::from_ne_bytes(value.to_ne_bytes()) as i128 +} + +/// Preserve a source diagnostic as the string expected by translation recipes. +fn stringify(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/mod.rs new file mode 100644 index 0000000000..4d1af87483 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tensor/mod.rs @@ -0,0 +1,9 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor-view construction, elementwise copies, and scaled GEMV mappings. + +pub(in crate::cute) mod elementwise; +pub(in crate::cute) mod gemv; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs new file mode 100644 index 0000000000..960ac6a229 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/cooperative.rs @@ -0,0 +1,1111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Independent address evaluation for cooperative-copy and matrix-load mappings. + +use crate::cute::{ + copy::emit::integer, + mma::ldmatrix::{matrix_address_plan, matrix_offset_fits_i32}, +}; + +#[cfg(test)] +mod tests { + //! Evaluate exported copy and matrix-load address arithmetic against independent layout expectations. + + use crate::render_module; + use crate::{CutlassFullCuteMlir22, MlirConsumerProfile}; + use cute_layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; + use dialect_cute::{ + attributes::CuteMatrixRoleAttr, + ops::{CuteCopyG2SOp, CuteLdmatrixOp}, + }; + use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType, address_space}, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::TypeHandle, + }; + + /// Construct the canonical swizzled shared tile used by cooperative-copy fixtures. + fn shared_layout() -> ComposedLayout { + ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + "(64,64):(64,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap() + } + + /// Append a typed source operation to the fixture block. + fn append(ctx: &Context, block: Ptr, op: O) { + op.get_operation().insert_at_back(block, ctx); + } + + /// Build the default copy or matrix-load fixture for the selected matrix role. + fn module(ctx: &mut Context, copy: bool, role: CuteMatrixRoleAttr) -> ModuleOp { + copy_module(ctx, copy, role, 64) + } + + /// Build a copy or matrix-load fixture with the selected row extent and default layout. + fn copy_module( + ctx: &mut Context, + copy: bool, + role: CuteMatrixRoleAttr, + rows: usize, + ) -> ModuleOp { + module_with_layout(ctx, copy, role, rows, shared_layout()) + } + + /// Build a copy or matrix-load fixture using the supplied shared layout. + fn module_with_layout( + ctx: &mut Context, + copy: bool, + role: CuteMatrixRoleAttr, + rows: usize, + layout: ComposedLayout, + ) -> ModuleOp { + module_with_copy_plan(ctx, copy, role, rows, layout, None) + } + + /// Atom bytes, thread/value/tile layouts, and composed shared layout for one copy fixture. + type CopyPlanFixture = (u32, Layout, Layout, Layout, ComposedLayout); + + /// Construct the typed test function for an optional explicit cooperative-copy plan. + fn module_with_copy_plan( + ctx: &mut Context, + copy: bool, + role: CuteMatrixRoleAttr, + rows: usize, + layout: ComposedLayout, + copy_plan: Option, + ) -> ModuleOp { + module_with_copy_storage(ctx, copy, role, rows, layout, copy_plan, None) + } + + /// Select integer storage while keeping pointer pointees and copy element identity aligned. + fn module_with_copy_storage( + ctx: &mut Context, + copy: bool, + role: CuteMatrixRoleAttr, + rows: usize, + layout: ComposedLayout, + copy_plan: Option, + integer_storage: Option<(u32, Signedness)>, + ) -> ModuleOp { + dialect_mir::register(ctx); + dialect_cute::register(ctx); + let module = ModuleOp::new(ctx, Identifier::try_from("attention_copy").unwrap()); + let element: TypeHandle = if let Some((width, signedness)) = integer_storage { + IntegerType::get(ctx, width, signedness).into() + } else { + MirFP16Type::get(ctx).into() + }; + let i32_: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let i64_: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let gmem: TypeHandle = MirPtrType::get(ctx, element, false, address_space::GENERIC).into(); + let smem: TypeHandle = MirPtrType::get(ctx, element, true, address_space::GENERIC).into(); + let slot: TypeHandle = MirPtrType::get(ctx, i32_, true, address_space::GENERIC).into(); + let types = if copy { + vec![gmem, i64_, i64_, i64_, i64_, i64_, smem, i64_, i32_] + } else { + vec![smem, i64_, i64_, i32_, slot] + }; + let ty = FunctionType::get(ctx, types.clone(), vec![]); + let op = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = MirFuncOp::new(ctx, op, TypeAttr::new(ty.into())); + func.set_symbol_name(ctx, Identifier::try_from("attention_copy").unwrap()); + module.append_operation(ctx, op, 0); + let block = BasicBlock::new(ctx, None, types); + block.insert_at_back(op.deref(ctx).get_region(0), ctx); + let args = block.deref(ctx).arguments().collect::>(); + if copy { + let (atom, thread, value, tile, shared) = copy_plan.unwrap_or_else(|| { + ( + 16, + "(16,8):(8,1)".parse().unwrap(), + format!("({},8):(8,1)", rows / 16).parse().unwrap(), + format!("({rows},64):(64,1)").parse().unwrap(), + ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + format!("({rows},64):(64,1)").parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ) + }); + let copy = CuteCopyG2SOp::new( + ctx, + args.try_into().unwrap(), + atom, + thread, + value, + tile, + shared, + 8, + element, + ); + append(ctx, block, copy); + } else { + let load = CuteLdmatrixOp::new(ctx, args.try_into().unwrap(), role, layout, element); + append(ctx, block, load); + } + let ret = Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + append(ctx, block, MirReturnOp::new(ret)); + module + } + + /// Check the canonical copy emits four atoms with runtime source-byte counts for edge zero filling. + #[test] + fn attention_copy_repeats_four_atoms_with_dynamic_zero_fill() { + let mut ctx = Context::new(); + let source = module(&mut ctx, true, CuteMatrixRoleAttr::A); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let text = render_module(&target).unwrap(); + assert_eq!( + text.matches("\"nvvm.cp.async.shared.global\"").count(), + 4, + "{text}" + ); + assert!(text.contains("#nvvm"), "{text}"); + assert!(text.contains("arith.select"), "{text}"); + assert!(text.contains("arith.xori"), "{text}"); + assert!(!text.contains("cp.async.commit"), "{text}"); + assert!(!text.contains("cp.async.wait"), "{text}"); + assert!(!text.contains("cute.copy_g2s"), "{text}"); + } + + /// Check A/B matrix roles select the reviewed normal or transposed ldmatrix instruction. + #[test] + fn attention_matrix_roles_select_normal_and_transposed_loads() { + for (role, number, layout) in [ + (CuteMatrixRoleAttr::A, 4, "row"), + (CuteMatrixRoleAttr::B, 2, "col"), + (CuteMatrixRoleAttr::BPair, 4, "col"), + (CuteMatrixRoleAttr::BPairKMajor, 4, "row"), + ] { + let mut ctx = Context::new(); + let source = module(&mut ctx, false, role); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let text = render_module(&target).unwrap(); + assert!(text.contains(&format!("num = {number} : i32")), "{text}"); + assert!( + text.contains(&format!("#nvvm.mma_layout<{layout}>")), + "{text}" + ); + assert_eq!(text.matches("\"llvm.store\"").count(), number, "{text}"); + assert!(!text.contains("cute.ldmatrix"), "{text}"); + } + } + + /// Find the device function's block through the executable GPU envelope. + fn function_block(op: &crate::MlirOperation) -> &crate::MlirBlock { + if op.name == "func.func" { + return &op.regions[0].blocks[0]; + } + op.regions + .iter() + .flat_map(|r| &r.blocks) + .flat_map(|b| &b.operations) + .find_map(|op| { + /// Recognize a nested target subtree containing a translated function. + fn has_function(op: &crate::MlirOperation) -> bool { + op.name == "func.func" + || op + .regions + .iter() + .flat_map(|r| &r.blocks) + .flat_map(|b| &b.operations) + .any(has_function) + } + has_function(op).then(|| function_block(op)) + }) + .expect("translated function") + } + + // Execute emitted scalar MLIR, observing the addresses and source byte + // counts handed to hardware. This checks the actual lowering, not a second + // copy of its address formula. No kernel or GPU runtime is involved. + /// Interpret exported integer address arithmetic and collect each memory instruction's runtime arguments. + fn memory_arguments(block: &crate::MlirBlock, args: &[u64]) -> Vec> { + use crate::{MlirAttribute as A, MlirFloatType, MlirType as T}; + let mut values = block + .arguments + .iter() + .zip(args) + .map(|(arg, value)| (arg.id, *value)) + .collect::>(); + let mut memory = vec![]; + for op in &block.operations { + let a = op + .operands + .iter() + .map(|v| values[&v.id]) + .collect::>(); + let property = |name: &str| match op.properties[name] { + A::Integer { value, .. } => value as u64, + _ => panic!("integer property {name}"), + }; + let result = match op.name.as_str() { + "arith.constant" => property("value"), + "arith.extui" | "llvm.addrspacecast" => a[0], + "arith.trunci" => u64::from(a[0] as u32), + "arith.addi" => a[0].wrapping_add(a[1]), + "arith.subi" => a[0].wrapping_sub(a[1]), + "arith.muli" => a[0].wrapping_mul(a[1]), + "arith.divui" => a[0] / a[1], + "arith.remui" => a[0] % a[1], + "arith.minui" => a[0].min(a[1]), + "arith.andi" => a[0] & a[1], + "arith.xori" => a[0] ^ a[1], + "arith.shrui" => a[0] >> a[1], + "arith.shli" => a[0] << a[1], + "arith.cmpi" => match property("predicate") { + 6 => u64::from(a[0] < a[1]), + 7 => u64::from(a[0] <= a[1]), + p => panic!("unexpected predicate {p}"), + }, + "arith.select" => { + if a[0] != 0 { + a[1] + } else { + a[2] + } + } + "llvm.getelementptr" => { + let width = match op.properties["elem_type"] { + A::Type(T::Integer(8)) => 1, + A::Type(T::Integer(16)) => 2, + A::Type(T::Integer(32)) => 4, + A::Type(T::Float(MlirFloatType::F16)) => 2, + ref ty => panic!("unexpected GEP element {ty:?}"), + }; + a[0].wrapping_add(a[1].wrapping_mul(width)) + } + "nvvm.cp.async.shared.global" => { + memory.push(a); + continue; + } + "nvvm.ldmatrix" => { + memory.push(a); + break; + } + "func.return" => break, + name => panic!("unexpected operation {name}"), + }; + // Arithmetic has the result's bit width, including wraparound. + // Pointers retain all 64 bits even when their offsets were i32. + let result = match op.results[0].ty { + T::Integer(width) if width < 64 => result & ((1u64 << width) - 1), + _ => result, + }; + values.insert(op.results[0].id, result); + } + memory + } + + /// Integer copy carriers preserve D128 byte addresses and every edge atom's source-byte count. + #[test] + fn attention_copy_integer16_preserves_storage_and_zero_fills_d128_edges() { + use crate::{MlirAttribute as A, MlirType as T}; + for signedness in [Signedness::Unsigned, Signedness::Signed] { + for atom in [4, 8, 16] { + let mut ctx = Context::new(); + let tile: Layout = "(8,128):(128,1)".parse().unwrap(); + let shared = ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + tile.clone(), + OffsetUnit::Elements, + ) + .unwrap(); + let source = module_with_copy_storage( + &mut ctx, + true, + CuteMatrixRoleAttr::A, + 8, + shared_layout(), + Some(( + atom, + "(8,16):(16,1)".parse().unwrap(), + "(1,8):(8,1)".parse().unwrap(), + tile, + shared, + )), + Some((16, signedness)), + ); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + assert!(block.operations.iter().any(|op| { + op.name == "llvm.getelementptr" + && op.properties["elem_type"] == A::Type(T::Integer(16)) + })); + for (rows, columns, tile_row, tile_column) in [ + (64, 128, 0, 0), + (19, 125, 2, 0), + (9, 131, 1, 1), + (0, 0, 0, 0), + (64, 128, u64::MAX, u64::MAX), + ] { + for thread in 0..128 { + let copies = memory_arguments( + block, + &[ + 0x10000, + rows, + columns, + 256, + tile_row, + tile_column, + 0x20000, + 1024, + thread, + ], + ); + let per_atom = u64::from(atom / 2); + assert_eq!(copies.len(), (8 / per_atom) as usize); + for (index, copy) in copies.iter().enumerate() { + let row = thread / 16; + let column = (thread % 16) * 8 + index as u64 * per_atom; + let bytes = (row * 128 + column) * 2; + let destination = 0x20000 + (bytes ^ ((bytes & 0x380) >> 3)); + let global_row = u128::from(tile_row) * 8 + u128::from(row); + let global_col = u128::from(tile_column) * 128 + u128::from(column); + let valid = + global_row < u128::from(rows) && global_col < u128::from(columns); + let source = if valid { + 0x10000 + (global_row as u64 * 256 + global_col as u64) * 2 + } else { + 0x10000 + }; + let valid_bytes = if valid { + (columns - global_col as u64).min(per_atom) * 2 + } else { + 0 + }; + assert_eq!( + *copy, + [destination, source, valid_bytes], + "signedness={signedness:?}, atom={atom}, thread={thread}, shape=({rows},{columns}), tile=({tile_row},{tile_column})" + ); + } + } + } + } + } + } + + /// The physical two-byte address calculation must reject other integer widths. + #[test] + fn attention_copy_rejects_other_integer_storage_widths() { + for width in [8, 32, 64] { + let mut ctx = Context::new(); + let source = module_with_copy_storage( + &mut ctx, + true, + CuteMatrixRoleAttr::A, + 16, + shared_layout(), + Some(( + if width == 8 { 4 } else { 16 }, + "(16,8):(8,1)".parse().unwrap(), + "(1,8):(8,1)".parse().unwrap(), + "(16,64):(64,1)".parse().unwrap(), + ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + "(16,64):(64,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + )), + Some((width, Signedness::Unsigned)), + ); + let error = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap_err(); + assert!(error.to_string().contains("16-bit"), "{error}"); + } + } + + /// Compare emitted destinations and valid source-byte counts with the verified thread/value assignment. + #[test] + fn attention_copy_matches_thread_layout_and_zero_fills_edge_bytes() { + let mut ctx = Context::new(); + let source = copy_module(&mut ctx, true, CuteMatrixRoleAttr::A, 16); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + for (rows, columns, tile_row, tile_column) in [ + (64, 64, 0, 0), + (19, 61, 1, 0), + (19, 67, 1, 1), + (0, 0, 0, 0), + (19, 61, 2, 1), + (64, 64, u64::MAX, u64::MAX), + ] { + for thread in 0..128 { + let copies = memory_arguments( + block, + &[ + 0x10000, + rows, + columns, + 128, + tile_row, + tile_column, + 0x20000, + 1024, + thread, + ], + ); + assert_eq!(copies.len(), 1); + let row = thread / 8; + let column = thread % 8 * 8; + let plain_bytes = (row * 64 + column) * 2; + let expected_shared = 0x20000 + (plain_bytes ^ ((plain_bytes & 0x380) >> 3)); + let global_row = u128::from(tile_row) * 16 + u128::from(row); + let global_column = u128::from(tile_column) * 64 + u128::from(column); + let in_bounds = + global_row < u128::from(rows) && global_column < u128::from(columns); + let expected_bytes = if in_bounds { + (columns - global_column as u64).min(8) * 2 + } else { + 0 + }; + let expected_global = if in_bounds { + 0x10000 + (global_row as u64 * 128 + global_column as u64) * 2 + } else { + 0x10000 + }; + assert_eq!( + copies[0], + vec![expected_shared, expected_global, expected_bytes], + "thread={thread}, rows={rows}, columns={columns}, tile=({tile_row},{tile_column})" + ); + } + } + } + + /// Check bounded thread IDs cover every copy atom and oversized tile coordinates still zero-fill safely. + #[test] + fn attention_copy_bounded_threads_preserve_all_atoms_and_u64_bounds() { + // A mask instead of remainder corrupts the six-thread fixtures; a + // warp-sized bound corrupts IDs32..1023. Each valid thread must keep + // its three atoms, including partial source sizes and zero-fill. + for (thread_rows, thread_columns, column_first) in [ + (1u64, 1u64, false), + (6, 1, false), + (8, 1, false), + (16, 8, false), + (1024, 1, false), + (2, 3, true), + ] { + for atom_bytes in [4u32, 8, 16] { + let values_per_atom = u64::from(atom_bytes) / 2; + let values_per_thread = 3 * values_per_atom; + let tile_columns = thread_columns * values_per_thread; + let thread: Layout = if column_first { + format!("({thread_rows},{thread_columns}):(1,{thread_rows})") + } else { + format!("({thread_rows},{thread_columns}):({thread_columns},1)") + } + .parse() + .unwrap(); + let value: Layout = format!("(1,{values_per_thread}):(0,1)").parse().unwrap(); + let tile: Layout = format!("({thread_rows},{tile_columns}):({tile_columns},1)") + .parse() + .unwrap(); + let shared = ComposedLayout::new( + Swizzle::new(3, 3, 3), + 8, + tile.clone(), + OffsetUnit::Elements, + ) + .unwrap(); + let shared_bytes = shared.to_byte_offsets(2).unwrap(); + let mut ctx = Context::new(); + let source = module_with_copy_plan( + &mut ctx, + true, + CuteMatrixRoleAttr::A, + 0, + shared_layout(), + Some((atom_bytes, thread, value, tile, shared)), + ); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + let cp_ops = block + .operations + .iter() + .filter(|op| op.name == "nvvm.cp.async.shared.global") + .collect::>(); + assert_eq!(cp_ops.len(), 3); + for op in cp_ops { + assert_eq!( + op.properties["size"], + super::integer(i128::from(atom_bytes), 32) + ); + } + for (rows, columns, tile_row, tile_column) in [ + (thread_rows, tile_columns, 0, 0), + (thread_rows - 1, tile_columns - 1, 0, 0), + (thread_rows + 1, tile_columns + values_per_atom - 1, 1, 1), + (0, 0, 0, 0), + (thread_rows, tile_columns, u64::MAX, u64::MAX), + ( + (1u64 << 32) + thread_rows, + tile_columns, + (1u64 << 32) / thread_rows, + 0, + ), + ( + thread_rows, + (1u64 << 32) + tile_columns, + 0, + (1u64 << 32) / tile_columns, + ), + ] { + let leading_dim = columns.div_ceil(8) * 8; + for thread in 0..thread_rows * thread_columns { + // Independent inversion of the two fixture thread maps. + let (local_row, column_group) = if column_first { + (thread % thread_rows, thread / thread_rows) + } else { + (thread / thread_columns, thread % thread_columns) + }; + let copies = memory_arguments( + block, + &[ + 0x1234_0001_0000, + rows, + columns, + leading_dim, + tile_row, + tile_column, + 0x5678_0002_0000, + 1 << 20, + thread, + ], + ); + assert_eq!(copies.len(), 3); + for (atom, copy) in copies.iter().enumerate() { + let local_column = + column_group * values_per_thread + atom as u64 * values_per_atom; + let destination = shared_bytes + .checked_call(&cute_layout::IntTuple::Leaf( + (local_row + thread_rows * local_column) as i64, + )) + .unwrap() as u64; + let global_row = u128::from(tile_row) * u128::from(thread_rows) + + u128::from(local_row); + let global_column = u128::from(tile_column) * u128::from(tile_columns) + + u128::from(local_column); + let valid = global_row < u128::from(rows) + && global_column < u128::from(columns); + let valid_bytes = if valid { + (columns - global_column as u64).min(values_per_atom) * 2 + } else { + 0 + }; + let source_offset = if valid { + (global_row as u64 * leading_dim + global_column as u64) * 2 + } else { + 0 + }; + assert_eq!( + copy, + &[ + 0x5678_0002_0000 + destination, + 0x1234_0001_0000 + source_offset, + valid_bytes + ], + "threads=({thread_rows},{thread_columns}), column_first={column_first}, atom_bytes={atom_bytes}, thread={thread}, atom={atom}, shape=({rows},{columns}), tile=({tile_row},{tile_column})" + ); + } + } + } + } + } + } + + /// Compare each emitted lane address with the independent byte-domain shared-layout swizzle. + #[test] + fn attention_matrix_lane_addresses_follow_byte_swizzle() { + for role in [ + CuteMatrixRoleAttr::A, + CuteMatrixRoleAttr::B, + CuteMatrixRoleAttr::BPair, + ] { + let mut ctx = Context::new(); + let source = module(&mut ctx, false, role); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + for lane in 0..32 { + let addresses = memory_arguments(block, &[0x20000, 1, 1, lane, 0x30000]); + let row = 16 + lane % 16; + let column = if role != CuteMatrixRoleAttr::B { + 16 + lane / 16 * 8 + } else { + 8 + }; + let plain_bytes = (row * 64 + column) * 2; + assert_eq!( + addresses[0], + [0x20000 + (plain_bytes ^ ((plain_bytes & 0x380) >> 3))] + ); + assert_eq!(addresses[0][0] % 16, 0); + } + } + } + + /// Check a proved shared-address range permits narrow arithmetic without changing byte offsets. + #[test] + fn attention_matrix_uses_i32_for_proven_shared_addresses() { + let mut ctx = Context::new(); + let source = module(&mut ctx, false, CuteMatrixRoleAttr::A); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + for op in block + .operations + .iter() + .take_while(|op| op.name != "nvvm.ldmatrix") + .filter(|op| { + matches!( + op.name.as_str(), + "arith.addi" + | "arith.muli" + | "arith.divui" + | "arith.remui" + | "arith.andi" + | "arith.xori" + ) + }) + { + assert_eq!( + op.results[0].ty, + crate::MlirType::Integer(32), + "{}", + op.name + ); + } + let source_gep = block + .operations + .iter() + .find(|op| op.name == "llvm.getelementptr") + .unwrap(); + assert_eq!(source_gep.operands[1].ty, crate::MlirType::Integer(64)); + } + + /// Check proved carry-free window/lane splitting preserves the full matrix-load address calculation. + #[test] + fn attention_matrix_factors_window_and_lane_pointer_offsets() { + for role in [ + CuteMatrixRoleAttr::A, + CuteMatrixRoleAttr::B, + CuteMatrixRoleAttr::BPair, + CuteMatrixRoleAttr::BPairKMajor, + ] { + let mut ctx = Context::new(); + let source = module(&mut ctx, false, role); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + let geps = block + .operations + .iter() + .take_while(|op| op.name != "nvvm.ldmatrix") + .filter(|op| op.name == "llvm.getelementptr") + .collect::>(); + assert_eq!(geps.len(), 2, "keep the uniform high offset additive"); + assert_eq!(geps[1].operands[0].id, geps[0].results[0].id); + let mut lane_dependent = std::collections::BTreeSet::from([block.arguments[3].id]); + for op in block + .operations + .iter() + .take_while(|op| op.name != "nvvm.ldmatrix") + { + if op.operands.iter().any(|v| lane_dependent.contains(&v.id)) { + lane_dependent.extend(op.results.iter().map(|v| v.id)); + } + } + assert!(!lane_dependent.contains(&geps[0].operands[1].id)); + assert!(lane_dependent.contains(&geps[1].operands[1].id)); + } + } + + /// Reject narrow-address proofs whose intermediate terms overflow despite cancellation in the final sum. + #[test] + fn attention_matrix_i32_proof_checks_each_bound_before_narrowing() { + let max = i64::from(i32::MAX); + for (outer, offset, expected) in [ + ("(16,16):(32,2)".to_owned(), max - 510, true), + ("(16,16):(32,2)".to_owned(), max - 509, false), + ("(16,16):(32,2)".to_owned(), -1, false), + ("(16,16):(-32,2)".to_owned(), 480, false), + (format!("(1,1):({max},0)"), 0, true), + (format!("(1,1):({},0)", max + 1), 0, false), + (format!("({max},1):(0,0)"), 0, true), + (format!("({},1):(0,0)", max + 1), 0, false), + ("(65536,65536):(0,0)".to_owned(), 0, false), + (format!("({},2):(0,0)", i64::MAX), 0, false), + ("(0,16):(32,2)".to_owned(), 0, false), + ("(-16,16):(32,2)".to_owned(), 0, false), + ("16:2".to_owned(), 0, false), + ("((4,4),(8,2)):((32,128),(2,16))".to_owned(), 16, true), + ] { + let layout = ComposedLayout::new( + Swizzle::IDENTITY, + offset, + outer.parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!( + super::matrix_offset_fits_i32(&layout), + expected, + "{layout:?}" + ); + } + for (swizzle, expected) in [ + (Swizzle::new(1, 29, 1), true), + (Swizzle::new(1, 29, -1), true), + (Swizzle::new(1, 30, 1), false), + (Swizzle::new(1, 30, -1), false), + (Swizzle::new(0, 32, 0), true), + ] { + let layout = ComposedLayout::new( + swizzle, + 0, + "(16,16):(32,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!( + super::matrix_offset_fits_i32(&layout), + expected, + "{layout:?}" + ); + } + assert!(!super::matrix_offset_fits_i32(&shared_layout())); + } + + /// Enumerate legal matrix windows and lanes and compare emitted addresses with checked layout evaluation. + #[test] + fn attention_matrix_addresses_match_checked_layout_for_every_window_and_lane() { + use cute_layout::IntTuple; + let layouts = [ + shared_layout(), + ComposedLayout::new( + Swizzle::IDENTITY, + 16, + "(32,32):(32,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + -16, + "(32,32):(32,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + 1024, + "(32,32):(-32,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(2, 3, -2), + 32, + "((4,8),(8,4)):((32,128),(1,8))".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + 64, + "((4,8),(8,4)):((32,128),(1,-8))".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + 1i64 << 32, + "(32,32):(64,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(1, 30, 1), + 0, + "(32,32):(64,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(1, 30, -1), + 0, + "(32,32):(64,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, 3), + 8192, + "((8,8),64):((128,1024),2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, 3), + 8192, + "((8,8),64):((1024,128),2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, -3), + 0, + "(64,64):(128,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, 3), + 16, + "(64,64):(128,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, 3), + 0, + "((3,16),64):((128,512),2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ComposedLayout::new( + Swizzle::new(3, 4, 3), + i64::from(i32::MAX) - 8191, + "(64,64):(128,2)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ]; + for (layout_index, layout) in layouts.into_iter().enumerate() { + let bytes = layout.to_byte_offsets(2).unwrap(); + let modes = layout.outer().modes(); + let rows = modes[0].checked_size().unwrap() as u64; + let columns = modes[1].checked_size().unwrap() as u64; + for role in [ + CuteMatrixRoleAttr::A, + CuteMatrixRoleAttr::B, + CuteMatrixRoleAttr::BPair, + CuteMatrixRoleAttr::BPairKMajor, + ] { + let mut ctx = Context::new(); + let source = module_with_layout(&mut ctx, false, role, 0, layout.clone()); + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &source) + .unwrap(); + let block = function_block(&target.root); + let source_geps = block + .operations + .iter() + .take_while(|op| op.name != "nvvm.ldmatrix") + .filter(|op| op.name == "llvm.getelementptr") + .count(); + let expected_geps = if [0, 9, 10, 11, 14].contains(&layout_index) { + 2 + } else { + 1 + }; + assert_eq!( + source_geps, expected_geps, + "factoring selection: {layout:?}, {role:?}" + ); + let expected_width = if [0, 1, 4].contains(&layout_index) || layout_index >= 9 { + 32 + } else { + 64 + }; + for op in block + .operations + .iter() + .take_while(|op| op.name != "nvvm.ldmatrix") + .filter(|op| { + matches!( + op.name.as_str(), + "arith.addi" + | "arith.muli" + | "arith.divui" + | "arith.remui" + | "arith.andi" + | "arith.xori" + ) + }) + { + assert_eq!( + op.results[0].ty, + crate::MlirType::Integer(expected_width), + "layout={layout:?}, op={}", + op.name + ); + } + let tile_columns = if role == CuteMatrixRoleAttr::B { 8 } else { 16 }; + // Hardware matrix order, independent of the emitter's arithmetic. + let quadrants = if role == CuteMatrixRoleAttr::BPairKMajor { + [(0, 0), (0, 8), (8, 0), (8, 8)] + } else if role == CuteMatrixRoleAttr::B { + [(0, 0), (8, 0), (0, 0), (8, 0)] + } else { + [(0, 0), (8, 0), (0, 8), (8, 8)] + }; + for window_row in 0..rows / 16 { + for window_column in 0..columns / tile_columns { + for lane in 0..32 { + let (row_half, column_half) = quadrants[lane as usize / 8]; + let row = window_row * 16 + row_half + lane % 8; + let column = window_column * tile_columns + column_half; + let offset = bytes + .checked_call(&IntTuple::Leaf((row + rows * column) as i64)) + .unwrap(); + let base = 0x1234_0002_0000u64; + let actual = memory_arguments( + block, + &[base, window_row, window_column, lane, 0x30000], + ); + assert_eq!( + actual, + [vec![base.wrapping_add(offset as u64)]], + "layout={layout:?}, role={role:?}, window=({window_row},{window_column}), lane={lane}" + ); + } + } + } + } + } + } + + /// Reject factoring when proof work exceeds its cap or the selected matrix window is incomplete. + #[test] + fn attention_matrix_factoring_caps_proof_and_requires_a_complete_window() { + for outer in ["(1024,2048):(4096,2)", "(8,64):(128,2)", "(16,8):(16,2)"] { + let layout = ComposedLayout::new( + Swizzle::IDENTITY, + 0, + outer.parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(); + assert!(super::matrix_offset_fits_i32(&layout)); + assert!( + super::matrix_address_plan(&layout, CuteMatrixRoleAttr::A).is_none(), + "{outer}" + ); + } + } + + /// Check K-major B loads arrange paired registers in the order expected by the following MMA. + #[test] + fn attention_k_major_b_pair_loads_contiguous_mma_register_pairs() { + let mut ctx = Context::new(); + let source = module(&mut ctx, false, CuteMatrixRoleAttr::BPairKMajor); + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &source).unwrap(); + let new_block = function_block(&target.root); + let old_source = module(&mut ctx, false, CuteMatrixRoleAttr::A); + let old_target = profile.translate_module(&ctx, &old_source).unwrap(); + let old_block = function_block(&old_target.root); + + // ldmatrix.x4 gets one row address from each lane: groups of eight + // lanes supply the four 8x8 submatrices. For K-major B, adjacent + // registers must cover K=0..7 and K=8..15 at the same N coordinate. + let quadrants = [(0, 0), (0, 8), (8, 0), (8, 8)]; + for window_row in 0..4 { + for window_column in 0..4 { + for (matrix, (row_half, column_half)) in quadrants.into_iter().enumerate() { + for matrix_row in 0..8 { + let lane = matrix as u64 * 8 + matrix_row; + let args = [0x20000, window_row, window_column, lane, 0x30000]; + let addresses = memory_arguments(new_block, &args); + assert_eq!(addresses.len(), 1); + let row = window_row * 16 + row_half + matrix_row; + let column = window_column * 16 + column_half; + let bytes = (row * 64 + column) * 2; + let expected = 0x20000 + (bytes ^ ((bytes & 0x380) >> 3)); + assert_eq!( + addresses[0], + [expected], + "window=({window_row},{window_column}), lane={lane}" + ); + assert_eq!(expected % 16, 0); + + // Independent compatibility check: the hardware's new + // register order equals the existing A-to-B permutation. + let old_matrix = [0, 2, 1, 3][matrix]; + let mut old_args = args; + old_args[3] = old_matrix * 8 + matrix_row; + assert_eq!(addresses, memory_arguments(old_block, &old_args)); + } + } + } + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/gemm.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/gemm.rs new file mode 100644 index 0000000000..db3b2de3cc --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/tests/gemm.rs @@ -0,0 +1,128 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Unit tests for shared GEMM mapping invariants. + +use crate::cute::{ + epilogue::sm120::{ + epilogue_gep_element_type, epilogue_sync_operations, f16x2_intrinsic_operands, + }, + mma::sm120::ldmatrix_swizzled_byte_offset, + pipeline::sm120::mbarrier_init, +}; +use crate::{MlirAttribute, MlirLocation, MlirType, MlirValueUse}; +use dialect_cute::attributes::CuteEpilogueSyncPhaseAttr; + +#[cfg(test)] +mod tests { + //! Focused checks for GEMM byte offsets, packed lanes, and synchronization ordering. + + use super::*; + use crate::{MlirFloatType, MlirValueId}; + + /// Construct a target SSA operand with a chosen identifier and type for recipe tests. + fn value(id: u64, ty: MlirType) -> MlirValueUse { + MlirValueUse { + id: MlirValueId(id), + ty, + } + } + + /// Check the byte-domain permutation preserves each contiguous 16-byte matrix-load segment. + #[test] + fn byte_domain_ldmatrix_swizzle_preserves_sixteen_byte_segments() { + for warp_row in 0..8_u64 { + for warp_column in 0..2_u64 { + for lane in 0..32_u64 { + let submatrix = lane / 8; + let row = warp_row * 16 + (submatrix % 2) * 8 + lane % 8; + let column = warp_column * 16 + (submatrix / 2) * 8; + let plain = row * 64 + column * 2; + assert_eq!( + ldmatrix_swizzled_byte_offset(plain) % 16, + 0, + "warp_row={warp_row} warp_column={warp_column} lane={lane}" + ); + } + } + } + + // The old element-domain mask toggled bit 3 for lane 1 after the + // layout was converted to bytes, reproducing the observed 0x6448 + // misaligned shared address. + let lane_one_plain = 64_u64; + let old = lane_one_plain ^ ((lane_one_plain & 0xc0) >> 3); + assert_eq!(old, 72); + assert_eq!(old % 16, 8); + assert_eq!(ldmatrix_swizzled_byte_offset(lane_one_plain), 64); + } + + /// Check low-first source lanes are reversed for the high-first native conversion intrinsic. + #[test] + fn packed_f16_intrinsic_reverses_low_first_source_operands() { + let low = value(7, MlirType::Float(MlirFloatType::F32)); + let high = value(8, MlirType::Float(MlirFloatType::F32)); + let [first, second] = f16x2_intrinsic_operands(low.clone(), high.clone()); + assert_eq!(first, high); + assert_eq!(second, low); + } + + /// Check ReadyForTma emits proxy publication before counted sync and Reusable emits only counted sync. + #[test] + fn epilogue_ready_publishes_then_syncs_while_reusable_only_syncs() { + let barrier_id = value(1, MlirType::Integer(32)); + let participants = value(2, MlirType::Integer(32)); + let ready = epilogue_sync_operations( + CuteEpilogueSyncPhaseAttr::ReadyForTma, + barrier_id.clone(), + participants.clone(), + &MlirLocation::Unknown, + ) + .unwrap(); + assert_eq!( + ready + .iter() + .map(|operation| operation.name.as_str()) + .collect::>(), + ["nvvm.fence.proxy", "nvvm.barrier.cta.sync"] + ); + assert_eq!( + ready[0].attributes.get("kind"), + Some(&MlirAttribute::dialect("#nvvm.proxy_kind").unwrap()) + ); + assert_eq!( + ready[0].attributes.get("space"), + Some(&MlirAttribute::dialect("#nvvm.shared_space").unwrap()) + ); + assert_eq!( + ready[1].operands, + vec![barrier_id.clone(), participants.clone()] + ); + + let reusable = epilogue_sync_operations( + CuteEpilogueSyncPhaseAttr::Reusable, + barrier_id, + participants, + &MlirLocation::Unknown, + ) + .unwrap(); + assert_eq!(reusable.len(), 1); + assert_eq!(reusable[0].name, "nvvm.barrier.cta.sync"); + assert_eq!( + epilogue_gep_element_type(), + MlirType::Float(MlirFloatType::F16) + ); + } + + /// Check barrier initialization remains unpredicated once source control flow has selected the initializer. + #[test] + fn pipeline_initialization_is_unpredicated_inside_the_source_branch() { + let barrier = value(3, MlirType::dialect("!llvm.ptr<3>").unwrap()); + let count = value(4, MlirType::Integer(32)); + let init = mbarrier_init(barrier.clone(), count.clone(), &MlirLocation::Unknown).unwrap(); + assert_eq!(init.name, "nvvm.mbarrier.init"); + assert_eq!(init.operands, vec![barrier, count]); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/grid_constant.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/grid_constant.rs new file mode 100644 index 0000000000..470b248c4a --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/grid_constant.rs @@ -0,0 +1,155 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Keep grid-constant values in read-only kernel parameter storage. + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::{MlirAttribute, MlirOperation, MlirType}; +use dialect_mir::{ops::MirFuncOp, types::MirPtrType}; +use pliron::{ + builtin::{attributes::TypeAttr, type_interfaces::FunctionTypeInterface}, + context::{Context, Ptr}, + operation::Operation, + r#type::{Typed, type_cast}, +}; +use reserved_oxide_symbols::{ + MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX, MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX, +}; + +/// Validated byte size and ABI alignment for one by-value grid-constant parameter. +pub(crate) struct Parameter { + /// Number of value bytes carried by the host launch argument. + size: u64, + /// Required launch-parameter alignment in bytes. + alignment: u64, +} + +impl Parameter { + /// Build byval, alignment, and grid-constant attributes for the launch argument. + pub(crate) fn attributes(&self) -> Result { + // The byte array sets the launch argument's size. `llvm.align` keeps + // Rust's alignment, which can exceed the byte array's alignment of 1. + // The kernel body still accesses fields using the original Rust layout. + // For a TMA descriptor: 128 bytes, aligned to 64 bytes. + Ok(MlirAttribute::Dictionary(BTreeMap::from([ + ( + "llvm.byval".into(), + MlirAttribute::Type(MlirType::dialect(format!( + "!llvm.array<{} x i8>", + self.size + ))?), + ), + ( + "llvm.align".into(), + MlirAttribute::Integer { + value: self.alignment.into(), + ty: MlirType::Integer(64), + }, + ), + ("nvvm.grid_constant".into(), MlirAttribute::Unit), + ]))) + } +} + +/// Validate grid-constant metadata and return parameters by Rust argument index. +/// +/// These indices precede slice flattening. Missing or inconsistent metadata is +/// an error: the launch must carry the value's bytes, not an ordinary pointer. +pub(crate) fn take_parameters( + ctx: &Context, + source: Ptr, + target: &mut MlirOperation, +) -> Result, String> { + let function = MirFuncOp::wrap(ctx, source).ok_or("expected mir.func")?; + let function_type = function.get_type(ctx); + let function_type = function_type.deref(ctx); + let function_type = type_cast::(&*function_type) + .ok_or("mir.func does not have a function type")?; + let inputs = function_type.arg_types(); + let mut indices = BTreeSet::new(); + for key in target.attributes.keys() { + let Some((prefix, suffix)) = [ + MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX, + MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX, + ] + .into_iter() + .find_map(|prefix| key.strip_prefix(prefix).map(|suffix| (prefix, suffix))) else { + continue; + }; + let index = suffix + .parse::() + .map_err(|_| format!("malformed grid-constant attribute `{key}`"))?; + if key != &format!("{prefix}{index}") || index >= inputs.len() { + return Err(format!( + "malformed or out-of-range grid-constant attribute `{key}`" + )); + } + indices.insert(index); + } + if !indices.is_empty() + && target.attributes.get("gpu_kernel") != Some(&MlirAttribute::String("true".into())) + { + return Err("grid-constant parameters require a kernel entry".into()); + } + let mut parameters = BTreeMap::new(); + for index in indices { + let pointee_key = format!("{MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX}{index}"); + let alignment_key = format!("{MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX}{index}"); + let ( + Some(MlirAttribute::Type(_)), + Some(MlirAttribute::Integer { + value: alignment, + ty: MlirType::Integer(64), + }), + ) = ( + target.attributes.remove(&pointee_key), + target.attributes.remove(&alignment_key), + ) + else { + return Err(format!( + "parameter {index} has malformed or incomplete grid-constant metadata" + )); + }; + let alignment = u64::try_from(alignment) + .map_err(|_| format!("parameter {index} grid-constant alignment does not fit u64"))?; + if !alignment.is_power_of_two() { + return Err(format!( + "parameter {index} grid-constant alignment must be a nonzero power of two" + )); + } + let pointer_type = inputs[index].deref(ctx); + let pointer = pointer_type + .downcast_ref::() + .ok_or_else(|| format!("grid-constant parameter {index} is not a pointer"))?; + if pointer.address_space() != 0 || pointer.is_mutable() { + return Err(format!( + "grid-constant parameter {index} must be an immutable generic pointer" + )); + } + let operation = source.deref(ctx); + let key = pointee_key + .as_str() + .try_into() + .map_err(|_| "invalid grid-constant metadata key")?; + let pointee = operation + .attributes + .get::(&key) + .ok_or("grid-constant pointee is not a type")? + .get_type(ctx); + if pointee != pointer.pointee { + return Err(format!( + "grid-constant parameter {index} pointee metadata differs from its pointer type" + )); + } + let size = crate::mir::memory::source_stored_size(ctx, pointee) + .filter(|&size| size > 0) + .ok_or_else(|| { + format!("grid-constant parameter {index} requires a sized, nonzero pointee") + })?; + parameters.insert(index, Parameter { size, alignment }); + } + Ok(parameters) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/argument_locations.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/argument_locations.rs new file mode 100644 index 0000000000..4b2c2b7bfe --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/argument_locations.rs @@ -0,0 +1,178 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compatibility for the published 0.1.0 printer's missing argument locations. +//! +//! Upstream owns operation syntax, SSA names, types, and location syntax. Its +//! private block printer has no location hook, so annotate only complete block +//! headers derived from the same graph after printing. Require each affected +//! header exactly once; upstream format changes must produce a diagnostic. +//! Checked dialect tokens contain no newlines, and upstream escapes newlines +//! in string literals, so a complete header line cannot be attribute content. +//! Remove this bridge when upstream exposes argument-location printing. + +use std::{collections::HashMap, fmt}; + +use pliron::{ + basic_block::BasicBlock, + common_traits::Named, + context::{Context, Ptr}, + location::Location, + printable::State, + r#type::Typed, + value::Value, +}; +use pliron_mlir_interop::{ + MlirPrinter, MlirPrinterT, + printers::{print_op_location, print_type}, +}; + +use super::tokens::printer_error; +use crate::TranslationError; + +/// Render one graph-derived header, optionally including known argument locations. +struct BlockHeader<'a> { + /// Target block whose unique names and actual types must match upstream output. + block: Ptr, + /// Known locations; absent entries retain upstream's implicit unknown location. + locations: Option<&'a HashMap>, +} + +impl MlirPrinterT for BlockHeader<'_> { + fn mlir_print( + &self, + ctx: &Context, + state: &State, + f: &mut fmt::Formatter<'_>, + ) -> pliron::result::Result<()> { + let block = self.block.deref(ctx); + write!(f, "^{}(", block.unique_name(ctx))?; + for (index, argument) in block.arguments().enumerate() { + if index > 0 { + write!(f, ", ")?; + } + write!(f, "%{}: ", argument.unique_name(ctx))?; + print_type(ctx, argument.get_type(ctx), state, f)?; + if let Some(location) = self + .locations + .and_then(|locations| locations.get(&argument)) + { + print_op_location(ctx, location, state, f)?; + } + } + write!(f, "):")?; + Ok(()) + } +} + +/// Restore known argument locations, rejecting any unexpected upstream header format. +pub(super) fn restore( + ctx: &Context, + text: String, + blocks: impl IntoIterator>, + argument_locations: &[(Value, Location)], +) -> Result { + let locations: HashMap<_, _> = argument_locations + .iter() + .filter(|(_, location)| !location.is_unknown()) + .map(|(value, location)| (*value, location)) + .collect(); + let mut headers = HashMap::new(); + for block in blocks { + if !block + .deref(ctx) + .arguments() + .any(|argument| locations.contains_key(&argument)) + { + continue; + } + let mut header = BlockHeader { + block, + locations: None, + }; + let original = String::try_from(MlirPrinter::new(ctx, &header)) + .map_err(|error| printer_error(ctx, error))?; + header.locations = Some(&locations); + let annotated = String::try_from(MlirPrinter::new(ctx, &header)) + .map_err(|error| printer_error(ctx, error))?; + headers.insert(original, (annotated, false)); + } + replace_headers(text, headers) +} + +/// Replace complete header lines exactly once, preserving all other bytes. +fn replace_headers( + text: String, + mut headers: HashMap, +) -> Result { + if headers.is_empty() { + return Ok(text); + } + let mut output = String::with_capacity(text.len()); + for line in text.split_inclusive('\n') { + let body = line.strip_suffix('\n').unwrap_or(line); + let header = body.trim_start_matches(' '); + if let Some((annotated, seen)) = headers.get_mut(header) { + if *seen { + return Err(TranslationError::InvalidTarget( + "upstream MLIR block header occurs more than once while restoring argument locations" + .into(), + )); + } + *seen = true; + output.push_str(&body[..body.len() - header.len()]); + output.push_str(annotated); + if line.ends_with('\n') { + output.push('\n'); + } + } else { + output.push_str(line); + } + } + if headers.values().any(|(_, seen)| !seen) { + return Err(TranslationError::InvalidTarget( + "upstream MLIR block header changed while restoring argument locations".into(), + )); + } + Ok(output) +} + +#[cfg(test)] +mod tests { + use super::replace_headers; + + #[test] + fn unexpected_header_format_is_rejected_instead_of_dropping_locations() { + let headers = [( + "^block1v1(%v0: i32):".into(), + ("^block1v1(%v0: i32 loc(\"argument\")):".into(), false), + )] + .into(); + assert!( + replace_headers(" ^block1v1(%v0 : i32):\n".into(), headers) + .unwrap_err() + .to_string() + .contains("block header changed") + ); + } + + #[test] + fn duplicate_headers_are_rejected_instead_of_annotating_both() { + let headers = [( + "^block1v1(%v0: i32):".into(), + ("^block1v1(%v0: i32 loc(\"argument\")):".into(), false), + )] + .into(); + assert!( + replace_headers( + " ^block1v1(%v0: i32):\n ^block1v1(%v0: i32):\n".into(), + headers, + ) + .unwrap_err() + .to_string() + .contains("block header occurs more than once") + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/graph.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/graph.rs new file mode 100644 index 0000000000..d6104d7220 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/graph.rs @@ -0,0 +1,222 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Build a real target SSA graph before invoking upstream generic printers. + +use std::{collections::HashMap, fmt, rc::Rc}; + +use pliron::{ + attribute::AttrObj, + basic_block::BasicBlock, + context::{Context, Ptr}, + derive::{op_interface_impl, pliron_op}, + ident, input_err_noloc, + location::Location, + op::Op, + operation::Operation, + printable::State, + value::Value, +}; +use pliron_mlir_interop::{ + MlirPrinter, MlirPrinterT, ToMlirOp, + printers::{GenericOp, print_op}, +}; + +use super::tokens::{ + dictionary_key, materialize_attribute, materialize_location, materialize_type, printer_error, +}; +use crate::{MlirBlockId, MlirModule, MlirOperation, MlirValueId, TranslationError}; + +struct OperationData { + mnemonic: String, + properties: Vec<(String, AttrObj)>, + attributes: Vec<(String, AttrObj)>, +} + +type OperationDataMap = HashMap, Rc>; + +#[pliron_op(name = "cudaoxide_export.target", format, verifier = "succ")] +struct TargetOp; + +#[op_interface_impl] +impl ToMlirOp for TargetOp { + fn to_mlir( + &self, + ctx: &Context, + state: &State, + f: &mut fmt::Formatter<'_>, + ) -> pliron::result::Result<()> { + // Release the state borrow before GenericOp changes indentation while + // recursively printing regions. The metadata itself is immutable. + let metadata = state + .aux_data_ref() + .get(&ident!("cudaoxide_export_operations")) + .and_then(|data| data.downcast_ref::()) + .and_then(|data| data.get(&self.get_operation()).cloned()); + let Some(metadata) = metadata else { + return input_err_noloc!("target operation has no MLIR mapping metadata"); + }; + let mut generic = GenericOp::new(self.get_operation(), &metadata.mnemonic); + for (key, value) in &metadata.properties { + generic = generic.prop_attr(key, &**value); + } + for (key, value) in &metadata.attributes { + // Release 0.1.0 exposes attr_raw but keeps attr_attr private. Let + // the upstream attribute printer supply its validated spelling. + generic = generic.attr_raw(key, String::try_from(MlirPrinter::new(ctx, &**value))?); + } + generic.print(ctx, state, f) + } +} + +struct TargetPrinter { + root: Ptr, + operations: OperationDataMap, +} + +impl MlirPrinterT for TargetPrinter { + fn mlir_print( + &self, + ctx: &Context, + state: &State, + f: &mut fmt::Formatter<'_>, + ) -> pliron::result::Result<()> { + state.aux_data_mut().insert( + ident!("cudaoxide_export_operations"), + Box::new(self.operations.clone()), + ); + print_op(ctx, self.root, state, f) + } +} + +#[derive(Default)] +struct GraphBuilder<'a> { + values: HashMap, + blocks: HashMap>, + pending: Vec<(&'a MlirOperation, Ptr)>, + operations: OperationDataMap, + argument_locations: Vec<(Value, Location)>, +} + +impl<'a> GraphBuilder<'a> { + fn allocate( + &mut self, + ctx: &mut Context, + plan: &'a MlirOperation, + ) -> Result, TranslationError> { + MlirOperation::new(plan.name.clone()).map_err(TranslationError::InvalidTarget)?; + let types = plan + .results + .iter() + .map(|result| materialize_type(ctx, &result.ty)) + .collect::>()?; + let operation = Operation::new( + ctx, + TargetOp::get_concrete_op_info(), + types, + vec![], + vec![], + plan.regions.len(), + ); + operation + .deref_mut(ctx) + .set_loc(materialize_location(ctx, &plan.location)?); + for (result, value) in plan.results.iter().zip(operation.deref(ctx).results()) { + self.values.insert(result.id, value); + } + let properties = plan + .properties + .iter() + .map(|(name, attr)| Ok((dictionary_key(name), materialize_attribute(ctx, attr)?))) + .collect::>()?; + let attributes = plan + .attributes + .iter() + .map(|(name, attr)| Ok((dictionary_key(name), materialize_attribute(ctx, attr)?))) + .collect::>()?; + self.operations.insert( + operation, + Rc::new(OperationData { + mnemonic: plan.name.clone(), + properties, + attributes, + }), + ); + self.pending.push((plan, operation)); + for (index, region_plan) in plan.regions.iter().enumerate() { + let region = operation.deref(ctx).get_region(index); + for block_plan in ®ion_plan.blocks { + let types = block_plan + .arguments + .iter() + .map(|arg| materialize_type(ctx, &arg.ty)) + .collect::>()?; + let block = BasicBlock::new(ctx, None, types); + block.insert_at_back(region, ctx); + self.blocks.insert(block_plan.id, block); + for (argument, value) in block_plan + .arguments + .iter() + .zip(block.deref(ctx).arguments()) + { + self.values.insert(argument.id, value); + self.argument_locations + .push((value, materialize_location(ctx, &argument.location)?)); + } + for child_plan in &block_plan.operations { + let child = self.allocate(ctx, child_plan)?; + child.insert_at_back(block, ctx); + } + } + } + Ok(operation) + } + + fn wire(&self, ctx: &Context) -> Result<(), TranslationError> { + // All definitions exist first, including forward CFG destinations and + // temporaries introduced by one-to-many mappings. + for (plan, operation) in &self.pending { + for operand in &plan.operands { + let value = self.values.get(&operand.id).ok_or_else(|| { + TranslationError::InvalidTarget(format!( + "operand {:?} has no definition", + operand.id + )) + })?; + Operation::push_operand(*operation, ctx, *value); + } + for successor in &plan.successors { + let block = self.blocks.get(successor).ok_or_else(|| { + TranslationError::InvalidTarget(format!( + "successor {successor:?} has no block definition" + )) + })?; + Operation::push_successor(*operation, ctx, *block); + } + } + Ok(()) + } +} + +pub(super) fn render(module: &MlirModule) -> Result { + let ctx = &mut Context::new(); + let mut graph = GraphBuilder::default(); + let root = graph.allocate(ctx, &module.root)?; + graph.wire(ctx)?; + let printer = TargetPrinter { + root, + operations: graph.operations, + }; + let text = String::try_from(MlirPrinter::new(ctx, &printer)) + .map_err(|error| printer_error(ctx, error))?; + let mut text = super::argument_locations::restore( + ctx, + text, + graph.blocks.values().copied(), + &graph.argument_locations, + )?; + text.push('\n'); + Ok(text) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/mod.rs new file mode 100644 index 0000000000..82fc427f75 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/mod.rs @@ -0,0 +1,24 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Materialize CUDA mapping plans for the upstream MLIR printer. + +use crate::{MlirModule, TranslationError}; + +mod argument_locations; +mod graph; +mod tokens; + +/// Published printer version and the downstream compatibility revision. +pub const INTEROP_REVISION: &str = "pliron-mlir-interop-0.1.0+cuda-oxide.2"; + +/// Render a checked mapping plan through the upstream MLIR printer. +pub fn render_module(module: &MlirModule) -> Result { + crate::translation::validate_target(module)?; + graph::render(module) +} + +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tests.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tests.rs new file mode 100644 index 0000000000..14fc4e1e87 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tests.rs @@ -0,0 +1,412 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use super::render_module; +use crate::{ + MlirAttribute, MlirBlock, MlirBlockArgument, MlirBlockId, MlirFloatType, MlirLocation, + MlirModule, MlirOperation, MlirRegion, MlirResult, MlirType, MlirValueId, MlirValueUse, +}; + +fn module(root: MlirOperation) -> MlirModule { + MlirModule { + root, + profile: "interop-test".into(), + } +} + +fn argument(id: u64, location: MlirLocation) -> MlirBlockArgument { + MlirBlockArgument { + id: MlirValueId(id), + ty: MlirType::Integer(32), + location, + } +} + +#[test] +fn generic_syntax_is_stable() { + let mut add = MlirOperation::new("arith.addi").unwrap(); + add.results.push(MlirResult { + id: MlirValueId(2), + ty: MlirType::Integer(32), + }); + add.operands = [0, 1] + .map(|id| MlirValueUse { + id: MlirValueId(id), + ty: MlirType::Integer(32), + }) + .into(); + let mut root = MlirOperation::new("probe.container").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![ + argument(0, MlirLocation::Unknown), + argument(1, MlirLocation::Unknown), + ], + operations: vec![add], + }], + }); + let plan = module(root); + let text = render_module(&plan).unwrap(); + assert_eq!(text, render_module(&plan).unwrap()); + assert!( + text.contains("%v2 = \"arith.addi\"(%v0, %v1) : (i32, i32) -> (i32)"), + "{text}" + ); +} + +#[test] +fn operation_properties_use_mlir_generic_syntax() { + let mut branch = MlirOperation::new("probe.branch").unwrap(); + branch.operands.push(MlirValueUse { + id: MlirValueId(0), + ty: MlirType::Integer(32), + }); + branch.successors = vec![MlirBlockId(1), MlirBlockId(2)]; + branch.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 0]), + ); + branch + .attributes + .insert("probe.marker".into(), MlirAttribute::Unit); + let mut root = MlirOperation::new("probe.container").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![ + MlirBlock { + id: MlirBlockId(0), + arguments: vec![argument(0, MlirLocation::Unknown)], + operations: vec![branch], + }, + MlirBlock { + id: MlirBlockId(1), + arguments: vec![], + operations: vec![], + }, + MlirBlock { + id: MlirBlockId(2), + arguments: vec![], + operations: vec![], + }, + ], + }); + let text = render_module(&module(root)).unwrap(); + assert!(text.contains("\"probe.branch\"(%v0) [^"), "{text}"); + assert!( + text.contains("<{operandSegmentSizes = array}> {probe.marker = unit}"), + "{text}" + ); +} + +#[test] +fn dense_i64_properties_render_without_stringly_typed_escapes() { + let mut root = MlirOperation::new("probe.indices").unwrap(); + root.properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![0, 2])); + let text = render_module(&module(root)).unwrap(); + assert!(text.contains("<{position = array}>"), "{text}"); +} + +#[test] +fn individual_argument_locations_and_nested_operation_locations_survive() { + let first = MlirLocation::Named { + name: "first".into(), + child: Box::new(MlirLocation::Unknown), + }; + let second = MlirLocation::FileLineCol { + file: "kernel.rs".into(), + line: 17, + column: 4, + }; + let mut child = MlirOperation::new("probe.child").unwrap(); + child.operands = [0, 1] + .map(|id| MlirValueUse { + id: MlirValueId(id), + ty: MlirType::Integer(32), + }) + .into(); + child.location = MlirLocation::CallSite { + callee: Box::new(first.clone()), + caller: Box::new(second.clone()), + }; + let mut root = MlirOperation::new("probe.container").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![argument(0, first), argument(1, second)], + operations: vec![child], + }], + }); + let text = render_module(&module(root)).unwrap(); + assert!( + text.contains("%v0: i32 loc(\"first\"), %v1: i32 loc(\"kernel.rs\":17:4)"), + "{text}" + ); + assert!( + text.contains("loc(callsite(\"first\" at \"kernel.rs\":17:4))"), + "{text}" + ); + assert!( + text.contains("\"probe.child\"(%v0, %v1) : (i32, i32) -> () loc(callsite("), + "argument locations must not change operand types: {text}" + ); +} + +#[test] +fn nested_blocks_preserve_known_and_unknown_argument_locations() { + let mut nested = MlirOperation::new("probe.nested").unwrap(); + nested.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(1), + arguments: vec![ + argument(1, MlirLocation::Unknown), + argument( + 2, + MlirLocation::Named { + name: "inner".into(), + child: Box::new(MlirLocation::Unknown), + }, + ), + ], + operations: vec![], + }], + }); + let mut root = MlirOperation::new("probe.container").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![argument( + 0, + MlirLocation::Named { + name: "outer".into(), + child: Box::new(MlirLocation::Unknown), + }, + )], + operations: vec![nested], + }], + }); + let text = render_module(&module(root)).unwrap(); + assert!(text.contains("%v0: i32 loc(\"outer\")):"), "{text}"); + assert!( + text.contains("%v1: i32, %v2: i32 loc(\"inner\")):"), + "{text}" + ); + assert!(!text.contains("loc(unknown)"), "{text}"); +} + +#[test] +fn block_header_text_inside_a_string_attribute_is_unchanged() { + let mut root = MlirOperation::new("probe.container").unwrap(); + root.attributes.insert( + "probe.text".into(), + MlirAttribute::String("^block1v1(%v0: i32):\nnext".into()), + ); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![argument( + 0, + MlirLocation::Named { + name: "argument".into(), + child: Box::new(MlirLocation::Unknown), + }, + )], + operations: vec![], + }], + }); + let text = render_module(&module(root)).unwrap(); + assert!( + text.contains("^block1v1(%v0: i32 loc(\"argument\")):"), + "{text}" + ); + assert!( + text.contains(r#"probe.text = "^block1v1(%v0: i32):\nnext""#), + "{text}" + ); +} + +#[test] +fn argument_locations_preserve_fused_metadata_and_callsite_escaping() { + let mut root = MlirOperation::new("probe.container").unwrap(); + root.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: MlirBlockId(0), + arguments: vec![argument( + 0, + MlirLocation::Fused { + metadata: Some(Box::new(MlirAttribute::String("a\r\"é".into()))), + locations: vec![MlirLocation::CallSite { + callee: Box::new(MlirLocation::Named { + name: "first".into(), + child: Box::new(MlirLocation::Unknown), + }), + caller: Box::new(MlirLocation::FileLineCol { + file: "kernel.rs".into(), + line: 17, + column: 4, + }), + }], + }, + )], + operations: vec![], + }], + }); + let text = render_module(&module(root)).unwrap(); + assert!( + text.contains( + r#"%v0: i32 loc(fused<"a\0D\"\C3\A9">[callsite("first" at "kernel.rs":17:4)])"# + ), + "{text}" + ); +} + +#[test] +fn target_types_and_exact_float_bits_survive() { + let mut root = MlirOperation::new("probe.types").unwrap(); + root.properties.insert( + "types".into(), + MlirAttribute::Array(vec![ + MlirAttribute::Type(MlirType::Tuple(vec![])), + MlirAttribute::Type(MlirType::Index), + MlirAttribute::Type(MlirType::Vector { + shape: vec![2, 4], + element: Box::new(MlirType::Float(MlirFloatType::F32)), + }), + MlirAttribute::Type(MlirType::Float(MlirFloatType::Other("f8E4M3FN".into()))), + MlirAttribute::Type(MlirType::dialect("!cute.layout<\"(4,8):(1,4)\">").unwrap()), + ]), + ); + root.properties.insert( + "value".into(), + MlirAttribute::Float { + value: "0xFFC01234".into(), + ty: MlirType::Float(MlirFloatType::F32), + }, + ); + let text = render_module(&module(root)).unwrap(); + assert!( + text.contains("[tuple<>, index, vector<2x4xf32>, f8E4M3FN, !cute.layout<\"(4,8):(1,4)\">]"), + "{text}" + ); + assert!(text.contains("value = 0xFFC01234 : f32"), "{text}"); +} + +#[test] +fn raw_type_spellings_reject_controls_when_enum_constructors_are_bypassed() { + for ty in [ + MlirType::Dialect("!probe.type<\n i32>".into()), + MlirType::Float(MlirFloatType::Other("f32\0".into())), + MlirType::Tuple(vec![MlirType::Function { + inputs: vec![MlirType::Vector { + shape: vec![2], + element: Box::new(MlirType::Dialect("!probe.type<\t i32>".into())), + }], + results: vec![], + }]), + ] { + let mut root = MlirOperation::new("probe.types").unwrap(); + root.results.push(MlirResult { + id: MlirValueId(0), + ty: ty.clone(), + }); + let error = render_module(&module(root)).expect_err("raw types must remain control-free"); + assert!( + error.to_string().contains("control character"), + "{ty:?}: {error}" + ); + } +} + +#[test] +fn raw_attribute_spellings_reject_controls_when_enum_constructors_are_bypassed() { + for attr in [ + MlirAttribute::Dialect("#probe.attr<\n^block1v1(%v0: i32):\n>".into()), + MlirAttribute::Float { + value: "0x7FC00000\r".into(), + ty: MlirType::Float(MlirFloatType::F32), + }, + MlirAttribute::Array(vec![MlirAttribute::Dictionary( + [( + "nested".into(), + MlirAttribute::Dialect("#probe.attr<\t>".into()), + )] + .into(), + )]), + MlirAttribute::Array(vec![MlirAttribute::Type(MlirType::Dialect( + "!probe.type<\u{85}>".into(), + ))]), + ] { + let mut root = MlirOperation::new("probe.attributes").unwrap(); + root.attributes.insert("probe.value".into(), attr.clone()); + let error = + render_module(&module(root)).expect_err("raw attributes must remain control-free"); + assert!( + error.to_string().contains("control character"), + "{attr:?}: {error}" + ); + } +} + +#[test] +fn fused_location_metadata_rejects_raw_control_characters() { + let mut root = MlirOperation::new("probe.location").unwrap(); + root.location = MlirLocation::Fused { + metadata: Some(Box::new(MlirAttribute::Array(vec![ + MlirAttribute::Dialect("#probe.metadata<\n>".into()), + ]))), + locations: vec![MlirLocation::Unknown], + }; + let error = + render_module(&module(root)).expect_err("raw location metadata must remain control-free"); + assert!(error.to_string().contains("control character"), "{error}"); +} + +#[test] +fn upstream_string_escaping_handles_carriage_return_unicode_and_quotes() { + let mut root = MlirOperation::new("probe.strings").unwrap(); + root.attributes + .insert("probe.text".into(), MlirAttribute::String("a\r\"é".into())); + let text = render_module(&module(root)).unwrap(); + assert!(text.contains(r#"probe.text = "a\0D\"\C3\A9""#), "{text}"); +} + +#[test] +fn hyphenated_dictionary_keys_and_symbol_names_are_quoted() { + let mut root = MlirOperation::new("probe.keys").unwrap(); + root.properties + .insert("property-key".into(), MlirAttribute::Unit); + root.attributes.insert( + "probe.dictionary-key".into(), + MlirAttribute::Dictionary( + [( + "nested-key".into(), + MlirAttribute::SymbolRef("symbol-name".into()), + )] + .into_iter() + .collect(), + ), + ); + let text = render_module(&module(root)).unwrap(); + assert!(text.contains(r#"<{"property-key" = unit}>"#), "{text}"); + assert!( + text.contains(r#""probe.dictionary-key" = {"nested-key" = @"symbol-name"}"#), + "{text}" + ); +} + +#[test] +fn undefined_operands_return_a_diagnostic() { + let mut root = MlirOperation::new("probe.invalid").unwrap(); + root.operands.push(MlirValueUse { + id: MlirValueId(77), + ty: MlirType::Integer(32), + }); + assert!( + render_module(&module(root)) + .unwrap_err() + .to_string() + .contains("has no definition") + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tokens.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tokens.rs new file mode 100644 index 0000000000..8df7b0ff2c --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/interop/tokens.rs @@ -0,0 +1,256 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Adapt target types and attributes to upstream printer interfaces. +//! +//! Native Pliron builtins use their upstream implementations. Target dialect +//! tokens and typed scalar lexemes retain the spelling selected by the CUDA +//! mapping packs; no source-dialect lowering is inferred here. + +use std::fmt; + +use pliron::{ + attribute::AttrObj, + builtin::{ + attributes::{BoolAttr, StringAttr, TypeAttr, UnitAttr, VecAttr}, + types::{BF16Type, FP16Type, FP32Type, FP64Type, FunctionType, IntegerType, Signedness}, + }, + combine::stream::position::SourcePosition, + context::Context, + derive::{attr_interface_impl, pliron_attr, pliron_type, type_interface_impl}, + location::{Location, Source}, + printable::{Printable, State}, + r#type::{Type, TypeHandle}, +}; +use pliron_mlir_interop::{ + MlirPrinter, ToMlirAttr, ToMlirType, + printers::{mlir_dense_array, mlir_string_literal}, +}; + +use crate::{MlirAttribute, MlirFloatType, MlirLocation, MlirType, TranslationError}; + +#[pliron_type( + name = "cudaoxide_export.token", + format = "$spelling", + verifier = "succ" +)] +#[derive(Hash, PartialEq, Eq, Debug)] +struct TokenType { + spelling: String, +} + +#[type_interface_impl] +impl ToMlirType for TokenType { + fn to_mlir( + &self, + _: &Context, + _: &State, + f: &mut fmt::Formatter<'_>, + ) -> pliron::result::Result<()> { + validate_token(&self.spelling)?; + write!(f, "{}", self.spelling)?; + Ok(()) + } +} + +#[pliron_attr( + name = "cudaoxide_export.token", + format = "$spelling", + verifier = "succ" +)] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +struct TokenAttr { + spelling: String, +} + +#[attr_interface_impl] +impl ToMlirAttr for TokenAttr { + fn to_mlir( + &self, + _: &Context, + _: &State, + f: &mut fmt::Formatter<'_>, + ) -> pliron::result::Result<()> { + validate_token(&self.spelling)?; + write!(f, "{}", self.spelling)?; + Ok(()) + } +} + +fn validate_token(spelling: &str) -> pliron::result::Result<()> { + // Public mapping enums can bypass their checked constructors. Enforce the + // same raw-token contract at printing, including recursively nested tokens + // and fused-location metadata. Builtin strings use upstream escaping. + if spelling.chars().any(char::is_control) { + return pliron::input_err_noloc!("raw MLIR token contains a control character"); + } + Ok(()) +} + +pub(super) fn printer_error(ctx: &Context, error: pliron::result::Error) -> TranslationError { + TranslationError::InvalidTarget(format!("upstream MLIR printer: {}", error.disp(ctx))) +} + +fn type_text(ctx: &Context, ty: &MlirType) -> Result { + let ty = materialize_type(ctx, ty)?; + String::try_from(MlirPrinter::new(ctx, &ty)).map_err(|error| printer_error(ctx, error)) +} + +fn attribute_text(ctx: &Context, attr: &MlirAttribute) -> Result { + let attr = materialize_attribute(ctx, attr)?; + String::try_from(MlirPrinter::new(ctx, &*attr)).map_err(|error| printer_error(ctx, error)) +} + +fn token_attr(spelling: String) -> AttrObj { + Box::new(TokenAttr { spelling }) +} + +pub(super) fn dictionary_key(key: &str) -> String { + let mut bytes = key.bytes(); + let bare = bytes + .next() + .is_some_and(|c| c.is_ascii_alphabetic() || c == b'_') + && bytes.all(|c| c.is_ascii_alphanumeric() || b"_.$".contains(&c)); + if bare { + key.to_owned() + } else { + mlir_string_literal(key) + } +} + +pub(super) fn materialize_type( + ctx: &Context, + ty: &MlirType, +) -> Result { + let spelling = match ty { + MlirType::Integer(0) => { + return Err(TranslationError::InvalidTarget( + "integer type has zero width".into(), + )); + } + MlirType::Integer(width) => { + return Ok(IntegerType::get(ctx, *width, Signedness::Signless).into()); + } + MlirType::Float(MlirFloatType::F16) => return Ok(FP16Type::get(ctx).into()), + MlirType::Float(MlirFloatType::BF16) => return Ok(BF16Type::get(ctx).into()), + MlirType::Float(MlirFloatType::F32) => return Ok(FP32Type::get(ctx).into()), + MlirType::Float(MlirFloatType::F64) => return Ok(FP64Type::get(ctx).into()), + MlirType::Function { inputs, results } => { + let inputs = inputs + .iter() + .map(|ty| materialize_type(ctx, ty)) + .collect::>()?; + let results = results + .iter() + .map(|ty| materialize_type(ctx, ty)) + .collect::>()?; + return Ok(FunctionType::get(ctx, inputs, results).into()); + } + MlirType::Index => "index".into(), + MlirType::Float(MlirFloatType::Other(spelling)) | MlirType::Dialect(spelling) => { + spelling.clone() + } + MlirType::Vector { shape, element } => { + let dimensions = shape.iter().map(|n| format!("{n}x")).collect::(); + format!("vector<{dimensions}{}>", type_text(ctx, element)?) + } + MlirType::Tuple(elements) => { + let elements = elements + .iter() + .map(|ty| type_text(ctx, ty)) + .collect::, _>>()?; + format!("tuple<{}>", elements.join(", ")) + } + }; + Ok(Type::instantiate(TokenType { spelling }, ctx).into()) +} + +pub(super) fn materialize_attribute( + ctx: &Context, + attr: &MlirAttribute, +) -> Result { + Ok(match attr { + MlirAttribute::Unit => Box::new(UnitAttr::new()), + MlirAttribute::Bool(value) => Box::new(BoolAttr::new(*value)), + MlirAttribute::String(value) => Box::new(StringAttr::new(value.clone())), + MlirAttribute::Type(ty) => Box::new(TypeAttr::new(materialize_type(ctx, ty)?)), + MlirAttribute::Array(values) => Box::new(VecAttr::new( + values + .iter() + .map(|attr| materialize_attribute(ctx, attr)) + .collect::>()?, + )), + // Preserve exact signed decimal and floating bit-pattern spellings, + // including index-typed integers that Pliron IntegerAttr cannot hold. + MlirAttribute::Integer { value, ty } => { + token_attr(format!("{value} : {}", type_text(ctx, ty)?)) + } + MlirAttribute::Float { value, ty } => { + token_attr(format!("{value} : {}", type_text(ctx, ty)?)) + } + MlirAttribute::DenseI32Array(values) => token_attr(mlir_dense_array("i32", values)), + MlirAttribute::DenseI64Array(values) => token_attr(mlir_dense_array("i64", values)), + MlirAttribute::Dictionary(values) => { + let values = values + .iter() + .map(|(key, value)| { + Ok(format!( + "{} = {}", + dictionary_key(key), + attribute_text(ctx, value)? + )) + }) + .collect::, TranslationError>>()?; + token_attr(format!("{{{}}}", values.join(", "))) + } + MlirAttribute::SymbolRef(symbol) => token_attr(format!("@{}", dictionary_key(symbol))), + MlirAttribute::Dialect(spelling) => token_attr(spelling.clone()), + }) +} + +pub(super) fn materialize_location( + ctx: &Context, + location: &MlirLocation, +) -> Result { + Ok(match location { + MlirLocation::Unknown => Location::Unknown, + MlirLocation::FileLineCol { file, line, column } => Location::SrcPos { + src: Source::new_from_file(ctx, file), + pos: SourcePosition { + line: (*line).try_into().map_err(|_| { + TranslationError::InvalidTarget( + "source line exceeds Pliron location range".into(), + ) + })?, + column: (*column).try_into().map_err(|_| { + TranslationError::InvalidTarget( + "source column exceeds Pliron location range".into(), + ) + })?, + }, + }, + MlirLocation::Fused { + metadata, + locations, + } => Location::Fused { + metadata: metadata + .as_deref() + .map(|attr| materialize_attribute(ctx, attr)) + .transpose()?, + locations: locations + .iter() + .map(|loc| materialize_location(ctx, loc)) + .collect::>()?, + }, + MlirLocation::Named { name, child } => Location::Named { + name: name.clone(), + child_loc: Box::new(materialize_location(ctx, child)?), + }, + MlirLocation::CallSite { callee, caller } => Location::CallSite { + callee: Box::new(materialize_location(ctx, callee)?), + caller: Box::new(materialize_location(ctx, caller)?), + }, + }) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs new file mode 100644 index 0000000000..114ee5c198 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/lib.rs @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! cuda-oxide mappings for textual MLIR export. +//! +//! This crate maps cuda-oxide operations to the types and operations accepted +//! by the pinned CUTLASS compiler. Upstream pliron-mlir-interop renders the +//! resulting target graph with its generic MLIR printer. +//! +//! Use a consumer profile to export a prepared Pliron context and module: +//! +//! ```no_run +//! use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +//! # use pliron::{builtin::ops::ModuleOp, context::Context}; +//! # fn export(ctx: &Context, module: &ModuleOp) -> Result<(), cuda_oxide_mlir_export::ProfileError> { +//! let profile = CutlassFullCuteMlir22::new("sm_120a")?; +//! let target = profile.translate_module(ctx, module)?; +//! let text = cuda_oxide_mlir_export::render_module(&target)?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Both concrete profiles verify the source and wrap device functions in the +//! `gpu.module` required to produce a cubin. Registering mapping packs alone +//! does not perform those steps. Exporting text needs no GPU or CUTLASS library. + +mod cute; +mod grid_constant; +mod interop; +mod mir; +mod nvvm; +mod packs; +mod primitives; +mod profile; +mod translation; + +pub use interop::{INTEROP_REVISION, render_module}; +pub use translation::*; + +pub use cute::register_cute_elementwise_pack; +pub(crate) use cute::{ + register_cute_attention_pack, register_cute_cutlass_profile_pack, register_cute_gemm_pack, + register_cute_sm100_pack, +}; +pub use mir::core::register_mir_core_pack; +pub use nvvm::{ + attention::register_nvvm_attention_pack, cluster::register_nvvm_cluster_pack, + numeric::register_nvvm_numeric_pack, sm120::register_nvvm_sm120_pack, + sregs::register_nvvm_sreg_pack, sync_restrict::register_nvvm_sync_restrict_pack, + tcgen05::register_nvvm_tcgen05_pack, +}; +pub use packs::register_builtin_pack; +pub use primitives::{CutlassPrimitivesMlir22, verify_primitives_source}; +pub use profile::{CutlassFullCuteMlir22, ExportManifest, MlirConsumerProfile, ProfileError}; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/checked.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/checked.rs new file mode 100644 index 0000000000..ab416655cb --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/checked.rs @@ -0,0 +1,1068 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Physical Rust carriers needed by checked instruction-level APIs. +//! Enum storage retains rustc's byte offsets, size and alignment. Niche encoding +//! remains unsupported rather than being approximated by a logical tag. +use crate::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use dialect_mir::{ + attributes::VariantIndexAttr, + ops::*, + types::{ + EnumCarrierKind, EnumLayoutKind, MirArrayType, MirEnumType, MirPtrType, MirStructType, + MirTupleType, + }, +}; +use pliron::{ + builtin::types::{IntegerType, Signedness}, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Register physical enum carriers, checked integer arithmetic, and unreachable mappings. +pub(crate) fn register(r: &mut TranslationRegistry) -> Result<(), TranslationError> { + r.register_type::(EnumType)?; + r.register_attribute::(DropAttribute)?; + r.register_operation::(EnumOp::Construct)?; + r.register_operation::(EnumOp::Discriminant)?; + r.register_operation::(EnumOp::Payload)?; + r.register_operation::(Checked("add"))?; + r.register_operation::(Checked("mul"))?; + r.register_operation::(Unreachable)?; + Ok(()) +} + +/// Validate reviewed enum layouts and preserve their physical size and alignment. +pub(crate) fn enum_storage_type(e: &MirEnumType) -> Result { + // Infallible core array helpers retain their uninhabited residual type in + // unreachable branches. It has a carrier for SSA typing, but no value can + // be constructed and no discriminant or payload may be read from it. + if e.layout_kind == EnumLayoutKind::Empty { + if e.total_size != 0 + || e.abi_align != 1 + || e.tag_offset != 0 + || e.carrier_kind != EnumCarrierKind::None + || e.carrier_width != 0 + || e.carrier_address_space != 0 + || e.variant_names.len() != e.variant_discriminants.len() + || e.variant_names.len() != e.variant_field_counts.len() + || e.variant_names.len() != e.variant_inhabited.len() + || e.variant_inhabited.iter().any(|&inhabited| inhabited != 0) + || e.variant_field_counts + .iter() + .map(|&n| u64::from(n)) + .sum::() + != e.all_field_types.len() as u64 + || e.all_field_types.len() != e.all_field_offsets.len() + || e.all_field_types.len() != e.all_field_sizes.len() + || e.all_field_offsets.iter().any(|&offset| offset != 0) + || e.all_field_sizes.iter().any(|&size| size != 0) + { + return Err( + "empty enum storage requires a zero-sized layout with no inhabited variants".into(), + ); + } + return MlirType::dialect("!llvm.struct<(array<0 x i8>, array<0 x i8>)>"); + } + if !matches!( + e.layout_kind, + EnumLayoutKind::Direct | EnumLayoutKind::Single + ) { + return Err("MLIR export supports only reviewed Direct/Single enum layouts".into()); + } + if !matches!(e.abi_align, 1 | 2 | 4 | 8) || !e.total_size.is_multiple_of(e.abi_align) { + return Err("enum storage needs known size and alignment of 1, 2, 4 or 8".into()); + } + if e.layout_kind == EnumLayoutKind::Direct + && (e.carrier_kind != EnumCarrierKind::Integer + || !matches!(e.carrier_width, 8 | 16 | 32 | 64) + || e.tag_offset + .checked_add(u64::from(e.carrier_width) / 8) + .is_none_or(|end| end > e.total_size)) + { + return Err("direct enum needs an in-bounds integer tag".into()); + } + if e.all_field_types.len() != e.all_field_offsets.len() + || e.all_field_types.len() != e.all_field_sizes.len() + { + return Err("enum lacks complete rustc field layout".into()); + } + for (&offset, &size) in e.all_field_offsets.iter().zip(&e.all_field_sizes) { + if offset + .checked_add(size) + .is_none_or(|end| end > e.total_size) + { + return Err("enum field exceeds storage".into()); + } + } + // Independent bytes keep inactive payload/padding undef separate from the + // initialized tag. A zero-length carrier supplies alignment without bytes. + MlirType::dialect(format!( + "!llvm.struct<(array<{} x i8>, array<0 x i{}>)>", + e.total_size, + e.abi_align * 8 + )) +} +/// Translate a reviewed enum layout without conflating inactive bytes with its initialized tag. +struct EnumType; +// These payloads require storage adapters, not merely a same-size load/store. +// Keep the initial enum pack fail-closed until those adapters are reviewed. +/// Recursively reject bool and named-space pointer payloads needing unimplemented byte adapters. +pub(super) fn check_payload_storage(ctx: &Context, ty: TypeHandle) -> Result<(), String> { + let t = ty.deref(ctx); + if t.downcast_ref::() + .is_some_and(|i| i.width() == 1) + { + return Err("enum bool payloads require canonical byte storage".into()); + } + if t.downcast_ref::() + .is_some_and(|p| p.address_space() != 0) + { + return Err("enum named-space pointers require generic storage adapters".into()); + } + let fields: &[TypeHandle] = if let Some(e) = t.downcast_ref::() { + &e.all_field_types + } else if let Some(s) = t.downcast_ref::() { + s.field_types() + } else if let Some(t) = t.downcast_ref::() { + t.get_types() + } else if let Some(a) = t.downcast_ref::() { + if a.size() != 0 { + check_payload_storage(ctx, a.element_type())?; + } + return Ok(()); + } else { + return Ok(()); + }; + for &field in fields { + check_payload_storage(ctx, field)?; + } + Ok(()) +} +impl TypeTranslation for EnumType { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _: &TranslationRegistry, + ) -> Result { + check_payload_storage(ctx, source)?; + enum_storage_type( + source + .deref(ctx) + .downcast_ref::() + .ok_or("expected enum")?, + ) + } +} + +/// Build an operation with the supplied result IDs, operands, and source location. +pub(super) fn op( + name: &str, + operands: Vec, + results: Vec, + loc: &MlirLocation, +) -> Result { + let mut op = MlirOperation::new(name)?; + op.operands = operands; + op.results = results; + op.location = loc.clone(); + Ok(op) +} +/// Allocate one SSA identifier and return matching result and operand handles. +fn fresh(s: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = s.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} +/// Append an integer constant of the requested bit width to the output sequence. +fn integer( + value: i128, + width: u32, + s: &mut TranslationSession<'_>, + loc: &MlirLocation, + out: &mut Vec, +) -> Result { + let ty = MlirType::Integer(width); + let (r, v) = fresh(s, ty.clone()); + let mut c = op("arith.constant", vec![], vec![r], loc)?; + c.properties + .insert("value".into(), MlirAttribute::Integer { value, ty }); + out.push(c); + Ok(v) +} +// Private markers are consumed by FunctionTranslation before MLIR is emitted. +// Only the scratch allocations below qualify; Rust/MIR allocas are untouched. +/// Private marker identifying the constant allocation count for enum scratch storage. +const ENUM_SCRATCH_COUNT: &str = "cuda_oxide.enum_scratch_count"; +/// Private marker identifying a scratch allocation eligible for entry-block hoisting. +const ENUM_SCRATCH_SLOT: &str = "cuda_oxide.enum_scratch_slot"; + +/// Make enum packing scratch visible to LLVM's entry-block scalar replacement. +/// +/// Each slot is private to one enum operation: its address is only used by +/// immediate byte-GEP/store/load operations. Moving the allocation, without +/// moving those accesses, preserves conditional initialization and loop values. +/// Distinct slots are never merged, and inactive payload/padding bytes retain +/// their independent byte representation. +pub(crate) fn hoist_enum_scratch(function: &mut MlirOperation) -> Result<(), String> { + /// Recognize either private marker used by enum scratch hoisting. + fn marked(op: &MlirOperation) -> bool { + op.attributes.contains_key(ENUM_SCRATCH_COUNT) + || op.attributes.contains_key(ENUM_SCRATCH_SLOT) + } + /// Reject marked scratch operations nested below the function's direct block operations. + fn reject_nested_markers(op: &MlirOperation) -> Result<(), String> { + for region in &op.regions { + for block in ®ion.blocks { + for nested in &block.operations { + if marked(nested) { + return Err("enum scratch in a nested region is not supported".into()); + } + reject_nested_markers(nested)?; + } + } + } + Ok(()) + } + if function.regions.len() != 1 { + return Err("enum scratch hoisting requires one function body region".into()); + } + let region = &mut function.regions[0]; + if region.blocks.is_empty() { + return Ok(()); // External declaration. + } + let mut prologue = Vec::new(); + for block in &mut region.blocks { + let mut body = Vec::new(); + let mut operations = core::mem::take(&mut block.operations).into_iter(); + while let Some(mut count) = operations.next() { + reject_nested_markers(&count)?; + if !marked(&count) { + body.push(count); + continue; + } + let valid_count = count.name == "arith.constant" + && count.attributes.len() == 1 + && count.attributes.get(ENUM_SCRATCH_COUNT) == Some(&MlirAttribute::Unit) + && count.operands.is_empty() + && count.results.len() == 1 + && count.results[0].ty == MlirType::Integer(32) + && count.properties.get("value") + == Some(&MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(32), + }) + && count.regions.is_empty() + && count.successors.is_empty(); + if !valid_count { + return Err("malformed enum scratch count marker".into()); + } + let mut allocation = operations + .next() + .ok_or("enum scratch count has no allocation")?; + let valid_allocation = allocation.name == "llvm.alloca" + && allocation.attributes.len() == 1 + && allocation.attributes.get(ENUM_SCRATCH_SLOT) == Some(&MlirAttribute::Unit) + && allocation.operands + == vec![MlirValueUse { + id: count.results[0].id, + ty: MlirType::Integer(32), + }] + && allocation.results.len() == 1 + && allocation.results[0].ty == MlirType::dialect("!llvm.ptr")? + && allocation.properties.contains_key("elem_type") + && allocation.regions.is_empty() + && allocation.successors.is_empty(); + if !valid_allocation { + return Err("malformed enum scratch allocation or count dependency".into()); + } + count.attributes.remove(ENUM_SCRATCH_COUNT); + allocation.attributes.remove(ENUM_SCRATCH_SLOT); + prologue.extend([count, allocation]); + } + block.operations = body; + } + prologue.append(&mut region.blocks[0].operations); + region.blocks[0].operations = prologue; + Ok(()) +} + +/// Allocate one private, aligned enum scratch slot and mark it for entry-block hoisting. +pub(super) fn slot( + ty: MlirType, + align: u64, + s: &mut TranslationSession<'_>, + loc: &MlirLocation, + out: &mut Vec, +) -> Result { + let one = integer(1, 32, s, loc, out)?; + out.last_mut() + .expect("integer appends its constant") + .attributes + .insert(ENUM_SCRATCH_COUNT.into(), MlirAttribute::Unit); + let (r, p) = fresh(s, MlirType::dialect("!llvm.ptr")?); + let mut a = op("llvm.alloca", vec![one], vec![r], loc)?; + a.attributes + .insert(ENUM_SCRATCH_SLOT.into(), MlirAttribute::Unit); + a.properties + .insert("elem_type".into(), MlirAttribute::Type(ty)); + a.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: align.into(), + ty: MlirType::Integer(64), + }, + ); + out.push(a); + Ok(p) +} +/// Offset an enum scratch pointer by its rustc-provided byte position. +fn byte_pointer( + base: MlirValueUse, + offset: u64, + s: &mut TranslationSession<'_>, + loc: &MlirLocation, + out: &mut Vec, +) -> Result { + let offset = i32::try_from(offset).map_err(|_| "enum offset exceeds supported GEP range")?; + let (r, p) = fresh(s, base.ty.clone()); + let mut gep = op("llvm.getelementptr", vec![base], vec![r], loc)?; + gep.properties.insert( + "elem_type".into(), + MlirAttribute::Type(MlirType::Integer(8)), + ); + gep.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![offset]), + ); + gep.properties.insert( + "noWrapFlags".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(32), + }, + ); + out.push(gep); + Ok(p) +} +/// Write one enum tag or payload to potentially unaligned byte storage. +pub(super) fn store( + value: MlirValueUse, + pointer: MlirValueUse, + loc: &MlirLocation, + out: &mut Vec, +) -> Result<(), String> { + let mut st = op("llvm.store", vec![value, pointer], vec![], loc)?; + // Variant payloads may be packed or overlap differently typed variants. + st.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(64), + }, + ); + out.push(st); + Ok(()) +} +/// Read one enum tag or payload from potentially unaligned byte storage. +pub(super) fn load( + pointer: MlirValueUse, + result: MlirResult, + loc: &MlirLocation, + out: &mut Vec, +) -> Result<(), String> { + let mut ld = op("llvm.load", vec![pointer], vec![result], loc)?; + ld.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(64), + }, + ); + out.push(ld); + Ok(()) +} +/// Select physical enum construction, tag extraction, or payload extraction. +enum EnumOp { + /// Initialize the active variant's tag and payload at their physical byte offsets. + Construct, + /// Read the variant tag from its rustc-provided storage position. + Discriminant, + /// Read a selected variant payload field from physical enum storage. + Payload, +} +impl OperationTranslation for EnumOp { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + s: &mut TranslationSession<'_>, + ) -> Result, String> { + if input.results.len() != 1 + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("invalid enum operation shape".into()); + } + let source_ty = if matches!(self, Self::Construct) { + source.deref(ctx).get_result(0).get_type(ctx) + } else { + source.deref(ctx).get_operand(0).get_type(ctx) + }; + let e = source_ty + .deref(ctx) + .downcast_ref::() + .ok_or("expected enum type")? + .clone(); + let storage_ty = enum_storage_type(&e)?; + if e.layout_kind == EnumLayoutKind::Empty { + return Err("cannot construct or inspect an uninhabited enum".into()); + } + let loc = &input.location; + let result = input.results[0].clone(); + let mut out = vec![]; + if matches!(self, Self::Discriminant) && e.layout_kind == EnumLayoutKind::Single { + let variant = e.single_variant as usize; + if e.variant_is_inhabited(variant) != Some(true) { + return Err("cannot read uninhabited enum discriminant".into()); + } + let bits = *e + .variant_discriminants + .get(variant) + .ok_or("missing single-variant discriminant")?; + let mut c = op("arith.constant", vec![], vec![result.clone()], loc)?; + c.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: bits.into(), + ty: result.ty, + }, + ); + return Ok(vec![c]); + } + let p = slot(storage_ty, e.abi_align, s, loc, &mut out)?; + match self { + Self::Construct => { + let variant = MirConstructEnumOp::new(source) + .get_attr_construct_enum_variant_index(ctx) + .ok_or("missing variant")? + .0 as usize; + if e.variant_is_inhabited(variant) != Some(true) { + return Err("cannot construct uninhabited enum variant".into()); + } + let count = *e + .variant_field_counts + .get(variant) + .ok_or("invalid variant")? as usize; + if count != input.operands.len() { + return Err("enum payload arity mismatch".into()); + } + if e.layout_kind == EnumLayoutKind::Direct { + let tag = integer( + e.variant_discriminants[variant].into(), + e.carrier_width, + s, + loc, + &mut out, + )?; + let ptr = byte_pointer(p.clone(), e.tag_offset, s, loc, &mut out)?; + store(tag, ptr, loc, &mut out)?; + } else if variant != e.single_variant as usize { + return Err("construction disagrees with single variant layout".into()); + } + let start = e.variant_field_counts[..variant] + .iter() + .map(|v| *v as usize) + .sum::(); + for (i, value) in input.operands.into_iter().enumerate() { + let size = + crate::mir::memory::source_stored_size(ctx, e.all_field_types[start + i]) + .ok_or("unknown enum field storage")?; + if size != e.all_field_sizes[start + i] { + return Err("enum payload storage size differs from rustc".into()); + } + if size == 0 { + continue; + } + let ptr = + byte_pointer(p.clone(), e.all_field_offsets[start + i], s, loc, &mut out)?; + store(value, ptr, loc, &mut out)?; + } + load(p, result, loc, &mut out)?; + } + Self::Discriminant => { + store(input.operands[0].clone(), p.clone(), loc, &mut out)?; + if result.ty != MlirType::Integer(e.carrier_width) { + return Err("direct enum tag/result widths differ".into()); + } + let ptr = byte_pointer(p, e.tag_offset, s, loc, &mut out)?; + load(ptr, result, loc, &mut out)?; + } + Self::Payload => { + let payload = MirEnumPayloadOp::new(source); + let variant = payload + .get_attr_payload_variant_index(ctx) + .ok_or("missing payload variant")? + .0 as usize; + let field = payload + .get_attr_payload_field_index(ctx) + .ok_or("missing payload field")? + .0 as usize; + if field + >= *e + .variant_field_counts + .get(variant) + .ok_or("invalid payload variant")? as usize + { + return Err("invalid payload field".into()); + } + let flat = e.variant_field_counts[..variant] + .iter() + .map(|v| *v as usize) + .sum::() + + field; + if crate::mir::memory::source_stored_size(ctx, e.all_field_types[flat]) + != Some(e.all_field_sizes[flat]) + { + return Err("enum payload storage size differs from rustc".into()); + } + if e.all_field_sizes[flat] == 0 { + out.push(op("llvm.mlir.poison", vec![], vec![result], loc)?); + } else { + store(input.operands[0].clone(), p.clone(), loc, &mut out)?; + let ptr = byte_pointer(p, e.all_field_offsets[flat], s, loc, &mut out)?; + load(ptr, result, loc, &mut out)?; + } + } + } + Ok(out) + } +} + +/// Emit signed or unsigned overflow arithmetic and rebuild the Rust result tuple. +struct Checked( + #[doc = "Arithmetic family, currently add or mul, used to select an overflow intrinsic."] + &'static str, +); +impl OperationTranslation for Checked { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + s: &mut TranslationSession<'_>, + ) -> Result, String> { + if input.operands.len() != 2 + || input.results.len() != 1 + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("invalid checked arithmetic shape".into()); + } + let ty = source.deref(ctx).get_operand(0).get_type(ctx); + let ty_ref = ty.deref(ctx); + let integer = ty_ref + .downcast_ref::() + .ok_or("checked arithmetic needs integers")?; + let width = integer.width(); + let sign = if integer.signedness() == Signedness::Signed { + "s" + } else { + "u" + }; + let result_ty = source.deref(ctx).get_result(0).get_type(ctx); + let result_ref = result_ty.deref(ctx); + let tuple = result_ref + .downcast_ref::() + .ok_or("checked arithmetic needs a tuple result")?; + if tuple.get_types().len() != 2 + || tuple.get_types()[0] != ty + || s.translate_type(tuple.get_types()[1]) + .map_err(|e| e.to_string())? + != MlirType::Integer(1) + { + return Err("checked tuple must contain value and overflow flag".into()); + } + let (r, v) = fresh( + s, + MlirType::dialect(format!("!llvm.struct<(i{width}, i1)>"))?, + ); + let mut call = op( + "llvm.call_intrinsic", + input.operands.clone(), + vec![r], + &input.location, + )?; + call.properties.insert( + "intrin".into(), + MlirAttribute::String(format!("llvm.{sign}{}.with.overflow.i{width}", self.0)), + ); + call.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + call.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![2, 0]), + ); + let mut out = vec![call]; + let mut fields = vec![]; + for (i, ty) in [MlirType::Integer(width), MlirType::Integer(1)] + .into_iter() + .enumerate() + { + let (r, field) = fresh(s, ty); + let mut extract = op( + "llvm.extractvalue", + vec![v.clone()], + vec![r], + &input.location, + )?; + extract.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![i as i64]), + ); + out.push(extract); + fields.push(field); + } + input.operands = fields; + out.extend( + crate::mir::memory::ConstructStructTranslation.translate(ctx, source, input, s)?, + ); + Ok(out) + } +} +/// Trap before marking control flow unreachable at the device contract boundary. +struct Unreachable; +impl OperationTranslation for Unreachable { + fn translate( + &self, + _: &Context, + _: Ptr, + input: OperationInput, + _: &mut TranslationSession<'_>, + ) -> Result, String> { + if !input.operands.is_empty() + || !input.results.is_empty() + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("invalid unreachable shape".into()); + } + Ok(vec![op( + "llvm.unreachable", + vec![], + vec![], + &input.location, + )?]) + } +} + +#[cfg(test)] +mod tests { + //! Enum byte-layout and scratch-hoisting regressions, including conditional and loop-local accesses. + + use super::*; + use dialect_mir::types::{EnumEncoding, EnumVariant}; + /// Construct the explicit Direct enum layout used by byte-storage rejection tests. + fn direct(ctx: &mut Context) -> MirEnumType { + let tag = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let word = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + MirEnumType::get_with_encoding( + ctx, + "ResultWord".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Ok".into(), vec![word], vec![8], vec![8]), + EnumVariant::new_with_layout("Err".into(), vec![tag], vec![1], vec![1]), + ], + EnumEncoding { + total_size: 16, + abi_align: 8, + layout_kind: EnumLayoutKind::Direct, + carrier_kind: EnumCarrierKind::Integer, + carrier_width: 8, + variant_inhabited: vec![1, 1], + ..Default::default() + }, + ) + .deref(ctx) + .clone() + } + /// Reject nested bool or named-space pointer payloads until their byte-storage adapters are supported. + #[test] + fn enum_payloads_reject_unreviewed_byte_and_pointer_adapters_recursively() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let array: TypeHandle = MirArrayType::get(&mut ctx, boolean, 4).into(); + assert!( + check_payload_storage(&ctx, array) + .unwrap_err() + .contains("bool") + ); + let pointer: TypeHandle = MirPtrType::get(&mut ctx, boolean, true, 3).into(); + let tuple: TypeHandle = MirTupleType::get(&mut ctx, vec![pointer]).into(); + assert!( + check_payload_storage(&ctx, tuple) + .unwrap_err() + .contains("named-space") + ); + let generic: TypeHandle = MirPtrType::get(&mut ctx, boolean, true, 0).into(); + check_payload_storage(&ctx, generic).unwrap(); // pointee is not stored inline + } + /// Check inactive payload and padding bytes remain independent of the initialized enum tag. + #[test] + fn enum_storage_keeps_padding_independent_of_initialized_tag_bytes() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let e = direct(&mut ctx); + assert_eq!( + enum_storage_type(&e).unwrap(), + MlirType::dialect("!llvm.struct<(array<16 x i8>, array<0 x i64>)>").unwrap() + ); + let mut single = e.clone(); + single.layout_kind = EnumLayoutKind::Single; + single.carrier_kind = EnumCarrierKind::None; + single.carrier_width = 0; + assert_eq!( + enum_storage_type(&single).unwrap(), + enum_storage_type(&e).unwrap() + ); + for kind in [ + EnumLayoutKind::Niche, + EnumLayoutKind::Unknown, + EnumLayoutKind::Empty, + ] { + let mut invalid = e.clone(); + invalid.layout_kind = kind; + assert!(enum_storage_type(&invalid).is_err()); + } + let mut invalid = e.clone(); + invalid.tag_offset = 16; + assert!(enum_storage_type(&invalid).is_err()); + let mut invalid = e.clone(); + invalid.all_field_offsets[0] = 12; + assert!(enum_storage_type(&invalid).is_err()); + } + /// Check scratch allocations move to entry while conditional construction and loop projections remain in place. + #[test] + fn conditional_enum_construct_and_loop_projections_use_entry_scratch() { + use crate::{CutlassPrimitivesMlir22, MlirConsumerProfile}; + use dialect_mir::attributes::FieldIndexAttr; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{ + OperandSegmentInterface, SingleBlockRegionInterface, SymbolOpInterface, + }, + ops::ModuleOp, + types::FunctionType, + }, + op::Op, + value::Value, + }; + /// Append a typed source operation to the fixture block and return its operation handle. + fn append( + ctx: &mut Context, + block: Ptr, + types: Vec, + values: Vec, + ) -> Ptr { + let op = Operation::new(ctx, O::get_concrete_op_info(), types, values, vec![], 0); + op.insert_at_back(block, ctx); + op + } + /// Create a conditional source branch with explicit destination blocks and no successor operands. + fn conditional( + ctx: &mut Context, + block: Ptr, + flag: Value, + left: Ptr, + right: Ptr, + lhs: Vec, + rhs: Vec, + ) { + let (operands, segments) = + MirCondBranchOp::compute_segment_sizes(vec![vec![flag], lhs, rhs]); + let op = Operation::new( + ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + operands, + vec![left, right], + 0, + ); + MirCondBranchOp::new(op).set_operand_segment_sizes(ctx, segments); + op.insert_at_back(block, ctx); + } + /// Find the translated function recursively through the module envelope. + fn find_function(op: &MlirOperation) -> Option<&MlirOperation> { + if op.name == "func.func" { + return Some(op); + } + op.regions + .iter() + .flat_map(|r| &r.blocks) + .flat_map(|b| &b.operations) + .find_map(find_function) + } + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let tag: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let flag: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let enumeration: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Result".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Ok".into(), vec![word], vec![8], vec![8]), + EnumVariant::new_with_layout("Err".into(), vec![tag], vec![1], vec![1]), + ], + 0, + 16, + 8, + ) + .into(); + let module = ModuleOp::new(&mut ctx, "enum_loop".try_into().unwrap()); + let ty = FunctionType::get(&ctx, vec![flag, word], vec![word]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, f, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "test".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let region = f.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![flag, word]); + let construct = BasicBlock::new(&mut ctx, None, vec![]); + let loop_body = BasicBlock::new(&mut ctx, None, vec![enumeration]); + let exit = BasicBlock::new(&mut ctx, None, vec![word]); + for block in [entry, construct, loop_body, exit] { + block.insert_at_back(region, &ctx); + } + let condition = entry.deref(&ctx).get_argument(0); + let input = entry.deref(&ctx).get_argument(1); + conditional( + &mut ctx, + entry, + condition, + construct, + exit, + vec![], + vec![input], + ); + let c = append::(&mut ctx, construct, vec![enumeration], vec![input]); + MirConstructEnumOp::new(c).set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + let value = c.deref(&ctx).get_result(0); + let branch = Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![value], + vec![loop_body], + 0, + ); + branch.insert_at_back(construct, &ctx); + let carried = loop_body.deref(&ctx).get_argument(0); + append::(&mut ctx, loop_body, vec![tag], vec![carried]); + let payload = append::(&mut ctx, loop_body, vec![word], vec![carried]); + let p = MirEnumPayloadOp::new(payload); + p.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); + p.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + let value = payload.deref(&ctx).get_result(0); + conditional( + &mut ctx, + loop_body, + condition, + loop_body, + exit, + vec![carried], + vec![value], + ); + let result = exit.deref(&ctx).get_argument(0); + append::(&mut ctx, exit, vec![], vec![result]); + let translated = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let f = find_function(&translated.root).unwrap(); + let blocks = &f.regions[0].blocks; + let allocas: Vec<_> = blocks[0] + .operations + .iter() + .filter(|o| o.name == "llvm.alloca") + .collect(); + assert_eq!( + allocas.len(), + 3, + "all three enum operations need distinct entry scratch" + ); + let ids: std::collections::HashSet<_> = allocas.iter().map(|o| o.results[0].id).collect(); + assert_eq!(ids.len(), 3); + assert!( + !blocks[0] + .operations + .iter() + .any(|o| o.name == "llvm.store" || o.name == "llvm.load") + ); + assert_eq!( + blocks[1] + .operations + .iter() + .filter(|o| o.name == "llvm.store") + .count(), + 2 + ); + assert_eq!( + blocks[2] + .operations + .iter() + .filter(|o| o.name == "llvm.store") + .count(), + 2 + ); + assert!( + blocks[1..] + .iter() + .flat_map(|b| &b.operations) + .all(|o| o.name != "llvm.alloca") + ); + assert_eq!(blocks[0].operations.last().unwrap().name, "cf.cond_br"); + assert_eq!( + blocks[2].operations.last().unwrap().successors, + vec![blocks[2].id, blocks[3].id] + ); + let text = crate::render_module(&translated).unwrap(); + assert!(!text.contains(ENUM_SCRATCH_COUNT) && !text.contains(ENUM_SCRATCH_SLOT)); + } + + /// Build a small target function with marked enum scratch for hoisting validation. + fn scratch_fixture() -> MlirOperation { + use crate::{MlirBlock, MlirBlockId, MlirRegion, MlirValueId}; + let mut count = MlirOperation::new("arith.constant").unwrap(); + count.results.push(MlirResult { + id: MlirValueId(0), + ty: MlirType::Integer(32), + }); + count.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(32), + }, + ); + count + .attributes + .insert(ENUM_SCRATCH_COUNT.into(), MlirAttribute::Unit); + let mut allocation = MlirOperation::new("llvm.alloca").unwrap(); + allocation.results.push(MlirResult { + id: MlirValueId(1), + ty: MlirType::dialect("!llvm.ptr").unwrap(), + }); + allocation.operands.push(MlirValueUse { + id: MlirValueId(0), + ty: MlirType::Integer(32), + }); + allocation.properties.insert( + "elem_type".into(), + MlirAttribute::Type(MlirType::Integer(64)), + ); + allocation + .attributes + .insert(ENUM_SCRATCH_SLOT.into(), MlirAttribute::Unit); + let mut f = MlirOperation::new("func.func").unwrap(); + f.regions.push(MlirRegion { + blocks: vec![ + MlirBlock { + id: MlirBlockId(0), + arguments: vec![], + operations: vec![MlirOperation::new("cf.br").unwrap()], + }, + MlirBlock { + id: MlirBlockId(1), + arguments: vec![], + operations: vec![count, allocation], + }, + ], + }); + f + } + + /// Check the hoister moves only its explicitly marked constant-size scratch allocations. + #[test] + fn enum_scratch_hoist_leaves_unmarked_and_dynamic_allocations_in_place() { + let mut f = scratch_fixture(); + let mut ordinary = f.regions[0].blocks[1].operations[1].clone(); + ordinary.attributes.clear(); + ordinary.operands[0].id = crate::MlirValueId(99); // Dynamic count. + f.regions[0].blocks[1].operations.push(ordinary.clone()); + hoist_enum_scratch(&mut f).unwrap(); + assert_eq!(f.regions[0].blocks[1].operations, vec![ordinary]); + assert_eq!(f.regions[0].blocks[0].operations.len(), 3); + } + + /// Reject malformed scratch markers or allocation dependencies that cannot safely move to entry. + #[test] + fn enum_scratch_hoist_rejects_malformed_markers_and_dependencies() { + for case in 0..5 { + let mut f = scratch_fixture(); + let operations = &mut f.regions[0].blocks[1].operations; + match case { + 0 => { + operations[0].properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 2, + ty: MlirType::Integer(32), + }, + ); + } + 1 => { + operations[1].operands[0].id = crate::MlirValueId(99); + } + 2 => { + operations[1].attributes.clear(); + } + 3 => { + operations.remove(0); + } + 4 => { + operations[0].attributes.insert( + ENUM_SCRATCH_COUNT.into(), + MlirAttribute::String("true".into()), + ); + } + _ => unreachable!(), + } + assert!(hoist_enum_scratch(&mut f).is_err(), "case {case}"); + } + let mut f = scratch_fixture(); + let mut nested = MlirOperation::new("func.func").unwrap(); + nested.regions = vec![f.regions[0].clone()]; + f.regions[0].blocks[1].operations = vec![nested]; + assert!( + hoist_enum_scratch(&mut f) + .unwrap_err() + .contains("nested region") + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs new file mode 100644 index 0000000000..66ebb09f11 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs @@ -0,0 +1,1957 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The part of `dialect-mir` that maps directly to standard MLIR dialects. +//! +//! Rust signedness lives on the Pliron integer type. MLIR integers are +//! signless, so operations that care about signedness choose an explicit +//! signed or unsigned `arith` operation before that information disappears. + +use crate::{ + AttributeTranslation, DropAttribute, FixedType, MlirAttribute, MlirBlockArgument, + MlirFloatType, MlirOperation, MlirResult, MlirType, MlirValueUse, OperationInput, + OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, +}; +use dialect_mir::{ + attributes::{ + InlineIntentAttr, MirCastKindAttr, MirFP16Attr, MirPointerKindAuthorityAttr, + ReferenceParamValidityAttr, UnrollAttr, + }, + ops::{ + MirAddOp, MirBitAndOp, MirBitOrOp, MirBitXorOp, MirCastOp, MirCondBranchOp, MirConstantOp, + MirDivOp, MirEqOp, MirFloatConstantOp, MirFuncOp, MirGeOp, MirGotoOp, MirGtOp, MirLeOp, + MirLtOp, MirMulOp, MirNeOp, MirNegOp, MirNotOp, MirRemOp, MirReturnOp, MirShlOp, MirShrOp, + MirSubOp, MirUndefOp, + }, + types::{MirDisjointSliceType, MirFP16Type, MirSliceType}, +}; +use pliron::{ + attribute::AttrObj, + builtin::{ + op_interfaces::{ATTR_KEY_SYM_NAME, OperandSegmentInterface, SymbolOpInterface}, + type_interfaces::{FloatTypeInterface, FunctionTypeInterface}, + types::{FP16Type, FP32Type, FP64Type, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed, type_cast, type_impls}, + utils::apfloat::Float, +}; + +/// Register standard-MLIR mappings for scalar MIR and arbitrary CFG. +pub fn register_mir_core_pack(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_type::(FixedType(MlirType::Float(MlirFloatType::F16)))?; + registry.register_attribute::(MirFp16AttributeTranslation)?; + // A `mir.func`'s `#[inline]` request. The function translation folds the + // keyword into the LLVM `passthrough` list, as `mir-lower` does natively. + registry.register_attribute::(InlineIntentTranslation)?; + registry.register_attribute::(DropAttribute)?; + // Rust-boundary provenance metadata consumed by the dialect-mir verifier; + // like the cast kind, it has no MLIR counterpart and is dropped at export. + registry.register_attribute::(DropAttribute)?; + // These optional optimization hints use Rust argument indices. Drop them + // until they can be remapped to the launch arguments after slice flattening. + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + + registry.register_operation::(FunctionTranslation)?; + registry.register_operation::(RenameOperation("func.return"))?; + registry.register_operation::(RenameOperation("cf.br"))?; + registry.register_operation::(ConditionalBranchTranslation)?; + + registry.register_operation::(ConstantTranslation::Integer)?; + registry.register_operation::(ConstantTranslation::Float)?; + registry.register_operation::(RenameOperation("ub.poison"))?; + + registry.register_operation::(NumericBinaryTranslation::Add)?; + registry.register_operation::(NumericBinaryTranslation::Sub)?; + registry.register_operation::(NumericBinaryTranslation::Mul)?; + registry.register_operation::(NumericBinaryTranslation::Div)?; + registry.register_operation::(NumericBinaryTranslation::Rem)?; + registry.register_operation::(IntegerBinaryTranslation("arith.andi"))?; + registry.register_operation::(IntegerBinaryTranslation("arith.ori"))?; + registry.register_operation::(IntegerBinaryTranslation("arith.xori"))?; + registry.register_operation::(ShiftTranslation { right: false })?; + registry.register_operation::(ShiftTranslation { right: true })?; + registry.register_operation::(NegTranslation)?; + registry.register_operation::(NotTranslation)?; + + registry.register_operation::(ComparisonTranslation::Lt)?; + registry.register_operation::(ComparisonTranslation::Le)?; + registry.register_operation::(ComparisonTranslation::Gt)?; + registry.register_operation::(ComparisonTranslation::Ge)?; + registry.register_operation::(ComparisonTranslation::Eq)?; + registry.register_operation::(ComparisonTranslation::Ne)?; + registry.register_operation::(CastTranslation)?; + crate::mir::memory::register_mir_memory_pack(registry)?; + crate::mir::checked::register(registry)?; + crate::mir::shared::register(registry)?; + Ok(()) +} + +/// Preserve the FP16 bit pattern when translating MIR half-precision constants. +struct MirFp16AttributeTranslation; + +impl AttributeTranslation for MirFp16AttributeTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected mir.fp16_attr".to_owned())?; + Ok(Some(MlirAttribute::Float { + value: format!("0x{:04X}", value.0.to_bits()), + ty: MlirType::Float(MlirFloatType::F16), + })) + } +} + +/// Carry a `mir.func`'s `#[inline]` request as the LLVM attribute keyword that +/// `mir-lower` emits for it on the native path. +struct InlineIntentTranslation; + +impl AttributeTranslation for InlineIntentTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let intent = source + .downcast_ref::() + .ok_or_else(|| "expected mir.inline_intent".to_owned())?; + let keyword = match intent { + InlineIntentAttr::Hint => "inlinehint", + InlineIntentAttr::Always => "alwaysinline", + InlineIntentAttr::Never => "noinline", + }; + Ok(Some(MlirAttribute::String(keyword.into()))) + } +} + +/// Forward translated structure to a standard MLIR operation spelling. +struct RenameOperation(#[doc = "Target MLIR operation spelling."] &'static str); + +impl OperationTranslation for RenameOperation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + Ok(vec![renamed(self.0, input)?]) + } +} + +/// Map MIR functions, launch contracts, slice ABI, and enum scratch to executable MLIR. +struct FunctionTranslation; + +impl OperationTranslation for FunctionTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mut target = renamed("func.func", input)?; + move_property(&mut target, ATTR_KEY_SYM_NAME.as_ref(), "sym_name")?; + move_property(&mut target, "mir_func_type", "function_type")?; + let grid_constants = crate::grid_constant::take_parameters(ctx, source, &mut target)?; + lower_alwaysinline_contract(&mut target)?; + lower_inline_intent_contract(&mut target)?; + let mut is_kernel = false; + if let Some(marker) = target.attributes.remove("gpu_kernel") { + match marker { + MlirAttribute::String(value) if value == "true" => { + is_kernel = true; + target + .attributes + .insert("cute.kernel".into(), MlirAttribute::Unit); + target + .attributes + .insert("gpu.kernel".into(), MlirAttribute::Unit); + move_exact_block_contract(&mut target)?; + move_launch_bounds_contract(&mut target)?; + move_cluster_contract(&mut target)?; + flatten_kernel_slice_abi(ctx, source, &mut target, session, &grid_constants)?; + } + other => { + return Err(format!( + "gpu_kernel must be the string `true`, got {other:?}" + )); + } + } + } + lower_helper_linkage(ctx, source, &mut target, is_kernel)?; + crate::mir::checked::hoist_enum_scratch(&mut target)?; + Ok(vec![target]) + } +} + +/// Keep the same public roots as native LLVM export, internalizing only local definitions. +/// +/// The original MIR symbol retains the reserved `#[device]` marker even when +/// target symbol legalization removes it. Kernel entries, device exports, and +/// declarations retain their external ABI. Internal linkage lets LLVM remove +/// zero-size arguments and results from ordinary Rust helper signatures before +/// libNVVM lowers calls to the PTX ABI, which cannot represent those values. +fn lower_helper_linkage( + ctx: &Context, + source: Ptr, + target: &mut MlirOperation, + is_kernel: bool, +) -> Result<(), String> { + for name in ["llvm.linkage", "sym_visibility"] { + if target.attributes.contains_key(name) { + return Err(format!( + "mir.func `{name}` conflicts with source-derived function linkage" + )); + } + } + let function = MirFuncOp::wrap(ctx, source).ok_or("expected mir.func")?; + let source_name = function.get_symbol_name(ctx); + let is_definition = target + .regions + .iter() + .any(|region| !region.blocks.is_empty()); + if is_definition + && !is_kernel + && !reserved_oxide_symbols::is_device_symbol(source_name.as_ref()) + { + target.attributes.insert( + "llvm.linkage".into(), + MlirAttribute::dialect("#llvm.linkage")?, + ); + } + Ok(()) +} + +/// Lower Rust's `#[inline(always)]` marker at the CUTLASS boundary. +/// +/// An arbitrary `alwaysinline = "true"` MLIR attribute survives conversion to +/// `llvm.func`, but LLVM translation does not interpret it. LLVM dialect's +/// `passthrough` spelling carries the explicit Rust request through CUTLASS, +/// including for pointer-bearing helper signatures. +fn lower_alwaysinline_contract(target: &mut MlirOperation) -> Result<(), String> { + let Some(marker) = target.attributes.remove("alwaysinline") else { + return Ok(()); + }; + match marker { + MlirAttribute::String(value) if value == "true" => { + let function_type = target + .properties + .get("function_type") + .ok_or_else(|| "alwaysinline function is missing function_type".to_owned())?; + let MlirAttribute::Type(_) = function_type else { + return Err("alwaysinline function_type is not a type".into()); + }; + if target.attributes.contains_key("passthrough") { + return Err("alwaysinline function already has a passthrough attribute".into()); + } + target.attributes.insert( + "passthrough".into(), + MlirAttribute::Array(vec![MlirAttribute::String("alwaysinline".into())]), + ); + Ok(()) + } + other => Err(format!( + "alwaysinline must be the string `true`, got {other:?}" + )), + } +} + +/// Lower the typed `#[inline]` request recorded on `mir.func` as `inline_intent`. +/// +/// `InlineIntentTranslation` has already turned the request into the keyword +/// `mir-lower` uses on the native path (`inlinehint`, `alwaysinline`, +/// `noinline`). Put it in the LLVM `passthrough` list so CUTLASS hands the +/// same request to LLVM, sharing the list with the older `alwaysinline` +/// marker when both are present. +fn lower_inline_intent_contract(target: &mut MlirOperation) -> Result<(), String> { + let Some(keyword) = target + .attributes + .remove(dialect_mir::ops::function::INLINE_INTENT_ATTR_KEY) + else { + return Ok(()); + }; + let MlirAttribute::String(keyword) = keyword else { + return Err(format!( + "inline_intent must translate to an LLVM attribute keyword, got {keyword:?}" + )); + }; + let entry = MlirAttribute::String(keyword); + match target.attributes.get_mut("passthrough") { + None => { + target + .attributes + .insert("passthrough".into(), MlirAttribute::Array(vec![entry])); + } + Some(MlirAttribute::Array(entries)) => { + if !entries.contains(&entry) { + entries.push(entry); + } + } + Some(other) => { + return Err(format!( + "function passthrough attribute is not an array: {other:?}" + )); + } + } + Ok(()) +} + +/// Move cuda-oxide's three scalar block-contract attributes to the one array +/// attribute consumed by CUTLASS/NVVM. A partial shape is never a valid launch +/// contract, so fail instead of quietly dropping an axis. +fn move_exact_block_contract(target: &mut MlirOperation) -> Result<(), String> { + let dimensions = ["reqntid_x", "reqntid_y", "reqntid_z"] + .map(|name| take_positive_i32_attribute(target, name)) + .into_iter() + .collect::, _>>()?; + + match dimensions.as_slice() { + [None, None, None] => Ok(()), + [Some(x), Some(y), Some(z)] => { + // ptxas rejects .maxntid and .reqntid on the same entry. The exact + // shape is stronger, matching the native exporter's choice. + target.attributes.remove("maxntid"); + target.attributes.insert( + "nvvm.reqntid".into(), + MlirAttribute::DenseI32Array(vec![*x, *y, *z]), + ); + Ok(()) + } + _ => Err("kernel reqntid_x/y/z must be present together".into()), + } +} + +/// Occupancy bounds are semantic NVVM attributes, not opaque source metadata. +/// In particular, minctasm determines the initial register budget needed by +/// setmaxnreg; leaving it unqualified lets ptxas discard register reallocation. +fn move_launch_bounds_contract(target: &mut MlirOperation) -> Result<(), String> { + if let Some(maximum) = take_positive_i32_attribute(target, "maxntid")? { + target.attributes.insert( + "nvvm.maxntid".into(), + MlirAttribute::DenseI32Array(vec![maximum, 1, 1]), + ); + } + if let Some(minimum) = take_positive_i32_attribute(target, "minctasm")? { + target.attributes.insert( + "nvvm.minctasm".into(), + MlirAttribute::Integer { + value: minimum as i128, + ty: MlirType::Integer(32), + }, + ); + } + Ok(()) +} + +/// Preserve the compile-time cluster contract for collective two-CTA MMA. +fn move_cluster_contract(target: &mut MlirOperation) -> Result<(), String> { + let dimensions = ["cluster_dim_x", "cluster_dim_y", "cluster_dim_z"] + .map(|name| take_positive_i32_attribute(target, name)) + .into_iter() + .collect::, _>>()?; + match dimensions.as_slice() { + [None, None, None] => Ok(()), + [Some(x), Some(y), Some(z)] => { + target.attributes.insert( + "nvvm.cluster_dim".into(), + MlirAttribute::DenseI32Array(vec![*x, *y, *z]), + ); + Ok(()) + } + _ => Err("kernel cluster_dim_x/y/z must be present together".into()), + } +} + +/// Remove a named launch attribute and require a positive i32-representable value. +fn take_positive_i32_attribute( + target: &mut MlirOperation, + name: &str, +) -> Result, String> { + let Some(attribute) = target.attributes.remove(name) else { + return Ok(None); + }; + let MlirAttribute::Integer { value, .. } = attribute else { + return Err(format!("{name} must be an integer attribute")); + }; + let value = i32::try_from(value).map_err(|_| format!("{name} does not fit in i32"))?; + if value <= 0 { + return Err(format!("{name} must be positive, got {value}")); + } + Ok(Some(value)) +} + +/// Pass each Rust slice as two launch arguments: pointer, then length. +/// Rebuild the slice at kernel entry so the body can keep using its original value. +/// +/// ```text +/// Rust arguments: slice, descriptor +/// Launch arguments: pointer, len, descriptor +/// Argument index: 0 1 2 +/// ``` +/// +/// Grid-constant attributes must follow the descriptor to its new argument index. +fn flatten_kernel_slice_abi( + ctx: &Context, + source: Ptr, + target: &mut MlirOperation, + session: &mut TranslationSession<'_>, + grid_constants: &std::collections::BTreeMap, +) -> Result<(), String> { + let source_function = + MirFuncOp::wrap(ctx, source).ok_or_else(|| "expected mir.func".to_owned())?; + let source_type = source_function.get_type(ctx); + let source_type_ref = source_type.deref(ctx); + let source_type = type_cast::(&*source_type_ref) + .ok_or_else(|| "mir.func type does not implement FunctionTypeInterface".to_owned())?; + let source_inputs = source_type.arg_types(); + + let function_type = target + .properties + .get("function_type") + .ok_or_else(|| "func.func is missing function_type".to_owned())?; + let MlirAttribute::Type(MlirType::Function { + inputs: target_inputs, + results: target_results, + }) = function_type + else { + return Err("func.func function_type is not an MLIR function type".into()); + }; + if source_inputs.len() != target_inputs.len() { + return Err(format!( + "mir.func has {} source inputs but {} translated inputs", + source_inputs.len(), + target_inputs.len() + )); + } + + let Some(region) = target.regions.first_mut() else { + return Err("kernel func.func has no body region".into()); + }; + let Some(entry) = region.blocks.first_mut() else { + return Err("kernel func.func has no entry block".into()); + }; + if entry.arguments.len() != target_inputs.len() { + return Err(format!( + "kernel entry has {} arguments but function_type has {} inputs", + entry.arguments.len(), + target_inputs.len() + )); + } + + let pointer_type = MlirType::dialect("!llvm.ptr")?; + let length_type = MlirType::Integer(64); + let mut flattened_types = Vec::new(); + let mut flattened_arguments = Vec::new(); + let mut prologue = Vec::new(); + let mut argument_attributes = Vec::new(); + + for (index, ((source_type, target_type), original_argument)) in source_inputs + .into_iter() + .zip(target_inputs.iter()) + .zip(entry.arguments.iter()) + .enumerate() + { + let source_type_ref = source_type.deref(ctx); + let is_plain_slice = source_type_ref.downcast_ref::().is_some(); + let disjoint_slice = source_type_ref.downcast_ref::(); + if let Some(disjoint_slice) = disjoint_slice + && !disjoint_slice.space_types().is_empty() + { + return Err( + "kernel ABI flattening for disjoint slices with runtime index-space fields is not implemented" + .into(), + ); + } + if !is_plain_slice && disjoint_slice.is_none() { + flattened_types.push(target_type.clone()); + flattened_arguments.push(original_argument.clone()); + argument_attributes.push(match grid_constants.get(&index) { + Some(parameter) => parameter.attributes()?, + None => MlirAttribute::Dictionary(Default::default()), + }); + continue; + } + + let expected = MlirType::dialect("!llvm.struct<(!llvm.ptr, i64)>")?; + if target_type != &expected || original_argument.ty != expected { + return Err(format!( + "kernel slice argument translated to {target_type:?}; expected {expected:?}" + )); + } + + let pointer_id = session.fresh_value(); + let length_id = session.fresh_value(); + flattened_types.extend([pointer_type.clone(), length_type.clone()]); + argument_attributes.extend([ + MlirAttribute::Dictionary(Default::default()), + MlirAttribute::Dictionary(Default::default()), + ]); + flattened_arguments.extend([ + MlirBlockArgument { + id: pointer_id, + ty: pointer_type.clone(), + location: original_argument.location.clone(), + }, + MlirBlockArgument { + id: length_id, + ty: length_type.clone(), + location: original_argument.location.clone(), + }, + ]); + + let undef_id = session.fresh_value(); + let partial_id = session.fresh_value(); + let mut undef = MlirOperation::new("llvm.mlir.undef")?; + undef.results.push(MlirResult { + id: undef_id, + ty: expected.clone(), + }); + undef.location = original_argument.location.clone(); + + let mut insert_pointer = MlirOperation::new("llvm.insertvalue")?; + insert_pointer.results.push(MlirResult { + id: partial_id, + ty: expected.clone(), + }); + insert_pointer.operands.extend([ + MlirValueUse { + id: undef_id, + ty: expected.clone(), + }, + MlirValueUse { + id: pointer_id, + ty: pointer_type.clone(), + }, + ]); + insert_pointer + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![0])); + insert_pointer.location = original_argument.location.clone(); + + let mut insert_length = MlirOperation::new("llvm.insertvalue")?; + insert_length.results.push(MlirResult { + // The aggregate argument no longer defines this ID, so the final + // insert can define it without changing any existing body use. + id: original_argument.id, + ty: expected.clone(), + }); + insert_length.operands.extend([ + MlirValueUse { + id: partial_id, + ty: expected.clone(), + }, + MlirValueUse { + id: length_id, + ty: length_type.clone(), + }, + ]); + insert_length + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![1])); + insert_length.location = original_argument.location.clone(); + prologue.extend([undef, insert_pointer, insert_length]); + } + + target.properties.insert( + "function_type".into(), + MlirAttribute::Type(MlirType::Function { + inputs: flattened_types, + results: target_results.clone(), + }), + ); + entry.arguments = flattened_arguments; + prologue.append(&mut entry.operations); + entry.operations = prologue; + if !grid_constants.is_empty() { + if target.properties.contains_key("arg_attrs") + || target.attributes.contains_key("arg_attrs") + { + return Err("grid-constant kernel already has argument attributes".into()); + } + target.properties.insert( + "arg_attrs".into(), + MlirAttribute::Array(argument_attributes), + ); + } + Ok(()) +} + +/// Rebuild conditional branch successor operand segments for cf.cond_br. +struct ConditionalBranchTranslation; + +impl OperationTranslation for ConditionalBranchTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let source = MirCondBranchOp::new(source); + let segment_sizes = source.get_operand_segment_sizes(ctx).0; + if segment_sizes.len() != 3 || segment_sizes[0] != 1 { + return Err(format!( + "mir.cond_br expected [1, true_args, false_args] operand segments, got {segment_sizes:?}" + )); + } + let mut target = renamed("cf.cond_br", input)?; + target.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array( + segment_sizes + .into_iter() + .map(|size| i32::try_from(size).map_err(|_| "operand segment exceeds i32")) + .collect::>()?, + ), + ); + Ok(vec![target]) + } +} + +/// Select integer-bit-preserving or floating-point constant translation. +enum ConstantTranslation { + /// Preserve integer width and bits in a signless arith.constant. + Integer, + /// Preserve the floating-point constant's target format. + Float, +} + +impl OperationTranslation for ConstantTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mut target = renamed("arith.constant", input)?; + let value = match self { + Self::Integer => target + .attributes + .remove("value") + .ok_or_else(|| "mir.constant has no translated `value`".to_owned())?, + Self::Float => { + let keys = ["float_value_f16", "float_value", "float_value_f64"]; + let present = keys + .into_iter() + .filter_map(|key| target.attributes.remove(key).map(|value| (key, value))) + .collect::>(); + if present.len() != 1 { + return Err(format!( + "mir.float_constant expected one value, found {}", + present.len() + )); + } + present.into_iter().next().unwrap().1 + } + }; + target.properties.insert("value".into(), value); + Ok(vec![target]) + } +} + +/// Select integer or floating arithmetic while retaining source signedness where needed. +enum NumericBinaryTranslation { + /// Emit addition in the source integer or floating-point domain. + Add, + /// Emit subtraction in the source integer or floating-point domain. + Sub, + /// Emit multiplication in the source integer or floating-point domain. + Mul, + /// Choose floating, signed-integer, or unsigned-integer division from the source type. + Div, + /// Choose floating, signed-integer, or unsigned-integer remainder from the source type. + Rem, +} + +impl OperationTranslation for NumericBinaryTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let scalar = source_operand_scalar(ctx, source, 0)?; + let name = match (self, scalar) { + (Self::Add, Scalar::Integer { .. }) => "arith.addi", + (Self::Add, Scalar::Float { .. }) => "arith.addf", + (Self::Sub, Scalar::Integer { .. }) => "arith.subi", + (Self::Sub, Scalar::Float { .. }) => "arith.subf", + (Self::Mul, Scalar::Integer { .. }) => "arith.muli", + (Self::Mul, Scalar::Float { .. }) => "arith.mulf", + (Self::Div, Scalar::Integer { signed: true, .. }) => "arith.divsi", + (Self::Div, Scalar::Integer { signed: false, .. }) => "arith.divui", + (Self::Div, Scalar::Float { .. }) => "arith.divf", + (Self::Rem, Scalar::Integer { signed: true, .. }) => "arith.remsi", + (Self::Rem, Scalar::Integer { signed: false, .. }) => "arith.remui", + (Self::Rem, Scalar::Float { .. }) => "arith.remf", + }; + Ok(vec![renamed(name, input)?]) + } +} + +/// Map a bitwise integer operation to its explicit target spelling. +struct IntegerBinaryTranslation(#[doc = "Target bitwise integer operation spelling."] &'static str); + +impl OperationTranslation for IntegerBinaryTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + source_integer(ctx, source, 0)?; + Ok(vec![renamed(self.0, input)?]) + } +} + +/// Match Rust shifts by casting and masking the count before selecting arithmetic or logical shift. +struct ShiftTranslation { + /// Select right shift; otherwise emit a left shift. + right: bool, +} + +impl OperationTranslation for ShiftTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let integer = source_integer(ctx, source, 0)?; + let name = if !self.right { + "arith.shli" + } else if integer.signedness() == Signedness::Signed { + "arith.shrsi" + } else { + "arith.shrui" + }; + let width = integer.width(); + let count = source_integer(ctx, source, 1)?; + let count_width = count.width(); + if !width.is_power_of_two() || width > 128 { + return Err(format!( + "Rust shift requires a power-of-two integer width up to 128, got {width}" + )); + } + // Rust permits a differently typed RHS, and release/wrapping shifts + // mask it by width-1. MLIR/LLVM require equal types and otherwise + // produce poison for oversized counts. Match the native MIR lowering. + let ty = MlirType::Integer(width); + let mut operations = vec![]; + if width != count_width { + let id = session.fresh_value(); + let mut cast = MlirOperation::new(if width > count_width { + "arith.extui" + } else { + "arith.trunci" + })?; + cast.operands.push(input.operands[1].clone()); + cast.results.push(MlirResult { id, ty: ty.clone() }); + cast.location = input.location.clone(); + operations.push(cast); + input.operands[1] = MlirValueUse { id, ty: ty.clone() }; + } + let mask_id = session.fresh_value(); + let mut mask = MlirOperation::new("arith.constant")?; + mask.results.push(MlirResult { + id: mask_id, + ty: ty.clone(), + }); + mask.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: i128::from(width - 1), + ty: ty.clone(), + }, + ); + mask.location = input.location.clone(); + operations.push(mask); + let count_id = session.fresh_value(); + let mut masked = MlirOperation::new("arith.andi")?; + masked.operands = vec![ + input.operands[1].clone(), + MlirValueUse { + id: mask_id, + ty: ty.clone(), + }, + ]; + masked.results.push(MlirResult { + id: count_id, + ty: ty.clone(), + }); + masked.location = input.location.clone(); + operations.push(masked); + input.operands[1] = MlirValueUse { id: count_id, ty }; + operations.push(renamed(name, input)?); + Ok(operations) + } +} + +/// Negate using integer zero subtraction or floating-point negation as appropriate. +struct NegTranslation; + +impl OperationTranslation for NegTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + match source_operand_scalar(ctx, source, 0)? { + Scalar::Float { .. } => Ok(vec![renamed("arith.negf", input)?]), + Scalar::Integer { .. } => { + let result_type = input + .results + .first() + .ok_or_else(|| "mir.neg has no result".to_owned())? + .ty + .clone(); + let zero_id = session.fresh_value(); + let mut zero = MlirOperation::new("arith.constant")?; + zero.results.push(crate::MlirResult { + id: zero_id, + ty: result_type.clone(), + }); + zero.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 0, + ty: result_type.clone(), + }, + ); + zero.location = input.location.clone(); + + let mut neg = renamed("arith.subi", input)?; + neg.operands.insert( + 0, + crate::MlirValueUse { + id: zero_id, + ty: result_type, + }, + ); + Ok(vec![zero, neg]) + } + } + } +} + +/// Complement every bit of the source integer width. +struct NotTranslation; + +impl OperationTranslation for NotTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let integer = source_integer(ctx, source, 0)?; + if integer.width() > 128 { + return Err(format!( + "mir.not wider than 128 bits is not supported yet (got i{})", + integer.width() + )); + } + let result_type = input + .results + .first() + .ok_or_else(|| "mir.not has no result".to_owned())? + .ty + .clone(); + let ones_id = session.fresh_value(); + let mut ones = MlirOperation::new("arith.constant")?; + ones.results.push(crate::MlirResult { + id: ones_id, + ty: result_type.clone(), + }); + ones.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: -1, + ty: result_type.clone(), + }, + ); + ones.location = input.location.clone(); + + let mut not = renamed("arith.xori", input)?; + not.operands.push(crate::MlirValueUse { + id: ones_id, + ty: result_type, + }); + Ok(vec![ones, not]) + } +} + +/// Choose ordered/unordered floating or signed/unsigned integer predicates from source types. +enum ComparisonTranslation { + /// Compare less-than with source signedness or ordered floating semantics. + Lt, + /// Compare less-than-or-equal with source signedness or ordered floating semantics. + Le, + /// Compare greater-than with source signedness or ordered floating semantics. + Gt, + /// Compare greater-than-or-equal with source signedness or ordered floating semantics. + Ge, + /// Compare equality, false for unordered floating operands. + Eq, + /// Compare inequality, true for unordered floating operands. + Ne, +} + +impl OperationTranslation for ComparisonTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let scalar = source_operand_scalar(ctx, source, 0)?; + let (name, predicate) = match scalar { + Scalar::Integer { signed, .. } => { + let predicate = match (self, signed) { + (Self::Eq, _) => 0, + (Self::Ne, _) => 1, + (Self::Lt, true) => 2, + (Self::Le, true) => 3, + (Self::Gt, true) => 4, + (Self::Ge, true) => 5, + (Self::Lt, false) => 6, + (Self::Le, false) => 7, + (Self::Gt, false) => 8, + (Self::Ge, false) => 9, + }; + ("arith.cmpi", predicate) + } + Scalar::Float { .. } => { + // Rust relational comparisons are ordered. `!=` is true for + // unordered values, so it uses UNE rather than ONE. + let predicate = match self { + Self::Eq => 1, + Self::Gt => 2, + Self::Ge => 3, + Self::Lt => 4, + Self::Le => 5, + Self::Ne => 13, + }; + ("arith.cmpf", predicate) + } + }; + let mut target = renamed(name, input)?; + target.properties.insert( + "predicate".into(), + MlirAttribute::Integer { + value: predicate, + ty: MlirType::Integer(64), + }, + ); + Ok(vec![target]) + } +} + +/// Translate scalar, bitwise, and pointer conversions without losing source address semantics. +struct CastTranslation; + +impl OperationTranslation for CastTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let kind = MirCastOp::new(source) + .get_attr_cast_kind(ctx) + .ok_or_else(|| "mir.cast has no cast_kind".to_owned())? + .clone(); + if kind == MirCastKindAttr::PointerCoercionUnsize { + return super::memory::unsize_array_pointer(ctx, source, input, session); + } + if kind == MirCastKindAttr::Transmute + && (super::storage::is_aggregate(ctx, source.deref(ctx).get_operand(0).get_type(ctx)) + || super::storage::is_aggregate(ctx, source.deref(ctx).get_result(0).get_type(ctx))) + { + return super::storage::transmute(ctx, source, input, session); + } + // Rust's pointer address (including ptr::addr's transmute) is generic, + // unlike the explicit NVVM shared-offset instruction boundary. + let pointer_transmute = kind == MirCastKindAttr::Transmute + && source_pointer_address_space(ctx, source, true).is_ok(); + if kind == MirCastKindAttr::PointerExposeAddress || pointer_transmute { + let space = source_pointer_address_space(ctx, source, true)?; + let Scalar::Integer { width, .. } = source_result_scalar(ctx, source, 0)? else { + return Err("pointer address requires an integer result".into()); + }; + if pointer_transmute && width != 64 { + return Err( + "pointer transmute requires the 64-bit Rust address representation".into(), + ); + } + let mut ops = vec![]; + if space != 0 { + let ty = MlirType::dialect("!llvm.ptr")?; + let id = session.fresh_value(); + let mut cast = MlirOperation::new("llvm.addrspacecast")?; + cast.operands = input.operands.clone(); + cast.results.push(MlirResult { id, ty: ty.clone() }); + cast.location = input.location.clone(); + ops.push(cast); + input.operands = vec![MlirValueUse { id, ty }]; + } + ops.push(renamed("llvm.ptrtoint", input)?); + return Ok(ops); + } + let name = match kind { + MirCastKindAttr::IntToInt => { + let source_scalar = source_operand_scalar(ctx, source, 0)?; + let destination_scalar = source_result_scalar(ctx, source, 0)?; + let ( + Scalar::Integer { + width: from, + signed, + .. + }, + Scalar::Integer { width: to, .. }, + ) = (source_scalar, destination_scalar) + else { + return Err("IntToInt cast does not have integer types".into()); + }; + if from < to { + if signed { "arith.extsi" } else { "arith.extui" } + } else if from > to { + "arith.trunci" + } else { + // Pliron signedness has disappeared from the target type; + // this remains a bit-for-bit identity operation. + "arith.bitcast" + } + } + MirCastKindAttr::IntToFloat => match source_operand_scalar(ctx, source, 0)? { + Scalar::Integer { signed: true, .. } => "arith.sitofp", + Scalar::Integer { signed: false, .. } => "arith.uitofp", + _ => return Err("IntToFloat cast does not have an integer source".into()), + }, + MirCastKindAttr::FloatToInt => match source_result_scalar(ctx, source, 0)? { + Scalar::Integer { signed: true, .. } => "arith.fptosi", + Scalar::Integer { signed: false, .. } => "arith.fptoui", + _ => return Err("FloatToInt cast does not have an integer result".into()), + }, + MirCastKindAttr::FloatToFloat => { + let source_scalar = source_operand_scalar(ctx, source, 0)?; + let destination_scalar = source_result_scalar(ctx, source, 0)?; + let (Scalar::Float { width: from }, Scalar::Float { width: to }) = + (source_scalar, destination_scalar) + else { + return Err("FloatToFloat cast does not have float types".into()); + }; + if from < to { + "arith.extf" + } else if from > to { + "arith.truncf" + } else { + "arith.bitcast" + } + } + MirCastKindAttr::Transmute => { + let source_scalar = source_operand_scalar(ctx, source, 0)?; + let destination_scalar = source_result_scalar(ctx, source, 0)?; + match (source_scalar, destination_scalar) { + (Scalar::Integer { width: from, .. }, Scalar::Float { width: to }) + | (Scalar::Float { width: from }, Scalar::Integer { width: to, .. }) + if from == to => + { + "arith.bitcast" + } + (from, to) => { + return Err(format!( + "scalar Transmute requires equal-width integer/float types, got {} bits to {} bits", + from.width(), + to.width() + )); + } + } + } + MirCastKindAttr::PtrToPtr | MirCastKindAttr::PointerCoercionMutToConst => { + let operand_type = source.deref(ctx).get_operand(0).get_type(ctx); + if operand_type.deref(ctx).is::() { + let result_type = source.deref(ctx).get_result(0).get_type(ctx); + if result_type.deref(ctx).is::() { + // Pointer provenance and mutability do not change the + // canonical {generic data pointer, i64 length} layout. + return Ok(vec![renamed("builtin.unrealized_conversion_cast", input)?]); + } + if source_pointer_address_space(ctx, source, false)? != 0 { + return Err( + "slice-to-pointer cast requires a generic pointer result".into() + ); + } + // Rust slice::as_ptr casts the fat pointer to its data + // pointer; extracting field zero must not load through it + // or accidentally reinterpret the length as an address. + let mut extract = renamed("llvm.extractvalue", input)?; + extract + .properties + .insert("position".into(), MlirAttribute::DenseI64Array(vec![0])); + return Ok(vec![extract]); + } + let source_address_space = source_pointer_address_space(ctx, source, true)?; + let result_address_space = source_pointer_address_space(ctx, source, false)?; + if source_address_space == result_address_space { + "builtin.unrealized_conversion_cast" + } else { + "llvm.addrspacecast" + } + } + MirCastKindAttr::PointerExposeAddress => { + source_pointer_address_space(ctx, source, true)?; + source_result_scalar(ctx, source, 0)?; + "llvm.ptrtoint" + } + MirCastKindAttr::PointerWithExposedProvenance => { + source_operand_scalar(ctx, source, 0)?; + source_pointer_address_space(ctx, source, false)?; + "llvm.inttoptr" + } + other => { + return Err(format!( + "cast kind {other:?} needs a pointer or aggregate mapping pack" + )); + } + }; + Ok(vec![renamed(name, input)?]) + } +} + +/// Recover a source MIR pointer's numeric address space before opaque-pointer translation. +fn source_pointer_address_space( + ctx: &Context, + source: Ptr, + operand: bool, +) -> Result { + let ty = if operand { + source.deref(ctx).get_operand(0).get_type(ctx) + } else { + source.deref(ctx).get_result(0).get_type(ctx) + }; + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .map(|pointer| pointer.address_space()) + .ok_or_else(|| "expected a pointer type".to_owned()) +} + +/// Source width and signedness needed to choose scalar MLIR operations. +#[derive(Clone, Copy)] +enum Scalar { + /// An integer's original bit width and signedness before MLIR erases signedness. + Integer { + #[doc = "Integer storage width in bits."] + width: u32, + #[doc = "Whether division, shifts, and comparisons use signed integer semantics."] + signed: bool, + }, + /// A source floating-point scalar with its original bit width. + Float { + #[doc = "Floating-point storage width in bits."] + width: u32, + }, +} + +impl Scalar { + /// Return the source scalar's bit width before signedness is erased. + fn width(self) -> u32 { + match self { + Self::Integer { width, .. } | Self::Float { width } => width, + } + } +} + +/// Read scalar classification from a selected source operand. +fn source_operand_scalar( + ctx: &Context, + source: Ptr, + index: usize, +) -> Result { + let ty = source.deref(ctx).get_operand(index).get_type(ctx); + scalar(ctx, ty) +} + +/// Read scalar classification from a selected source result. +fn source_result_scalar( + ctx: &Context, + source: Ptr, + index: usize, +) -> Result { + let ty = source.deref(ctx).get_result(index).get_type(ctx); + scalar(ctx, ty) +} + +/// Classify reviewed integer and floating-point source types or reject unsupported scalars. +fn scalar(ctx: &Context, ty: TypeHandle) -> Result { + let ty_ref = ty.deref(ctx); + if let Some(integer) = ty_ref.downcast_ref::() { + return Ok(Scalar::Integer { + width: integer.width(), + // Under the shared integer contract, signless integers (notably bool) take + // the unsigned path for operations where signedness matters. + signed: integer.signedness() == Signedness::Signed, + }); + } + if !type_impls::(&*ty_ref) { + return Err("expected a scalar integer or floating-point type".into()); + } + let width = if ty_ref.is::() || ty_ref.is::() { + 16 + } else if ty_ref.is::() { + 32 + } else if ty_ref.is::() { + 64 + } else { + return Err("unsupported floating-point type".into()); + }; + Ok(Scalar::Float { width }) +} + +/// Require an integer source value and retain its bit width and signedness. +fn source_integer( + ctx: &Context, + source: Ptr, + index: usize, +) -> Result { + let ty = source.deref(ctx).get_operand(index).get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .cloned() + .ok_or_else(|| "expected an integer operand".to_owned()) +} + +/// Retain translated operation structure while changing its target operation name. +fn renamed(name: &str, input: OperationInput) -> Result { + let mut target = MlirOperation::new(name)?; + target.results = input.results; + target.operands = input.operands; + target.successors = input.successors; + target.regions = input.regions; + target.attributes = input.attributes; + target.location = input.location; + Ok(target) +} + +/// Remove an optional source attribute and reattach it under the target property name. +fn move_property( + operation: &mut MlirOperation, + source_name: &str, + target_name: &str, +) -> Result<(), String> { + let value = operation + .attributes + .remove(source_name) + .ok_or_else(|| format!("missing required `{source_name}` property"))?; + operation.properties.insert(target_name.into(), value); + Ok(()) +} + +#[cfg(test)] +mod tests { + //! Scalar, launch-contract, and control-flow mapping tests using ordinary MIR fixtures. + + use std::num::NonZero; + + /// Reject incomplete cluster shapes and dimensions that are zero, negative, or too large. + #[test] + fn cluster_contract_rejects_partial_and_nonpositive_dimensions() { + use crate::{MlirAttribute, MlirOperation, MlirType}; + let mut kernel = MlirOperation::new("func.func").unwrap(); + kernel.attributes.insert( + "cluster_dim_x".into(), + MlirAttribute::Integer { + value: 2, + ty: MlirType::Integer(32), + }, + ); + assert!( + super::move_cluster_contract(&mut kernel) + .unwrap_err() + .contains("present together") + ); + kernel.attributes.insert( + "cluster_dim_x".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(32), + }, + ); + assert!( + super::move_cluster_contract(&mut kernel) + .unwrap_err() + .contains("positive") + ); + } + + use crate::{ + CutlassFullCuteMlir22, MlirConsumerProfile, + profile::render_mapping_module_without_cutlass_envelope, + }; + use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAddOp, MirCastOp, MirCondBranchOp, MirConstantOp, MirDivOp, MirFloatConstantOp, + MirFuncOp, MirGotoOp, MirLtOp, MirReturnOp, + }, + types::MirPtrType, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{FPSingleAttr, IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{ + OperandSegmentInterface, SingleBlockRegionInterface, SymbolOpInterface, + }, + ops::ModuleOp, + types::{FP32Type, FP64Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::{TypeHandle, TypedHandle}, + utils::apint::APInt, + value::Value, + }; + + use dialect_mir::attributes::InlineIntentAttr; + + /// Create a source context with the MIR dialect needed by scalar translation tests. + fn context() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + ctx + } + + /// Create a typed MIR function with an entry block and append it to the test module. + fn append_function( + ctx: &mut Context, + module: &ModuleOp, + name: &str, + inputs: Vec, + results: Vec, + ) -> (Ptr, pliron::context::Ptr) { + let function_type = FunctionType::get(ctx, inputs, results); + let operation = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(ctx, operation, TypeAttr::new(function_type.into())); + function.set_symbol_name(ctx, Identifier::try_from(name).unwrap()); + module.append_operation(ctx, operation, 0); + let region = operation.deref(ctx).get_region(0); + (operation, region) + } + + /// Construct a source operation of the requested concrete kind and append it to the block. + fn operation( + ctx: &mut Context, + block: Ptr, + results: Vec, + operands: Vec, + ) -> Ptr { + let operation = + Operation::new(ctx, O::get_concrete_op_info(), results, operands, vec![], 0); + operation.insert_at_back(block, ctx); + operation + } + + /// Export the scalar fixture and render only the operation-mapping module. + fn translate(ctx: &mut Context, module: &ModuleOp) -> String { + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(ctx, module).unwrap(); + render_mapping_module_without_cutlass_envelope( + &target, + module.get_symbol_name(ctx).as_ref(), + ) + } + + /// Obtain a source integer type with explicit width and signedness. + fn integer(ctx: &mut Context, width: u32, signedness: Signedness) -> TypedHandle { + IntegerType::get(ctx, width, signedness) + } + + /// Check Rust shift counts are converted and masked before MLIR receives equal-width operands. + #[test] + fn shifts_cast_and_mask_counts_before_emitting_equal_width_operands() { + use dialect_mir::ops::{MirShlOp, MirShrOp}; + for (width, count_width, signed, right, name, cast) in [ + (64, 32, false, true, "arith.shrui", Some("arith.extui")), + (16, 32, false, false, "arith.shli", Some("arith.trunci")), + (32, 64, true, true, "arith.shrsi", Some("arith.trunci")), + (32, 32, false, false, "arith.shli", None), + ] { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("shift").unwrap()); + let lhs = integer( + &mut ctx, + width, + if signed { + Signedness::Signed + } else { + Signedness::Unsigned + }, + ); + let rhs = integer(&mut ctx, count_width, Signedness::Unsigned); + let (_, region) = append_function( + &mut ctx, + &module, + "shift", + vec![lhs.into(), rhs.into()], + vec![lhs.into()], + ); + let entry = BasicBlock::new(&mut ctx, None, vec![lhs.into(), rhs.into()]); + entry.insert_at_back(region, &ctx); + let args = entry.deref(&ctx).arguments().collect(); + let shift = if right { + operation::(&mut ctx, entry, vec![lhs.into()], args) + } else { + operation::(&mut ctx, entry, vec![lhs.into()], args) + }; + let result = shift.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![result]); + let text = translate(&mut ctx, &module); + let shift_line = text + .lines() + .find(|line| line.contains(&format!("\"{name}\""))) + .unwrap(); + assert!( + shift_line.contains(&format!(": (i{width}, i{width}) -> (i{width})")), + "{text}" + ); + assert!( + text.contains(&format!("value = {} : i{width}", width - 1)), + "{text}" + ); + assert!(text.contains("\"arith.andi\""), "{text}"); + if let Some(cast) = cast { + assert!(text.contains(&format!("\"{cast}\"")), "{text}"); + } + } + } + + /// Check signed and unsigned arithmetic select distinct target instructions after type signedness erasure. + #[test] + fn scalar_arithmetic_keeps_signedness_in_the_operation_names() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("core").unwrap()); + let u64_type = integer(&mut ctx, 64, Signedness::Unsigned); + let i64_type = integer(&mut ctx, 64, Signedness::Signed); + let i1_type = integer(&mut ctx, 1, Signedness::Signless); + let f32_type = FP32Type::get(&ctx); + let f64_type = FP64Type::get(&ctx); + + let (_, region) = append_function( + &mut ctx, + &module, + "scalar", + vec![ + u64_type.into(), + u64_type.into(), + i64_type.into(), + i64_type.into(), + f32_type.into(), + ], + vec![ + u64_type.into(), + i64_type.into(), + i1_type.into(), + i1_type.into(), + f32_type.into(), + f64_type.into(), + ], + ); + let entry = BasicBlock::new( + &mut ctx, + None, + vec![ + u64_type.into(), + u64_type.into(), + i64_type.into(), + i64_type.into(), + f32_type.into(), + ], + ); + entry.insert_at_back(region, &ctx); + let arguments = entry.deref(&ctx).arguments().collect::>(); + + let four = operation::(&mut ctx, entry, vec![u64_type.into()], vec![]); + MirConstantOp::new(four).set_attr_value( + &ctx, + IntegerAttr::new(u64_type, APInt::from_u128(4, NonZero::new(64).unwrap())), + ); + let four_value = four.deref(&ctx).get_result(0); + + let unsigned_div = operation::( + &mut ctx, + entry, + vec![u64_type.into()], + vec![arguments[0], four_value], + ); + let signed_div = operation::( + &mut ctx, + entry, + vec![i64_type.into()], + vec![arguments[2], arguments[3]], + ); + let unsigned_lt = operation::( + &mut ctx, + entry, + vec![i1_type.into()], + vec![arguments[0], arguments[1]], + ); + let signed_lt = operation::( + &mut ctx, + entry, + vec![i1_type.into()], + vec![arguments[2], arguments[3]], + ); + + let one = operation::(&mut ctx, entry, vec![f32_type.into()], vec![]); + MirFloatConstantOp::new(one).set_attr_float_value(&ctx, FPSingleAttr::from(1.0)); + let one_value = one.deref(&ctx).get_result(0); + let sum = operation::( + &mut ctx, + entry, + vec![f32_type.into()], + vec![arguments[4], one_value], + ); + let cast = + operation::(&mut ctx, entry, vec![f64_type.into()], vec![arguments[4]]); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::FloatToFloat); + + let returns = vec![ + unsigned_div.deref(&ctx).get_result(0), + signed_div.deref(&ctx).get_result(0), + unsigned_lt.deref(&ctx).get_result(0), + signed_lt.deref(&ctx).get_result(0), + sum.deref(&ctx).get_result(0), + cast.deref(&ctx).get_result(0), + ]; + operation::(&mut ctx, entry, vec![], returns); + + let text = translate(&mut ctx, &module); + let expected = r#""builtin.module"() <{sym_name = "core"}> ({ + ^bb0: + "func.func"() <{function_type = (i64, i64, i64, i64, f32) -> (i64, i64, i1, i1, f32, f64), sym_name = "scalar"}> ({ + ^bb1(%v0: i64, %v1: i64, %v2: i64, %v3: i64, %v4: f32): + %v5 = "arith.constant"() <{value = 4 : i64}> : () -> (i64) + %v6 = "arith.divui"(%v0, %v5) : (i64, i64) -> (i64) + %v7 = "arith.divsi"(%v2, %v3) : (i64, i64) -> (i64) + %v8 = "arith.cmpi"(%v0, %v1) <{predicate = 6 : i64}> : (i64, i64) -> (i1) + %v9 = "arith.cmpi"(%v2, %v3) <{predicate = 2 : i64}> : (i64, i64) -> (i1) + %v10 = "arith.constant"() <{value = 0x3F800000 : f32}> : () -> (f32) + %v11 = "arith.addf"(%v4, %v10) : (f32, f32) -> (f32) + %v12 = "arith.extf"(%v4) : (f32) -> (f64) + "func.return"(%v6, %v7, %v8, %v9, %v11, %v12) : (i64, i64, i1, i1, f32, f64) -> () + }) {llvm.linkage = #llvm.linkage} : () -> () +}) : () -> () +"#; + assert_eq!( + crate::profile::normalize_mlir_identifiers(&text), + crate::profile::normalize_mlir_identifiers(expected) + ); + } + + /// Check same-width scalar transmutation preserves bits rather than performing numeric conversion. + #[test] + fn equal_width_integer_float_transmute_maps_to_bitcast() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("transmute").unwrap()); + let u32_type = integer(&mut ctx, 32, Signedness::Unsigned); + let f32_type = FP32Type::get(&ctx); + let (_, region) = append_function( + &mut ctx, + &module, + "bits_to_float", + vec![u32_type.into()], + vec![f32_type.into()], + ); + let entry = BasicBlock::new(&mut ctx, None, vec![u32_type.into()]); + entry.insert_at_back(region, &ctx); + let input = entry.deref(&ctx).arguments().next().unwrap(); + let cast = operation::(&mut ctx, entry, vec![f32_type.into()], vec![input]); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let cast_result = cast.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![cast_result]); + + let text = translate(&mut ctx, &module); + assert!( + text.contains(r#""arith.bitcast"(%v0) : (i32) -> (f32)"#), + "{text}" + ); + } + + /// Check FP16 bit extraction and reconstruction use exact bit-preserving scalar operations. + #[test] + fn rust_f16_to_bits_and_from_bits_preserve_the_scalar_bit_pattern() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("half_bits").unwrap()); + let half: TypeHandle = dialect_mir::types::MirFP16Type::get(&ctx).into(); + let bits: TypeHandle = integer(&mut ctx, 16, Signedness::Unsigned).into(); + let (_, region) = + append_function(&mut ctx, &module, "half_bits", vec![half], vec![bits, half]); + let entry = BasicBlock::new(&mut ctx, None, vec![half]); + entry.insert_at_back(region, &ctx); + let input = entry.deref(&ctx).get_argument(0); + let to_bits = operation::(&mut ctx, entry, vec![bits], vec![input]); + MirCastOp::new(to_bits).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let bits_value = to_bits.deref(&ctx).get_result(0); + let from_bits = operation::(&mut ctx, entry, vec![half], vec![bits_value]); + MirCastOp::new(from_bits).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let half_value = from_bits.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![bits_value, half_value]); + let text = translate(&mut ctx, &module); + assert!( + text.contains(r#""arith.bitcast"(%v0) : (f16) -> (i16)"#), + "{text}" + ); + assert!( + text.contains(r#""arith.bitcast"(%v1) : (i16) -> (f16)"#), + "{text}" + ); + } + + /// Check slice metadata remains intact until the source operation requests only its pointer. + #[test] + // Reproduce the importer's Rust reference-to-raw-slice and as_ptr boundaries. + #[allow(clippy::disallowed_methods)] + fn slice_as_ptr_preserves_metadata_until_extracting_the_data_pointer() { + use dialect_mir::{ + attributes::MirPointerKindAuthorityAttr, + types::{MirFP16Type, MirPointerKind, MirSliceType}, + }; + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("slice_pointer").unwrap()); + let half: TypeHandle = MirFP16Type::get(&ctx).into(); + let slice: TypeHandle = + MirSliceType::get_with_kind(&mut ctx, half, MirPointerKind::SharedRef).into(); + let raw_slice: TypeHandle = + MirSliceType::get_with_kind(&mut ctx, half, MirPointerKind::RawConst).into(); + let pointer: TypeHandle = + MirPtrType::get_with_kind(&mut ctx, half, false, 0, MirPointerKind::RawConst).into(); + let (_, region) = append_function( + &mut ctx, + &module, + "slice_pointer", + vec![slice], + vec![raw_slice, pointer], + ); + let entry = BasicBlock::new(&mut ctx, None, vec![slice]); + entry.insert_at_back(region, &ctx); + let input = entry.deref(&ctx).get_argument(0); + let raw = operation::(&mut ctx, entry, vec![raw_slice], vec![input]); + MirCastOp::new(raw).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + MirCastOp::new(raw) + .set_attr_pointer_kind_authority(&ctx, MirPointerKindAuthorityAttr::RawAddress); + let raw_value = raw.deref(&ctx).get_result(0); + let data = operation::(&mut ctx, entry, vec![pointer], vec![raw_value]); + MirCastOp::new(data).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + MirCastOp::new(data) + .set_attr_pointer_kind_authority(&ctx, MirPointerKindAuthorityAttr::RustCast); + let pointer_value = data.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![raw_value, pointer_value]); + let text = translate(&mut ctx, &module); + assert!(text.contains(r#""builtin.unrealized_conversion_cast"(%v0) : (!llvm.struct<(!llvm.ptr, i64)>) -> (!llvm.struct<(!llvm.ptr, i64)>)"#), "{text}"); + assert!(text.contains(r#""llvm.extractvalue"(%v1) <{position = array}> : (!llvm.struct<(!llvm.ptr, i64)>) -> (!llvm.ptr)"#), "{text}"); + assert!(!text.contains("llvm.load"), "{text}"); + assert!(!text.contains("llvm.ptrtoint"), "{text}"); + } + + /// Preserve native public roots together while internalizing only local definitions. + #[test] + fn helper_linkage_preserves_native_public_roots_and_declarations() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("linkage").unwrap()); + let cases = [ + ("helper".to_owned(), false, true, true), + ("kernel".to_owned(), true, true, false), + ( + reserved_oxide_symbols::device_symbol("exported"), + false, + true, + false, + ), + ( + format!("{}legacy", reserved_oxide_symbols::LEGACY_DEVICE_PREFIX), + false, + true, + false, + ), + ("external_declaration".to_owned(), false, false, false), + ( + reserved_oxide_symbols::device_extern_symbol("external"), + false, + false, + false, + ), + ]; + for (name, kernel, definition, _) in &cases { + let (function, region) = append_function(&mut ctx, &module, name, vec![], vec![]); + if *kernel { + function.deref_mut(&ctx).attributes.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + } + if *definition { + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + } + } + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let functions = + &target.root.regions[0].blocks[0].operations[0].regions[0].blocks[0].operations; + assert_eq!(functions.len(), cases.len()); + for (function, (_, _, _, internal)) in functions.iter().zip(&cases) { + assert_eq!( + function.attributes.contains_key("llvm.linkage"), + *internal, + "{function:?}" + ); + } + } + + /// A zero-size value keeps its Rust type while its local helper gains LLVM linkage. + #[test] + fn helper_linkage_internalizes_zero_size_signatures() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("zst_linkage").unwrap()); + let unit: TypeHandle = dialect_mir::types::MirTupleType::get_with_layout( + &mut ctx, + vec![], + vec![], + vec![], + 0, + 1, + ) + .into(); + let (_, region) = append_function(&mut ctx, &module, "zst_helper", vec![unit], vec![unit]); + let entry = BasicBlock::new(&mut ctx, None, vec![unit]); + entry.insert_at_back(region, &ctx); + let value = entry.deref(&ctx).get_argument(0); + operation::(&mut ctx, entry, vec![], vec![value]); + let text = translate(&mut ctx, &module); + assert!( + text.contains("llvm.linkage = #llvm.linkage"), + "{text}" + ); + assert!( + text.contains("function_type = (!llvm.struct<()>) -> (!llvm.struct<()>)"), + "{text}" + ); + } + + /// Reject unmodeled linkage overrides and retain malformed kernel-marker errors. + #[test] + fn helper_linkage_rejects_conflicting_metadata() { + for (name, value) in [("sym_visibility", "private"), ("gpu_kernel", "false")] { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("bad_linkage").unwrap()); + let (function, region) = append_function(&mut ctx, &module, "helper", vec![], vec![]); + function.deref_mut(&ctx).attributes.set( + Identifier::try_from(name).unwrap(), + StringAttr::new(value.into()), + ); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + let result = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&ctx, &module); + assert!(result.is_err(), "accepted conflicting {name}"); + } + } + + /// Check explicit inline(always) survives lowering for a scalar-only helper. + #[test] + fn alwaysinline_pointer_free_helper_uses_llvm_passthrough() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("inline").unwrap()); + let (function, region) = append_function(&mut ctx, &module, "helper", vec![], vec![]); + function.deref_mut(&ctx).attributes.set( + Identifier::try_from("alwaysinline").unwrap(), + StringAttr::new("true".into()), + ); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + + let text = translate(&mut ctx, &module); + assert!(text.contains(r#"passthrough = ["alwaysinline"]"#), "{text}"); + assert!(!text.contains(r#"alwaysinline = "true""#), "{text}"); + } + + /// Check explicit inline(always) also survives lowering for a pointer-bearing helper. + #[test] + fn alwaysinline_pointer_helper_uses_llvm_passthrough() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("inline_ptr").unwrap()); + let f32_type: TypeHandle = FP32Type::get(&ctx).into(); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, f32_type, false).into(); + let (function, region) = + append_function(&mut ctx, &module, "pointer_helper", vec![pointer], vec![]); + function.deref_mut(&ctx).attributes.set( + Identifier::try_from("alwaysinline").unwrap(), + StringAttr::new("true".into()), + ); + let entry = BasicBlock::new(&mut ctx, None, vec![pointer]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + + let text = translate(&mut ctx, &module); + assert!(text.contains(r#"passthrough = ["alwaysinline"]"#), "{text}"); + assert!(!text.contains(r#"alwaysinline = "true""#), "{text}"); + } + + /// Check each typed `#[inline]` request lowers to the LLVM passthrough keyword + /// the native path uses, and that the request itself does not reach the text. + #[test] + fn inline_intent_uses_llvm_passthrough() { + for (intent, keyword) in [ + (InlineIntentAttr::Hint, "inlinehint"), + (InlineIntentAttr::Always, "alwaysinline"), + (InlineIntentAttr::Never, "noinline"), + ] { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("inline_request").unwrap()); + let (function, region) = append_function(&mut ctx, &module, "helper", vec![], vec![]); + MirFuncOp::wrap(&ctx, function) + .unwrap() + .set_inline_intent(&mut ctx, intent); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + + let text = translate(&mut ctx, &module); + assert!( + text.contains(&format!(r#"passthrough = ["{keyword}"]"#)), + "{intent:?}: {text}" + ); + assert!(!text.contains("inline_intent"), "{intent:?}: {text}"); + } + } + + /// Check the typed request and the older string marker produce one keyword, not two. + #[test] + fn inline_intent_shares_the_passthrough_list_with_the_alwaysinline_marker() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("inline_both").unwrap()); + let (function, region) = append_function(&mut ctx, &module, "helper", vec![], vec![]); + function.deref_mut(&ctx).attributes.set( + Identifier::try_from("alwaysinline").unwrap(), + StringAttr::new("true".into()), + ); + MirFuncOp::wrap(&ctx, function) + .unwrap() + .set_inline_intent(&mut ctx, InlineIntentAttr::Always); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + operation::(&mut ctx, entry, vec![], vec![]); + + let text = translate(&mut ctx, &module); + assert!(text.contains(r#"passthrough = ["alwaysinline"]"#), "{text}"); + assert_eq!(text.matches("alwaysinline").count(), 1, "{text}"); + } + + /// Check an unstructured MIR branch graph retains successor edges and operand segment sizes. + #[test] + fn arbitrary_mir_cfg_maps_to_cf_with_explicit_segments() { + let mut ctx = context(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("cfg").unwrap()); + let i1_type = integer(&mut ctx, 1, Signedness::Signless); + let u64_type = integer(&mut ctx, 64, Signedness::Unsigned); + let (_, region) = append_function( + &mut ctx, + &module, + "choose", + vec![i1_type.into(), u64_type.into()], + vec![u64_type.into()], + ); + let entry = BasicBlock::new(&mut ctx, None, vec![i1_type.into(), u64_type.into()]); + let on_true = BasicBlock::new(&mut ctx, None, vec![u64_type.into()]); + let on_false = BasicBlock::new(&mut ctx, None, vec![u64_type.into()]); + entry.insert_at_back(region, &ctx); + on_true.insert_at_back(region, &ctx); + on_false.insert_at_back(region, &ctx); + + let entry_arguments = entry.deref(&ctx).arguments().collect::>(); + let (operands, segments) = MirCondBranchOp::compute_segment_sizes(vec![ + vec![entry_arguments[0]], + vec![entry_arguments[1]], + vec![entry_arguments[1]], + ]); + let branch = Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + operands, + vec![on_true, on_false], + 0, + ); + MirCondBranchOp::new(branch).set_operand_segment_sizes(&ctx, segments); + branch.insert_at_back(entry, &ctx); + + let true_value = on_true.deref(&ctx).get_argument(0); + operation::(&mut ctx, on_true, vec![], vec![true_value]); + let false_value = on_false.deref(&ctx).get_argument(0); + let goto = Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![false_value], + vec![on_true], + 0, + ); + goto.insert_at_back(on_false, &ctx); + + let text = translate(&mut ctx, &module); + let expected = r#""builtin.module"() <{sym_name = "cfg"}> ({ + ^bb0: + "func.func"() <{function_type = (i1, i64) -> (i64), sym_name = "choose"}> ({ + ^bb1(%v0: i1, %v1: i64): + "cf.cond_br"(%v0, %v1, %v1) [^bb2, ^bb3] <{operandSegmentSizes = array}> : (i1, i64, i64) -> () + ^bb2(%v2: i64): + "func.return"(%v2) : (i64) -> () + ^bb3(%v3: i64): + "cf.br"(%v3) [^bb2] : (i64) -> () + }) {llvm.linkage = #llvm.linkage} : () -> () +}) : () -> () +"#; + assert_eq!( + crate::profile::normalize_mlir_identifiers(&text), + crate::profile::normalize_mlir_identifiers(expected) + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs new file mode 100644 index 0000000000..a0a5ea5c3a --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs @@ -0,0 +1,2068 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! MIR aggregates, pointers, calls, and memory in standard MLIR dialects. +//! +//! MIR pointers become opaque LLVM pointers. Rust slices keep their two-word +//! ABI as an LLVM struct, and fixed arrays and structs keep their physical +//! field order. Operations then use LLVM's aggregate and memory operations. + +use crate::{ + DropAttribute, MlirAttribute, MlirFloatType, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use dialect_mir::{ + attributes::{CompilerResultBundleAttr, FieldIndexAttr, MutabilityAttr}, + ops::{ + MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirCallOp, MirCastOp, MirConstructArrayOp, + MirConstructSliceOp, MirConstructStructOp, MirConstructTupleOp, MirExtractArrayElementOp, + MirExtractFieldOp, MirFieldAddrOp, MirLoadOp, MirPtrOffsetOp, MirRefOp, MirStorageDeadOp, + MirStorageLiveOp, MirStoreOp, + }, + types::{ + MirArrayType, MirDisjointSliceType, MirEnumType, MirFP16Type, MirPtrType, MirSliceType, + MirStructType, MirTupleType, MirUnionType, + }, +}; +use llvm_export::ops::LocalMemoryProvenanceAttr; +use pliron::{ + builtin::types::{FP16Type, FP32Type, FP64Type, IntegerType, UnitType}, + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Register reviewed pointer, physical aggregate, call, and memory mappings. +pub(crate) fn register_mir_memory_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_type::(PointerTypeTranslation)?; + registry.register_type::(SliceTypeTranslation)?; + registry.register_type::(DisjointSliceTypeTranslation)?; + registry.register_type::(ArrayTypeTranslation)?; + registry.register_type::(StructTypeTranslation)?; + registry.register_type::(TupleTypeTranslation)?; + super::storage::register(registry)?; + + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + // This marks an importer-created ABI aggregate for optional forwarding. + // If it survives preparation, preserve the aggregate normally. + registry.register_attribute::(DropAttribute)?; + + registry.register_operation::(AllocaTranslation)?; + registry.register_operation::(ReferenceTranslation)?; + registry.register_operation::(LoadTranslation)?; + registry.register_operation::(StoreTranslation)?; + registry.register_operation::(PointerOffsetTranslation)?; + registry.register_operation::(ExtractFieldTranslation)?; + registry.register_operation::(ConstructStructTranslation)?; + registry.register_operation::(ConstructStructTranslation)?; + registry.register_operation::(ConstructArrayTranslation)?; + registry.register_operation::(ConstructSliceTranslation)?; + registry.register_operation::(FieldAddressTranslation)?; + registry.register_operation::(ArrayElementAddressTranslation)?; + registry.register_operation::(ExtractArrayElementTranslation)?; + registry.register_operation::(CallTranslation)?; + registry.register_operation::(AssertTranslation)?; + registry.register_operation::(DropOperation)?; + registry.register_operation::(DropOperation)?; + Ok(()) +} + +/// Erase a MIR pointee while preserving its LLVM address space. +struct PointerTypeTranslation; + +impl TypeTranslation for PointerTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let pointer = source_ref + .downcast_ref::() + .ok_or_else(|| "expected mir.ptr".to_owned())?; + llvm_pointer_type(pointer.address_space()) + } +} + +/// Represent a Rust slice as a data pointer plus 64-bit element count. +struct SliceTypeTranslation; + +impl TypeTranslation for SliceTypeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + llvm_struct_type(&[llvm_pointer_type(0)?, MlirType::Integer(64)]) + } +} + +/// Preserve the disjoint slice's runtime pointer and extent carrier. +struct DisjointSliceTypeTranslation; + +impl TypeTranslation for DisjointSliceTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let slice = source_ref + .downcast_ref::() + .ok_or_else(|| "expected mir.disjoint_slice".to_owned())?; + let mut fields = vec![llvm_pointer_type(0)?, MlirType::Integer(64)]; + for &field in slice.space_types() { + fields.push(registry.translate_type(ctx, field)?); + } + llvm_struct_type(&fields) + } +} + +/// Translate fixed-length array storage with its element type and length. +struct ArrayTypeTranslation; + +impl TypeTranslation for ArrayTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let array = source_ref + .downcast_ref::() + .ok_or_else(|| "expected mir.array".to_owned())?; + llvm_array_type( + array.size(), + ®istry.translate_type(ctx, array.element_type())?, + ) + } +} + +/// Build an LLVM struct with verified Rust field order, offsets, and padding. +struct StructTypeTranslation; + +impl TypeTranslation for StructTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + aggregate_plan(ctx, source, |field| registry.translate_type(ctx, field)).map(|plan| plan.ty) + } +} + +/// Build an LLVM struct retaining the tuple's physical Rust layout. +struct TupleTypeTranslation; + +impl TypeTranslation for TupleTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + aggregate_plan(ctx, source, |field| registry.translate_type(ctx, field)).map(|plan| plan.ty) + } +} + +/// Allocate local storage with source element type, count, and explicit alignment. +struct AllocaTranslation; + +impl OperationTranslation for AllocaTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.alloca", &input)?; + let result = only_result("mir.alloca", &input)?; + let pointee = pointer_pointee(ctx, source.deref(ctx).get_result(0).get_type(ctx))?; + let target_pointee = session + .translate_type(pointee) + .map_err(|error| error.to_string())?; + + let count_id = session.fresh_value(); + let mut count = MlirOperation::new("arith.constant")?; + count.results.push(MlirResult { + id: count_id, + ty: MlirType::Integer(32), + }); + count.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(32), + }, + ); + count.location = input.location.clone(); + + let mut alloca = MlirOperation::new("llvm.alloca")?; + alloca.results.push(result); + alloca.operands.push(MlirValueUse { + id: count_id, + ty: MlirType::Integer(32), + }); + alloca + .properties + .insert("elem_type".into(), MlirAttribute::Type(target_pointee)); + if let Some(alignment) = source_abi_align(ctx, pointee) { + alloca.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: alignment.into(), + ty: MlirType::Integer(64), + }, + ); + } + alloca.location = input.location; + Ok(vec![count, alloca]) + } +} + +/// Give an SSA value fresh storage at the reference operation's control-flow position. +/// +/// These addresses can escape to helper calls or be compared. Unlike private +/// packing scratch, their allocation must not be hoisted or shared. +struct ReferenceTranslation; + +impl OperationTranslation for ReferenceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + MirRefOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ensure_empty_attributes("mir.ref", &input)?; + if input.operands.len() != 1 || !input.regions.is_empty() || !input.successors.is_empty() { + return Err("mir.ref requires one SSA value and one pointer result".into()); + } + let result = only_result("mir.ref", &input)?; + let value_type = source.deref(ctx).get_operand(0).get_type(ctx); + super::checked::check_payload_storage(ctx, value_type)?; + let alignment = source_abi_align(ctx, value_type) + .ok_or("mir.ref requires a known source ABI alignment")?; + let count_id = session.fresh_value(); + let mut count = MlirOperation::new("arith.constant")?; + count.results.push(MlirResult { + id: count_id, + ty: MlirType::Integer(32), + }); + count.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(32), + }, + ); + count.location = input.location.clone(); + + let pointer = MlirValueUse { + id: result.id, + ty: result.ty.clone(), + }; + let mut allocation = MlirOperation::new("llvm.alloca")?; + allocation.operands.push(MlirValueUse { + id: count_id, + ty: MlirType::Integer(32), + }); + allocation.results.push(result); + allocation.properties.insert( + "elem_type".into(), + MlirAttribute::Type(input.operands[0].ty.clone()), + ); + add_alignment(&mut allocation, Some(alignment)); + allocation.location = input.location.clone(); + + let mut store = MlirOperation::new("llvm.store")?; + store.operands = vec![input.operands[0].clone(), pointer]; + add_alignment(&mut store, Some(alignment)); + memory_ordering(&mut store); + store.location = input.location; + Ok(vec![count, allocation, store]) + } +} + +/// Load a verified pointee while preserving the source memory access contract. +struct LoadTranslation; + +impl OperationTranslation for LoadTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + input.attributes.remove("mir_load_volatile"); + ensure_empty_attributes("mir.load", &input)?; + let mut target = renamed("llvm.load", input)?; + memory_ordering(&mut target); + if MirLoadOp::new(source).is_volatile(ctx) { + target + .properties + .insert("volatile_".into(), MlirAttribute::Unit); + } + let pointee = pointer_pointee(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + add_alignment(&mut target, source_abi_align(ctx, pointee)); + Ok(vec![target]) + } +} + +/// Store a value to its verified pointee with explicit alignment and memory ordering. +struct StoreTranslation; + +impl OperationTranslation for StoreTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + input.attributes.remove("mir_store_volatile"); + ensure_empty_attributes("mir.store", &input)?; + if input.operands.len() != 2 { + return Err(format!( + "mir.store expected two operands, got {}", + input.operands.len() + )); + } + input.operands.swap(0, 1); + let mut target = renamed("llvm.store", input)?; + memory_ordering(&mut target); + if MirStoreOp::new(source).is_volatile(ctx) { + target + .properties + .insert("volatile_".into(), MlirAttribute::Unit); + } + let pointee = pointer_pointee(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + add_alignment(&mut target, source_abi_align(ctx, pointee)); + Ok(vec![target]) + } +} + +/// Compute a pointer offset in source pointee elements. +struct PointerOffsetTranslation; + +impl OperationTranslation for PointerOffsetTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + input.attributes.remove("mir_ptr_offset_inbounds"); + ensure_empty_attributes("mir.ptr_offset", &input)?; + let ptr_ty = source.deref(ctx).get_operand(0).get_type(ctx); + let pointee = pointer_pointee(ctx, ptr_ty)?; + let is_shared = ptr_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|ptr| ptr.address_space == dialect_mir::types::address_space::SHARED); + let target_pointee = session + .translate_type(pointee) + .map_err(|error| error.to_string())?; + let inbounds = MirPtrOffsetOp::new(source).is_inbounds(ctx); + let mut operations = Vec::new(); + if is_shared && input.operands.len() == 2 && input.operands[1].ty == MlirType::Integer(64) { + let trunc_id = session.fresh_value(); + let mut trunc = MlirOperation::new("arith.trunci")?; + trunc.operands = vec![input.operands[1].clone()]; + trunc.results = vec![MlirResult { + id: trunc_id, + ty: MlirType::Integer(32), + }]; + trunc.location = input.location.clone(); + operations.push(trunc); + input.operands[1] = MlirValueUse { + id: trunc_id, + ty: MlirType::Integer(32), + }; + } + let mut target = renamed("llvm.getelementptr", input)?; + gep_properties( + &mut target, + target_pointee, + vec![i32::MIN], + if inbounds { 3 } else { 0 }, + ); + operations.push(target); + Ok(operations) + } +} + +/// Extract a logical field using its physical LLVM aggregate slot. +struct ExtractFieldTranslation; + +impl OperationTranslation for ExtractFieldTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.extract_field", &input)?; + let index = MirExtractFieldOp::new(source) + .get_attr_index(ctx) + .ok_or_else(|| "mir.extract_field has no index".to_owned())? + .0 as usize; + let operand_type = source.deref(ctx).get_operand(0).get_type(ctx); + let operand_ref = operand_type.deref(ctx); + + if operand_ref.is::() { + return super::storage::extract_union_field(ctx, source, input, session, index); + } + + if operand_ref.downcast_ref::().is_some() { + if index != 0 { + return Err("scalar newtype extraction only accepts field zero".into()); + } + return Ok(vec![renamed("builtin.unrealized_conversion_cast", input)?]); + } + + let slot = if operand_ref.downcast_ref::().is_some() { + (index < 2).then_some(Some(index)) + } else if let Some(slice) = operand_ref.downcast_ref::() { + (index < slice.field_count()).then_some(Some(index)) + } else if let Some(array) = operand_ref.downcast_ref::() { + (index < array.size() as usize).then_some(Some(index)) + } else if operand_ref.downcast_ref::().is_some() + || operand_ref.downcast_ref::().is_some() + { + aggregate_plan(ctx, operand_type, |field| { + session + .translate_type(field) + .map_err(|error| error.to_string()) + })? + .decl_to_slot + .get(index) + .copied() + } else { + return Err("mir.extract_field needs a slice, array, struct, tuple, or scalar".into()); + } + .ok_or_else(|| format!("mir.extract_field index {index} is out of bounds"))?; + + let Some(slot) = slot else { + let result = only_result("mir.extract_field", &input)?; + let mut poison = MlirOperation::new("llvm.mlir.poison")?; + poison.results.push(result); + poison.location = input.location; + return Ok(vec![poison]); + }; + + let mut target = renamed("llvm.extractvalue", input)?; + target.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![slot as i64]), + ); + Ok(vec![target]) + } +} + +/// Build a physical aggregate in Rust declaration order while retaining layout padding. +pub(crate) struct ConstructStructTranslation; + +/// Build a fixed array preserving every scalar register's source order. +struct ConstructArrayTranslation; + +impl OperationTranslation for ConstructArrayTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + MirConstructArrayOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ensure_empty_attributes("mir.construct_array", &input)?; + let result = only_result("mir.construct_array", &input)?; + let mut poison = MlirOperation::new("llvm.mlir.poison")?; + let mut current = MlirValueUse { + id: if input.operands.is_empty() { + result.id + } else { + session.fresh_value() + }, + ty: result.ty.clone(), + }; + poison.results.push(MlirResult { + id: current.id, + ty: current.ty.clone(), + }); + poison.location = input.location.clone(); + let mut operations = vec![poison]; + for (index, operand) in input.operands.iter().enumerate() { + let id = if index + 1 == input.operands.len() { + result.id + } else { + session.fresh_value() + }; + let mut insert = MlirOperation::new("llvm.insertvalue")?; + insert.operands = vec![current, operand.clone()]; + insert.results.push(MlirResult { + id, + ty: result.ty.clone(), + }); + insert.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + insert.location = input.location.clone(); + operations.push(insert); + current = MlirValueUse { + id, + ty: result.ty.clone(), + }; + } + Ok(operations) + } +} + +impl OperationTranslation for ConstructStructTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.construct_struct", &input)?; + let source_type = source.deref(ctx).get_result(0).get_type(ctx); + let plan = aggregate_plan(ctx, source_type, |field| { + session + .translate_type(field) + .map_err(|error| error.to_string()) + })?; + if input.operands.len() != plan.decl_to_slot.len() { + return Err(format!( + "mir.construct_struct has {} operands for {} fields", + input.operands.len(), + plan.decl_to_slot.len() + )); + } + let source_result = only_result("mir.construct_struct", &input)?; + let live_fields = plan + .memory_order + .iter() + .copied() + .filter(|&field| plan.decl_to_slot[field].is_some()) + .collect::>(); + + let mut operations = Vec::with_capacity(live_fields.len() + 1); + let poison_id = if live_fields.is_empty() { + source_result.id + } else { + session.fresh_value() + }; + let mut poison = MlirOperation::new("llvm.mlir.poison")?; + poison.results.push(MlirResult { + id: poison_id, + ty: source_result.ty.clone(), + }); + poison.location = input.location.clone(); + operations.push(poison); + + let mut current = MlirValueUse { + id: poison_id, + ty: source_result.ty.clone(), + }; + for (position, field) in live_fields.iter().copied().enumerate() { + let result_id = if position + 1 == live_fields.len() { + source_result.id + } else { + session.fresh_value() + }; + let mut insert = MlirOperation::new("llvm.insertvalue")?; + insert.results.push(MlirResult { + id: result_id, + ty: source_result.ty.clone(), + }); + insert.operands.push(current); + insert.operands.push(input.operands[field].clone()); + insert.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![plan.decl_to_slot[field].unwrap() as i64]), + ); + insert.location = input.location.clone(); + operations.push(insert); + current = MlirValueUse { + id: result_id, + ty: source_result.ty.clone(), + }; + } + Ok(operations) + } +} + +/// Construct a slice's pointer and length without changing the pointer's semantic kind. +struct ConstructSliceTranslation; + +impl OperationTranslation for ConstructSliceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + MirConstructSliceOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + build_slice_carrier(input, session) + } +} + +/// Attach an array's static extent without changing its data pointer or ownership kind. +pub(super) fn unsize_array_pointer( + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, +) -> Result, String> { + MirCastOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + ensure_empty_attributes("array-to-slice coercion", &input)?; + if input.operands.len() != 1 { + return Err("array-to-slice coercion requires one pointer".into()); + } + let from = source.deref(ctx).get_operand(0).get_type(ctx); + let from_ref = from.deref(ctx); + let pointer = from_ref + .downcast_ref::() + .ok_or("expected array pointer")?; + let pointee = pointer.pointee.deref(ctx); + let array = pointee + .downcast_ref::() + .ok_or("only direct array-to-slice coercion is supported")?; + let to = source.deref(ctx).get_result(0).get_type(ctx); + let to_ref = to.deref(ctx); + let slice = to_ref + .downcast_ref::() + .ok_or("expected slice result")?; + if pointer.address_space() != 0 + || array.element_type() != slice.element_ty + || pointer.kind != slice.kind + || pointer.is_mutable != slice.is_mutable + { + return Err( + "array-to-slice coercion must preserve element type and pointer kind in generic memory" + .into(), + ); + } + let id = session.fresh_value(); + let mut length = MlirOperation::new("arith.constant")?; + length.results.push(MlirResult { + id, + ty: MlirType::Integer(64), + }); + length.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: i128::from(array.size()), + ty: MlirType::Integer(64), + }, + ); + length.location = input.location.clone(); + input.operands.push(MlirValueUse { + id, + ty: MlirType::Integer(64), + }); + let mut operations = vec![length]; + operations.extend(build_slice_carrier(input, session)?); + Ok(operations) +} + +/// Assemble a verified generic pointer and 64-bit extent into the canonical slice carrier. +fn build_slice_carrier( + input: OperationInput, + session: &mut TranslationSession<'_>, +) -> Result, String> { + ensure_empty_attributes("mir.construct_slice", &input)?; + if input.operands.len() != 2 || !input.regions.is_empty() || !input.successors.is_empty() { + return Err("mir.construct_slice requires a pointer and length".into()); + } + let result = only_result("mir.construct_slice", &input)?; + if input.operands[0].ty != llvm_pointer_type(0)? + || input.operands[1].ty != MlirType::Integer(64) + { + return Err("mir.construct_slice requires a generic pointer and 64-bit length".into()); + } + let mut current = MlirValueUse { + id: session.fresh_value(), + ty: result.ty.clone(), + }; + let mut poison = MlirOperation::new("llvm.mlir.poison")?; + poison.results.push(MlirResult { + id: current.id, + ty: current.ty.clone(), + }); + poison.location = input.location.clone(); + let mut operations = vec![poison]; + for (index, operand) in input.operands.into_iter().enumerate() { + let id = if index == 1 { + result.id + } else { + session.fresh_value() + }; + let mut insert = MlirOperation::new("llvm.insertvalue")?; + insert.operands = vec![current, operand]; + insert.results.push(MlirResult { + id, + ty: result.ty.clone(), + }); + insert.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + insert.location = input.location.clone(); + operations.push(insert); + current = MlirValueUse { + id, + ty: result.ty.clone(), + }; + } + Ok(operations) +} + +/// Compute a logical field address through the verified physical aggregate layout. +struct FieldAddressTranslation; + +impl OperationTranslation for FieldAddressTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + // `aggregate_ty` is the build-time-stamped aggregate fact; this + // translation consumes it as the GEP element type source, so it must + // not survive as a target attribute. + let mut input = input; + input.attributes.remove("aggregate_ty"); + ensure_empty_attributes("mir.field_addr", &input)?; + let index = MirFieldAddrOp::new(source) + .get_attr_field_index(ctx) + .ok_or_else(|| "mir.field_addr has no field_index".to_owned())? + .0 as usize; + let pointee = match MirFieldAddrOp::new(source).get_attr_aggregate_ty(ctx) { + Some(attr) => attr.get_type(ctx), + None => pointer_pointee(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?, + }; + if let Some(union) = pointee.deref(ctx).downcast_ref::() { + MirFieldAddrOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + super::storage::union_storage_type(ctx, union)?; + if index >= union.field_count() { + return Err("mir.field_addr union field index is out of bounds".into()); + } + // All union views begin at byte zero. Keep a distinct result for + // the selected field while preserving the pointer address space. + let mut target = renamed("llvm.getelementptr", input)?; + gep_properties(&mut target, MlirType::Integer(8), vec![0], 0); + return Ok(vec![target]); + } + let plan = aggregate_plan(ctx, pointee, |field| { + session + .translate_type(field) + .map_err(|error| error.to_string()) + })?; + let slot = plan + .decl_to_slot + .get(index) + .ok_or_else(|| format!("mir.field_addr index {index} is out of bounds"))?; + let (element_type, indices) = match slot { + Some(slot) => (plan.ty, vec![0, *slot as i32]), + None => (MlirType::Integer(8), vec![0]), + }; + let mut target = renamed("llvm.getelementptr", input)?; + gep_properties(&mut target, element_type, indices, 0); + Ok(vec![target]) + } +} + +/// Compute the address of a dynamic element within a fixed array. +struct ArrayElementAddressTranslation; + +impl OperationTranslation for ArrayElementAddressTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.array_element_addr", &input)?; + let pointee = pointer_pointee(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + if pointee.deref(ctx).downcast_ref::().is_none() { + return Err("mir.array_element_addr pointer does not point to an array".into()); + } + let target_array = session + .translate_type(pointee) + .map_err(|error| error.to_string())?; + let mut target = renamed("llvm.getelementptr", input)?; + gep_properties(&mut target, target_array, vec![0, i32::MIN], 0); + Ok(vec![target]) + } +} + +/// Extract an array element while preserving constant and dynamic index behavior. +struct ExtractArrayElementTranslation; + +impl OperationTranslation for ExtractArrayElementTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("mir.extract_array_element", &input)?; + if input.operands.len() != 2 { + return Err("mir.extract_array_element expected array and index operands".into()); + } + let array_source_type = source.deref(ctx).get_operand(0).get_type(ctx); + let array_ref = array_source_type.deref(ctx); + let array = array_ref + .downcast_ref::() + .ok_or_else(|| "mir.extract_array_element operand is not an array".to_owned())?; + let target_array = session + .translate_type(array_source_type) + .map_err(|error| error.to_string())?; + let result = only_result("mir.extract_array_element", &input)?; + // Constant projections are register extraction. The memory fallback + // below remains for dynamic indexing, never for a fixed TMEM lane. + if let Some(constant) = source + .deref(ctx) + .get_operand(1) + .defining_op() + .and_then(|op| Operation::get_op::(op, ctx)) + .and_then(|op| op.get_attr_value(ctx)) + { + let bits = constant.value(); + if bits.bw() <= 64 && bits.to_u64() < array.size() { + let mut extract = MlirOperation::new("llvm.extractvalue")?; + extract.operands.push(input.operands[0].clone()); + extract.results.push(result); + extract.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![bits.to_u64() as i64]), + ); + extract.location = input.location; + return Ok(vec![extract]); + } + } + let pointer = llvm_pointer_type(0)?; + let count_id = session.fresh_value(); + let slot_id = session.fresh_value(); + let element_ptr_id = session.fresh_value(); + + let mut count = MlirOperation::new("arith.constant")?; + count.results.push(MlirResult { + id: count_id, + ty: MlirType::Integer(32), + }); + count.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(32), + }, + ); + count.location = input.location.clone(); + + let mut alloca = MlirOperation::new("llvm.alloca")?; + alloca.results.push(MlirResult { + id: slot_id, + ty: pointer.clone(), + }); + alloca.operands.push(MlirValueUse { + id: count_id, + ty: MlirType::Integer(32), + }); + alloca.properties.insert( + "elem_type".into(), + MlirAttribute::Type(target_array.clone()), + ); + add_alignment(&mut alloca, source_abi_align(ctx, array_source_type)); + alloca.location = input.location.clone(); + + let mut store = MlirOperation::new("llvm.store")?; + store.operands.push(input.operands[0].clone()); + store.operands.push(MlirValueUse { + id: slot_id, + ty: pointer.clone(), + }); + memory_ordering(&mut store); + add_alignment(&mut store, source_abi_align(ctx, array_source_type)); + store.location = input.location.clone(); + + let mut gep = MlirOperation::new("llvm.getelementptr")?; + gep.results.push(MlirResult { + id: element_ptr_id, + ty: pointer.clone(), + }); + gep.operands.push(MlirValueUse { + id: slot_id, + ty: pointer.clone(), + }); + gep.operands.push(input.operands[1].clone()); + gep_properties(&mut gep, target_array, vec![0, i32::MIN], 0); + gep.location = input.location.clone(); + + let mut load = MlirOperation::new("llvm.load")?; + load.results.push(result); + load.operands.push(MlirValueUse { + id: element_ptr_id, + ty: pointer, + }); + memory_ordering(&mut load); + add_alignment(&mut load, source_abi_align(ctx, array.element_type())); + load.location = input.location; + Ok(vec![count, alloca, store, gep, load]) + } +} + +/// Map reviewed Rust intrinsics directly and preserve ordinary function call signatures. +struct CallTranslation; + +impl OperationTranslation for CallTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + input.attributes.remove("callee"); + ensure_empty_attributes("mir.call", &input)?; + let callee = MirCallOp::new(source) + .get_attr_callee(ctx) + .ok_or_else(|| "mir.call has no callee".to_owned())? + .as_str() + .to_owned(); + if callee.is_empty() || callee.chars().any(char::is_whitespace) { + return Err(format!("mir.call has invalid callee {callee:?}")); + } + // Intrinsic placeholders are compiler operations, never external symbols. + // Preserve source type checking before dropping the internal callee name. + use dialect_mir::rust_intrinsics::{ + CALLEE_MAXNUM_NSZ_F32, CALLEE_SATURATING_ADD, CALLEE_SATURATING_SUB, + CALLEE_SELECT_UNPREDICTABLE, PLACEHOLDER_PREFIX, + }; + if callee.starts_with(PLACEHOLDER_PREFIX) { + MirCallOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + if callee == CALLEE_MAXNUM_NSZ_F32 { + let scalar = MlirType::Float(MlirFloatType::F32); + if input.operands.len() != 2 + || input.results.len() != 1 + || input.operands.iter().any(|value| value.ty != scalar) + || input.results[0].ty != scalar + || !input.attributes.is_empty() + { + return Err("maximum_number_nsz_f32 requires exactly two FP32 operands and one FP32 result".into()); + } + // Native PTX max.f32 ignores either quiet or signaling NaNs + // when the other operand is numeric, preserving Rust's max + // contract. Do not use llvm.maxnum (different sNaN semantics) + // or enable NVVM's NaN-propagation/FTZ/absolute-value modifiers. + return Ok(vec![renamed("nvvm.fmax", input)?]); + } + if callee == CALLEE_SELECT_UNPREDICTABLE { + let supported = match &input.results[0].ty { + MlirType::Integer(_) | MlirType::Float(_) | MlirType::Vector { .. } => true, + MlirType::Dialect(t) => t == "!llvm.ptr" || t.starts_with("!llvm.ptr<"), + _ => false, + }; + if !supported { + return Err("select_unpredictable currently requires a scalar, pointer, or vector result".into()); + } + // The branch-weight hint does not change device semantics. + return Ok(vec![renamed("arith.select", input)?]); + } + if matches!( + callee.as_str(), + CALLEE_SATURATING_ADD | CALLEE_SATURATING_SUB + ) { + let ty = source.deref(ctx).get_operand(0).get_type(ctx); + let ty = ty.deref(ctx); + let integer = ty + .downcast_ref::() + .ok_or("saturation requires integers")?; + let sign = if integer.is_signed() { "s" } else { "u" }; + let operation = if callee == CALLEE_SATURATING_ADD { + "add" + } else { + "sub" + }; + let intrinsic = format!("llvm.{sign}{operation}.sat.i{}", integer.width()); + let mut target = renamed("llvm.call_intrinsic", input)?; + target + .properties + .insert("intrin".into(), MlirAttribute::String(intrinsic)); + target.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + target.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![2, 0]), + ); + return Ok(vec![target]); + } + return Err(format!( + "Rust intrinsic {callee} has no CUTLASS MLIR translation" + )); + } + let unit_result = if MirCallOp::new(source).has_logical_unit_result_bridge(ctx) { + input.results.pop() + } else { + None + }; + let mut target = renamed("func.call", input)?; + target + .properties + .insert("callee".into(), MlirAttribute::SymbolRef(callee)); + if let Some(result) = unit_result { + // MIR lets the logical unit flow through a destination place. + // MLIR's call must instead match the callee's void signature. + let mut carrier = MlirOperation::new("ub.poison")?; + carrier.results.push(result); + carrier.location = target.location.clone(); + Ok(vec![target, carrier]) + } else { + Ok(vec![target]) + } + } +} + +/// Preserve Rust's mandatory trap independently of optional backend assertions. +struct AssertTranslation; + +impl OperationTranslation for AssertTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + Operation::get_op::(source, ctx) + .ok_or("expected mir.assert")? + .verify(ctx) + .map_err(|error| error.to_string())?; + ensure_empty_attributes("mir.assert", &input)?; + if input.operands.len() != 1 + || input.operands[0].ty != MlirType::Integer(1) + || !input.successors.is_empty() + || !input.results.is_empty() + || !input.regions.is_empty() + { + return Err("mir.assert expected one i1 condition and no successors or results".into()); + } + // CUTLASS's cute-to-nvvm pipeline removes cf.assert by default when + // enable-assertions=false. Rust bounds/overflow checks are mandatory + // source effects. Expose the failure branch and its nonreturning trap + // explicitly so later optimizations retain the same control flow. + let true_value = MlirValueUse { + id: session.fresh_value(), + ty: MlirType::Integer(1), + }; + let mut constant = MlirOperation::new("arith.constant")?; + constant.results.push(MlirResult { + id: true_value.id, + ty: true_value.ty.clone(), + }); + constant.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(1), + }, + ); + constant.location = input.location.clone(); + + let failure = MlirValueUse { + id: session.fresh_value(), + ty: MlirType::Integer(1), + }; + let mut invert = MlirOperation::new("arith.xori")?; + invert.operands = vec![input.operands[0].clone(), true_value]; + invert.results.push(MlirResult { + id: failure.id, + ty: failure.ty.clone(), + }); + invert.location = input.location.clone(); + + let mut trap = MlirOperation::new("llvm.intr.trap")?; + trap.location = input.location.clone(); + let mut yield_op = MlirOperation::new("scf.yield")?; + yield_op.location = input.location.clone(); + let mut conditional = MlirOperation::new("scf.if")?; + conditional.operands.push(failure); + conditional.regions = vec![ + crate::MlirRegion { + blocks: vec![crate::MlirBlock { + id: session.fresh_block(), + arguments: vec![], + operations: vec![trap, yield_op], + }], + }, + crate::MlirRegion { blocks: vec![] }, + ]; + conditional.location = input.location; + + // Any following MIR branch is translated independently. The structured + // failure region only adds the conditional trap before that continuation. + Ok(vec![constant, invert, conditional]) + } +} + +/// Erase a source-only operation after checking that it has no runtime results. +struct DropOperation; + +impl OperationTranslation for DropOperation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + ensure_empty_attributes("MIR lifetime marker", &input)?; + if !input.operands.is_empty() + || !input.results.is_empty() + || !input.successors.is_empty() + || !input.regions.is_empty() + { + return Err("MIR lifetime marker unexpectedly carries IR data".into()); + } + Ok(vec![]) + } +} + +/// Validated physical LLVM aggregate plus mappings from Rust declaration indices. +#[derive(Debug)] +pub(crate) struct AggregatePlan { + /// LLVM aggregate type including any explicit padding slots. + pub(crate) ty: MlirType, + /// Physical slot for each declared field; zero-sized fields have no slot. + pub(crate) decl_to_slot: Vec>, + /// Rust declaration indices ordered by their physical positions in memory. + pub(crate) memory_order: Vec, +} + +/// Rebuild Rust aggregate layout and reject overlaps, packed offsets, or incompatible total size. +pub(crate) fn aggregate_plan( + ctx: &Context, + source: TypeHandle, + mut translate_type: impl FnMut(TypeHandle) -> Result, +) -> Result { + let source_ref = source.deref(ctx); + let (fields, memory_order, field_offsets, total_size) = + if let Some(structure) = source_ref.downcast_ref::() { + ( + structure.field_types().to_vec(), + structure.memory_order(), + structure.field_offsets().to_vec(), + structure.total_size(), + ) + } else if let Some(tuple) = source_ref.downcast_ref::() { + ( + tuple.get_types().to_vec(), + tuple.memory_order(), + tuple.field_offsets().to_vec(), + tuple.total_size(), + ) + } else { + return Err("expected a MIR struct or tuple".into()); + }; + + validate_memory_order(&memory_order, fields.len())?; + let explicit = !field_offsets.is_empty() && total_size > 0; + if explicit && field_offsets.len() != fields.len() { + return Err(format!( + "aggregate records {} offsets for {} fields", + field_offsets.len(), + fields.len() + )); + } + + let translated = fields + .iter() + .map(|&field| translate_type(field)) + .collect::, _>>()?; + let mut target_fields = Vec::new(); + let mut decl_to_slot = vec![None; fields.len()]; + let mut offset = 0u64; + let mut max_alignment = 1u64; + + for &declaration in &memory_order { + let source_field = fields[declaration]; + if source_is_zero_sized(ctx, source_field) { + continue; + } + let field_alignment = source_natural_alignment(ctx, source_field).ok_or_else(|| { + format!("cannot determine natural alignment of aggregate field {declaration}") + })?; + max_alignment = max_alignment.max(field_alignment); + let target_offset = if explicit { + field_offsets[declaration] + } else { + align_up(offset, field_alignment)? + }; + if offset > target_offset { + return Err(format!( + "aggregate field {declaration} overlaps a previous field at byte {target_offset}" + )); + } + if offset < target_offset { + target_fields.push(llvm_array_type( + target_offset - offset, + &MlirType::Integer(8), + )?); + offset = target_offset; + } + let llvm_offset = align_up(offset, field_alignment)?; + if llvm_offset != target_offset { + return Err(format!( + "aggregate field {declaration} is at byte {target_offset} in Rust but an unpacked LLVM struct would place it at byte {llvm_offset}; packed by-value aggregates are not supported" + )); + } + decl_to_slot[declaration] = Some(target_fields.len()); + target_fields.push(translated[declaration].clone()); + offset = target_offset + .checked_add(source_stored_size(ctx, source_field).ok_or_else(|| { + format!("cannot determine stored size of aggregate field {declaration}") + })?) + .ok_or_else(|| "aggregate size overflow".to_owned())?; + } + + if explicit { + if offset > total_size { + return Err(format!( + "aggregate fields occupy {offset} bytes but Rust reports {total_size}" + )); + } + if offset < total_size { + target_fields.push(llvm_array_type(total_size - offset, &MlirType::Integer(8))?); + offset = total_size; + } + let natural_size = align_up(offset, max_alignment)?; + if natural_size != total_size { + return Err(format!( + "aggregate is {total_size} bytes in Rust but its unpacked LLVM form is {natural_size} bytes" + )); + } + } + + Ok(AggregatePlan { + ty: llvm_struct_type(&target_fields)?, + decl_to_slot, + memory_order, + }) +} + +/// Require a complete, nonduplicated permutation of all declared field indices. +fn validate_memory_order(order: &[usize], fields: usize) -> Result<(), String> { + if order.len() != fields { + return Err(format!( + "aggregate memory order has {} entries for {fields} fields", + order.len() + )); + } + let mut seen = vec![false; fields]; + for &field in order { + if field >= fields || seen[field] { + return Err(format!( + "aggregate memory order {order:?} is not a permutation of 0..{fields}" + )); + } + seen[field] = true; + } + Ok(()) +} + +/// Determine whether a reviewed source type occupies no stored bytes. +fn source_is_zero_sized(ctx: &Context, source: TypeHandle) -> bool { + let source_ref = source.deref(ctx); + if let Some(union) = source_ref.downcast_ref::() { + return union.total_size() == 0; + } + if let Some(e) = source_ref.downcast_ref::() { + return e.total_size == 0 && crate::mir::checked::enum_storage_type(e).is_ok(); + } + if source_ref.downcast_ref::().is_some() { + return true; + } + if let Some(array) = source_ref.downcast_ref::() { + return array.size() == 0 || source_is_zero_sized(ctx, array.element_type()); + } + if let Some(structure) = source_ref.downcast_ref::() { + return structure + .field_types() + .iter() + .all(|&field| source_is_zero_sized(ctx, field)); + } + if let Some(tuple) = source_ref.downcast_ref::() { + return tuple + .get_types() + .iter() + .all(|&field| source_is_zero_sized(ctx, field)); + } + false +} + +/// Return the reviewed source type's storage size in bytes, or None when unknown. +pub(crate) fn source_stored_size(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + if let Some(union) = source_ref.downcast_ref::() { + return Some(union.total_size()); + } + if let Some(e) = source_ref.downcast_ref::() { + return crate::mir::checked::enum_storage_type(e) + .ok() + .map(|_| e.total_size); + } + if let Some(integer) = source_ref.downcast_ref::() { + return Some(u64::from(integer.width()).div_ceil(8)); + } + if source_ref.downcast_ref::().is_some() + || source_ref.downcast_ref::().is_some() + { + return Some(2); + } + if source_ref.downcast_ref::().is_some() { + return Some(4); + } + if source_ref.downcast_ref::().is_some() { + return Some(8); + } + if source_ref.downcast_ref::().is_some() { + return Some(8); + } + if source_ref.downcast_ref::().is_some() { + return Some(16); + } + if let Some(slice) = source_ref.downcast_ref::() { + let mut size = 16u64; + let mut alignment = 8u64; + for &field in slice.space_types() { + let field_alignment = source_natural_alignment(ctx, field)?; + size = align_up(size, field_alignment).ok()?; + size = size.checked_add(source_stored_size(ctx, field)?)?; + alignment = alignment.max(field_alignment); + } + return align_up(size, alignment).ok(); + } + if let Some(array) = source_ref.downcast_ref::() { + return source_stored_size(ctx, array.element_type())?.checked_mul(array.size()); + } + if let Some(structure) = source_ref.downcast_ref::() { + if structure.total_size() > 0 { + return Some(structure.total_size()); + } + return natural_aggregate_size(ctx, structure.field_types(), &structure.memory_order()); + } + if let Some(tuple) = source_ref.downcast_ref::() { + if tuple.total_size() > 0 { + return Some(tuple.total_size()); + } + return natural_aggregate_size(ctx, tuple.get_types(), &tuple.memory_order()); + } + if source_ref.downcast_ref::().is_some() { + return Some(0); + } + None +} + +/// Return the unpacked LLVM-compatible storage alignment in bytes when known. +pub(crate) fn source_natural_alignment(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + if let Some(union) = source_ref.downcast_ref::() { + return Some(union.abi_align()); + } + if let Some(e) = source_ref.downcast_ref::() { + return crate::mir::checked::enum_storage_type(e) + .ok() + .map(|_| e.abi_align); + } + if let Some(integer) = source_ref.downcast_ref::() { + return Some( + u64::from(integer.width()) + .div_ceil(8) + .next_power_of_two() + .min(16), + ); + } + if source_ref.downcast_ref::().is_some() + || source_ref.downcast_ref::().is_some() + { + return Some(2); + } + if source_ref.downcast_ref::().is_some() { + return Some(4); + } + if source_ref.downcast_ref::().is_some() + || source_ref.downcast_ref::().is_some() + || source_ref.downcast_ref::().is_some() + || source_ref.downcast_ref::().is_some() + { + return Some(8); + } + if let Some(array) = source_ref.downcast_ref::() { + return source_natural_alignment(ctx, array.element_type()); + } + if let Some(structure) = source_ref.downcast_ref::() { + return structure + .field_types() + .iter() + .filter(|&&field| !source_is_zero_sized(ctx, field)) + .filter_map(|&field| source_natural_alignment(ctx, field)) + .max() + .or(Some(1)); + } + if let Some(tuple) = source_ref.downcast_ref::() { + return tuple + .get_types() + .iter() + .filter(|&&field| !source_is_zero_sized(ctx, field)) + .filter_map(|&field| source_natural_alignment(ctx, field)) + .max() + .or(Some(1)); + } + if source_ref.downcast_ref::().is_some() { + return Some(1); + } + None +} + +/// Read explicit Rust aggregate or enum alignment, falling back to natural alignment. +pub(super) fn source_abi_align(ctx: &Context, source: TypeHandle) -> Option { + let source_ref = source.deref(ctx); + if let Some(array) = source_ref.downcast_ref::() { + return source_abi_align(ctx, array.element_type()); + } + if let Some(structure) = source_ref.downcast_ref::() + && structure.abi_align > 0 + { + return Some(structure.abi_align); + } + if let Some(tuple) = source_ref.downcast_ref::() + && tuple.abi_align() > 0 + { + return Some(tuple.abi_align()); + } + source_natural_alignment(ctx, source) +} + +/// Compute padded byte size for fields in their physical memory order. +fn natural_aggregate_size(ctx: &Context, fields: &[TypeHandle], order: &[usize]) -> Option { + let mut offset = 0u64; + let mut alignment = 1u64; + for &field in order { + if source_is_zero_sized(ctx, fields[field]) { + continue; + } + let field_alignment = source_natural_alignment(ctx, fields[field])?; + offset = align_up(offset, field_alignment).ok()?; + offset = offset.checked_add(source_stored_size(ctx, fields[field])?)?; + alignment = alignment.max(field_alignment); + } + align_up(offset, alignment).ok() +} + +/// Round a byte offset up to a nonzero power-of-two alignment, rejecting overflow. +fn align_up(value: u64, alignment: u64) -> Result { + if alignment == 0 || !alignment.is_power_of_two() { + return Err(format!("invalid alignment {alignment}")); + } + value + .checked_add(alignment - 1) + .map(|value| value & !(alignment - 1)) + .ok_or_else(|| "alignment overflow".to_owned()) +} + +/// Require a MIR pointer and recover its source pointee type. +fn pointer_pointee(ctx: &Context, pointer: TypeHandle) -> Result { + let pointer_ref = pointer.deref(ctx); + pointer_ref + .downcast_ref::() + .map(|pointer| pointer.pointee) + .ok_or_else(|| "expected mir.ptr".to_owned()) +} + +/// Spell an opaque LLVM pointer in the requested numeric address space. +fn llvm_pointer_type(address_space: u32) -> Result { + if address_space == 0 { + MlirType::dialect("!llvm.ptr") + } else { + MlirType::dialect(format!("!llvm.ptr<{address_space}>")) + } +} + +/// Spell an LLVM fixed array with its element count and translated element type. +fn llvm_array_type(size: u64, element: &MlirType) -> Result { + MlirType::dialect(format!("!llvm.array<{size} x {}>", type_spelling(element))) +} + +/// Spell an unpacked LLVM struct in the supplied physical field order. +fn llvm_struct_type(fields: &[MlirType]) -> Result { + let fields = fields + .iter() + .map(type_spelling) + .collect::>() + .join(", "); + MlirType::dialect(format!("!llvm.struct<({fields})>")) +} + +/// Render a target type for embedding in an LLVM dialect type spelling. +pub(crate) fn type_spelling(ty: &MlirType) -> String { + match ty { + MlirType::Index => "index".into(), + MlirType::Integer(width) => format!("i{width}"), + MlirType::Float(MlirFloatType::F16) => "f16".into(), + MlirType::Float(MlirFloatType::BF16) => "bf16".into(), + MlirType::Float(MlirFloatType::F32) => "f32".into(), + MlirType::Float(MlirFloatType::F64) => "f64".into(), + MlirType::Float(MlirFloatType::Other(spelling)) => spelling.clone(), + MlirType::Vector { shape, element } => format!( + "vector<{}{}>", + shape + .iter() + .map(|value| format!("{value}x")) + .collect::(), + type_spelling(element) + ), + MlirType::Function { inputs, results } => format!( + "({}) -> ({})", + inputs + .iter() + .map(type_spelling) + .collect::>() + .join(", "), + results + .iter() + .map(type_spelling) + .collect::>() + .join(", ") + ), + MlirType::Tuple(elements) => format!( + "tuple<{}>", + elements + .iter() + .map(type_spelling) + .collect::>() + .join(", ") + ), + MlirType::Dialect(spelling) => spelling.clone(), + } +} + +/// Require exactly one translated result and return its SSA definition. +fn only_result(name: &str, input: &OperationInput) -> Result { + match input.results.as_slice() { + [result] => Ok(result.clone()), + results => Err(format!("{name} expected one result, got {}", results.len())), + } +} + +/// Reject source attributes not consumed by the recipe. +fn ensure_empty_attributes(name: &str, input: &OperationInput) -> Result<(), String> { + if input.attributes.is_empty() { + Ok(()) + } else { + Err(format!( + "{name} has targetless attributes: {:?}", + input.attributes.keys().collect::>() + )) + } +} + +/// Retain translated operation structure while changing its target operation name. +fn renamed(name: &str, input: OperationInput) -> Result { + let mut target = MlirOperation::new(name)?; + target.results = input.results; + target.operands = input.operands; + target.successors = input.successors; + target.regions = input.regions; + target.attributes = input.attributes; + target.location = input.location; + Ok(target) +} + +/// Mark an ordinary load/store as non-atomic in the LLVM dialect contract. +fn memory_ordering(operation: &mut MlirOperation) { + operation.properties.insert( + "ordering".into(), + MlirAttribute::Integer { + value: 0, + ty: MlirType::Integer(64), + }, + ); +} + +/// Attach a byte alignment property when the source records one. +fn add_alignment(operation: &mut MlirOperation, alignment: Option) { + if let Some(alignment) = alignment { + operation.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: alignment.into(), + ty: MlirType::Integer(64), + }, + ); + } +} + +/// Set GEP element and index metadata without promising no-wrap behavior. +fn gep_properties( + operation: &mut MlirOperation, + element_type: MlirType, + indices: Vec, + no_wrap_flags: i128, +) { + operation + .properties + .insert("elem_type".into(), MlirAttribute::Type(element_type)); + operation.properties.insert( + "noWrapFlags".into(), + MlirAttribute::Integer { + value: no_wrap_flags, + ty: MlirType::Integer(32), + }, + ); + operation.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(indices), + ); +} + +#[cfg(test)] +mod tests { + //! Physical aggregate, dynamic array, and effectful assertion translation checks. + + use crate::{ + CutlassFullCuteMlir22, MlirConsumerProfile, + profile::render_mapping_module_without_cutlass_envelope, + }; + use dialect_mir::{ + attributes::{CompilerResultBundleAttr, FieldIndexAttr}, + ops::{ + MirAddOp, MirAllocaOp, MirArrayElementAddrOp, MirAssertOp, MirConstructArrayOp, + MirExtractArrayElementOp, MirExtractFieldOp, MirFieldAddrOp, MirFuncOp, MirLoadOp, + MirReturnOp, MirStoreOp, + }, + types::{MirArrayType, MirPtrType, MirStructType}, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness, UnitType}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::{TypeHandle, TypedHandle}, + value::Value, + }; + + /// Construct a typed source operation and append it to the test block. + fn operation( + ctx: &mut Context, + block: Ptr, + results: Vec, + operands: Vec, + ) -> Ptr { + let operation = + Operation::new(ctx, O::get_concrete_op_info(), results, operands, vec![], 0); + operation.insert_at_back(block, ctx); + operation + } + + /// Obtain an unsigned source integer type with the requested bit width. + fn integer(ctx: &mut Context, width: u32) -> TypedHandle { + IntegerType::get(ctx, width, Signedness::Unsigned) + } + + /// Rust assertions must trap even when the downstream compiler disables optional assertions. + #[test] + fn assertion_remains_an_effectful_check_before_following_operations() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("assertion").unwrap()); + let condition = integer(&mut ctx, 1); + let word = integer(&mut ctx, 32); + let pointer = MirPtrType::get(&mut ctx, word.into(), true, 0); + let inputs = vec![condition.into(), word.into(), pointer.into()]; + let ty = FunctionType::get(&ctx, inputs.clone(), vec![word.into()]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(ty.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("checked_add").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, inputs); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let condition = block.deref(&ctx).get_argument(0); + let word = block.deref(&ctx).get_argument(1); + let pointer = block.deref(&ctx).get_argument(2); + operation::(&mut ctx, block, vec![], vec![condition]); + let result_type = integer(&mut ctx, 32); + let sum = + operation::(&mut ctx, block, vec![result_type.into()], vec![word, word]); + let sum = sum.deref(&ctx).get_result(0); + operation::(&mut ctx, block, vec![], vec![pointer, sum]); + operation::(&mut ctx, block, vec![], vec![sum]); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_mapping_module_without_cutlass_envelope(&target, "assertion"); + assert_eq!(text.matches("\"llvm.intr.trap\"").count(), 1, "{text}"); + assert_eq!(text.matches("\"scf.if\"").count(), 1, "{text}"); + assert!( + text.find("\"scf.if\"").unwrap() < text.find("\"llvm.store\"").unwrap(), + "{text}" + ); + let gpu_module = &target.root.regions[0].blocks[0].operations[0]; + let function = &gpu_module.regions[0].blocks[0].operations[0]; + let entry = &function.regions[0].blocks[0]; + let inversion = &entry.operations[1]; + assert_eq!(inversion.name, "arith.xori"); + assert_eq!(inversion.operands[0].id, entry.arguments[0].id); + let true_constant = &entry.operations[0]; + assert_eq!(true_constant.name, "arith.constant"); + assert_eq!( + true_constant.properties.get("value"), + Some(&crate::MlirAttribute::Integer { + value: 1, + ty: crate::MlirType::Integer(1), + }), + ); + assert_eq!(inversion.operands[1].id, true_constant.results[0].id); + let conditional = &entry.operations[2]; + assert_eq!(conditional.operands[0].id, inversion.results[0].id); + let failure = &conditional.regions[0].blocks[0].operations; + assert_eq!(failure[0].name, "llvm.intr.trap"); + assert_eq!(failure[1].name, "scf.yield"); + assert!(conditional.regions[1].blocks.is_empty()); + assert!(!text.contains("\"cf.assert\""), "{text}"); + assert!(!text.contains("\"cf.br\""), "{text}"); + } + + /// Conditions wider than a boolean cannot be reinterpreted as an assertion predicate. + #[test] + fn assertion_rejects_non_boolean_conditions() { + for width in [8, 32] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let module = + ModuleOp::new(&mut ctx, Identifier::try_from("invalid_assertion").unwrap()); + let condition = integer(&mut ctx, width); + let ty = FunctionType::get(&ctx, vec![condition.into()], vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(ty.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("invalid_condition").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, vec![condition.into()]); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let condition = block.deref(&ctx).get_argument(0); + operation::(&mut ctx, block, vec![], vec![condition]); + operation::(&mut ctx, block, vec![], vec![]); + let error = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap_err(); + assert!(error.to_string().contains("i1"), "{error}"); + } + } + + /// Check compiler-created arrays retain scalar lane order and valid zero-length representations. + #[test] + fn compiler_result_array_preserves_lane_order_and_empty_arrays() { + for count in [0, 4, 32] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("registers").unwrap()); + let word = integer(&mut ctx, 32); + let array = MirArrayType::get(&mut ctx, word.into(), count); + let args = vec![word.into(); count as usize]; + let ty = FunctionType::get(&ctx, args.clone(), vec![array.into()]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(ty.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("bundle").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let values = (0..count as usize) + .map(|i| block.deref(&ctx).get_argument(i)) + .collect(); + let bundle = + operation::(&mut ctx, block, vec![array.into()], values); + bundle.deref_mut(&ctx).attributes.set( + Identifier::try_from("compiler_result_bundle").unwrap(), + CompilerResultBundleAttr(true), + ); + let value = bundle.deref(&ctx).get_result(0); + operation::(&mut ctx, block, vec![], vec![value]); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_mapping_module_without_cutlass_envelope(&target, "registers"); + assert_eq!( + text.matches("\"llvm.insertvalue\"").count(), + count as usize, + "{text}" + ); + for i in 0..count { + let position = format!("position = array"); + let insert = text.lines().find(|line| line.contains(&position)).unwrap(); + assert!(insert.contains(&format!(", %v{i})")), "{text}"); + } + assert!(!text.contains("compiler_result_bundle"), "{text}"); + assert!(!text.contains("llvm.alloca"), "{text}"); + } + } + + /// Check logical fields map to physical aggregate slots and dynamic array accesses retain element semantics. + #[test] + fn aggregate_slots_and_dynamic_array_access_map_to_llvm() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, Identifier::try_from("memory").unwrap()); + let f32_type = FP32Type::get(&ctx); + let u64_type = integer(&mut ctx, 64); + let unit_type = UnitType::get(&ctx); + let array_type = MirArrayType::get(&mut ctx, f32_type.into(), 4); + let structure_type = MirStructType::get_with_full_layout( + &mut ctx, + "RegisterTile".into(), + vec!["marker".into(), "values".into()], + vec![unit_type.into(), array_type.into()], + vec![0, 1], + vec![0, 0], + 16, + 4, + ); + let structure_pointer = MirPtrType::get_generic(&mut ctx, structure_type.into(), true); + let array_pointer = MirPtrType::get_generic(&mut ctx, array_type.into(), true); + let element_pointer = MirPtrType::get_generic(&mut ctx, f32_type.into(), true); + + let function_type = FunctionType::get( + &ctx, + vec![structure_type.into(), u64_type.into()], + vec![f32_type.into()], + ); + let function_operation = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new( + &mut ctx, + function_operation, + TypeAttr::new(function_type.into()), + ); + function.set_symbol_name(&mut ctx, Identifier::try_from("read_tile").unwrap()); + function_operation.deref_mut(&ctx).attributes.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, function_operation, 0); + let region = function_operation.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![structure_type.into(), u64_type.into()]); + entry.insert_at_back(region, &ctx); + let arguments = entry.deref(&ctx).arguments().collect::>(); + + let alloca = + operation::(&mut ctx, entry, vec![structure_pointer.into()], vec![]); + let slot = alloca.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![slot, arguments[0]]); + + let field = + operation::(&mut ctx, entry, vec![array_pointer.into()], vec![slot]); + MirFieldAddrOp::new(field).set_attr_field_index(&ctx, FieldIndexAttr(1)); + MirFieldAddrOp::new(field).set_attr_aggregate_ty( + &ctx, + pliron::builtin::attributes::TypeAttr::new(structure_type.into()), + ); + let field_result = field.deref(&ctx).get_result(0); + let element = operation::( + &mut ctx, + entry, + vec![element_pointer.into()], + vec![field_result, arguments[1]], + ); + let element_result = element.deref(&ctx).get_result(0); + let loaded = + operation::(&mut ctx, entry, vec![f32_type.into()], vec![element_result]); + let loaded_result = loaded.deref(&ctx).get_result(0); + + let array = operation::( + &mut ctx, + entry, + vec![array_type.into()], + vec![arguments[0]], + ); + MirExtractFieldOp::new(array).set_attr_index(&ctx, FieldIndexAttr(1)); + let array_result = array.deref(&ctx).get_result(0); + let extracted = operation::( + &mut ctx, + entry, + vec![f32_type.into()], + vec![array_result, arguments[1]], + ); + let extracted_result = extracted.deref(&ctx).get_result(0); + let sum = operation::( + &mut ctx, + entry, + vec![f32_type.into()], + vec![loaded_result, extracted_result], + ); + let sum_result = sum.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![sum_result]); + + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_mapping_module_without_cutlass_envelope(&target, "memory"); + let expected = r#""builtin.module"() <{sym_name = "memory"}> ({ + ^bb0: + "func.func"() <{function_type = (!llvm.struct<(!llvm.array<4 x f32>)>, i64) -> (f32), sym_name = "read_tile"}> ({ + ^bb1(%v0: !llvm.struct<(!llvm.array<4 x f32>)>, %v1: i64): + %v9 = "arith.constant"() <{value = 1 : i32}> : () -> (i32) + %v2 = "llvm.alloca"(%v9) <{alignment = 4 : i64, elem_type = !llvm.struct<(!llvm.array<4 x f32>)>}> : (i32) -> (!llvm.ptr) + "llvm.store"(%v0, %v2) <{alignment = 4 : i64, ordering = 0 : i64}> : (!llvm.struct<(!llvm.array<4 x f32>)>, !llvm.ptr) -> () + %v3 = "llvm.getelementptr"(%v2) <{elem_type = !llvm.struct<(!llvm.array<4 x f32>)>, noWrapFlags = 0 : i32, rawConstantIndices = array}> : (!llvm.ptr) -> (!llvm.ptr) + %v4 = "llvm.getelementptr"(%v3, %v1) <{elem_type = !llvm.array<4 x f32>, noWrapFlags = 0 : i32, rawConstantIndices = array}> : (!llvm.ptr, i64) -> (!llvm.ptr) + %v5 = "llvm.load"(%v4) <{alignment = 4 : i64, ordering = 0 : i64}> : (!llvm.ptr) -> (f32) + %v6 = "llvm.extractvalue"(%v0) <{position = array}> : (!llvm.struct<(!llvm.array<4 x f32>)>) -> (!llvm.array<4 x f32>) + %v10 = "arith.constant"() <{value = 1 : i32}> : () -> (i32) + %v11 = "llvm.alloca"(%v10) <{alignment = 4 : i64, elem_type = !llvm.array<4 x f32>}> : (i32) -> (!llvm.ptr) + "llvm.store"(%v6, %v11) <{alignment = 4 : i64, ordering = 0 : i64}> : (!llvm.array<4 x f32>, !llvm.ptr) -> () + %v12 = "llvm.getelementptr"(%v11, %v1) <{elem_type = !llvm.array<4 x f32>, noWrapFlags = 0 : i32, rawConstantIndices = array}> : (!llvm.ptr, i64) -> (!llvm.ptr) + %v7 = "llvm.load"(%v12) <{alignment = 4 : i64, ordering = 0 : i64}> : (!llvm.ptr) -> (f32) + %v8 = "arith.addf"(%v5, %v7) : (f32, f32) -> (f32) + "func.return"(%v8) : (f32) -> () + }) {cute.kernel = unit, gpu.kernel = unit} : () -> () +}) : () -> () +"#; + assert_eq!( + crate::profile::normalize_mlir_identifiers(&text), + crate::profile::normalize_mlir_identifiers(expected) + ); + } + + #[test] + fn shared_pointer_offset_truncates_i64_index_to_i32() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let u32_ty = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let shared_ptr: TypeHandle = MirPtrType::get( + &mut ctx, + u32_ty.into(), + true, + dialect_mir::types::address_space::SHARED, + ) + .into(); + + let module = ModuleOp::new(&mut ctx, Identifier::try_from("memory").unwrap()); + let function_type = + FunctionType::get(&ctx, vec![shared_ptr, u64_ty.into()], vec![u32_ty.into()]); + let function_operation = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new( + &mut ctx, + function_operation, + TypeAttr::new(function_type.into()), + ); + function.set_symbol_name(&mut ctx, Identifier::try_from("shared_offset").unwrap()); + function_operation.deref_mut(&ctx).attributes.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, function_operation, 0); + let region = function_operation.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![shared_ptr, u64_ty.into()]); + entry.insert_at_back(region, &ctx); + let ptr_arg = entry.deref(&ctx).get_argument(0); + let idx_arg = entry.deref(&ctx).get_argument(1); + let offset = operation::( + &mut ctx, + entry, + vec![shared_ptr], + vec![ptr_arg, idx_arg], + ); + let offset_res = offset.deref(&ctx).get_result(0); + let loaded = operation::(&mut ctx, entry, vec![u32_ty.into()], vec![offset_res]); + let loaded_res = loaded.deref(&ctx).get_result(0); + operation::(&mut ctx, entry, vec![], vec![loaded_res]); + + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_mapping_module_without_cutlass_envelope(&target, "memory"); + assert!( + text.contains("\"arith.trunci\"") + && text.contains(": (!llvm.ptr<3>, i32) -> (!llvm.ptr<3>)"), + "shared pointer offset must truncate i64 index to i32 before GEP:\n{text}" + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/mod.rs new file mode 100644 index 0000000000..beaed27ffd --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/mod.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Private MIR scalar, aggregate, checked-storage, and shared-memory mappings. + +pub(crate) mod checked; +pub(crate) mod core; +pub(crate) mod memory; +pub(crate) mod shared; +pub(crate) mod storage; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/shared.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/shared.rs new file mode 100644 index 0000000000..d899197755 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/shared.rs @@ -0,0 +1,334 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Static shared declarations, deduplicated inside the executable GPU module. +use crate::{ + MlirAttribute, MlirBlock, MlirOperation, MlirRegion, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; +use dialect_mir::ops::MirSharedAllocOp; +use pliron::{ + context::{Context, Ptr}, + operation::Operation, + r#type::Typed, +}; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; + +/// Register source static shared allocations for later module-level hoisting. +pub(crate) fn register(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_operation::(SharedAllocation) +} +/// Emit a named shared global and an address while preserving element identity and alignment. +struct SharedAllocation; +impl OperationTranslation for SharedAllocation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + s: &mut TranslationSession<'_>, + ) -> Result, String> { + if !input.operands.is_empty() + || input.results.len() != 1 + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("invalid shared allocation shape".into()); + } + for key in input.attributes.keys() { + if !matches!( + key.as_str(), + "elem_type" | "size" | "alloc_key" | "source_key" | "source_name" | "mir_alignment" + ) { + return Err(format!("unsupported shared allocation attribute {key}")); + } + } + let allocation = MirSharedAllocOp::new(source); + let element = allocation + .get_attr_elem_type(ctx) + .ok_or("missing shared element type")? + .get_type(ctx); + let size = allocation + .get_attr_size(ctx) + .ok_or("missing shared size")? + .value() + .to_u64(); + if size == 0 { + return Err("empty static shared allocation".into()); + } + let alignment = allocation + .get_alignment_value(ctx) + .filter(|align| *align != 0) + .or_else(|| crate::mir::memory::source_natural_alignment(ctx, element)) + .ok_or("unknown shared alignment")?; + if alignment == 0 || !alignment.is_power_of_two() { + return Err("invalid shared alignment".into()); + } + let key = allocation + .get_attr_alloc_key(ctx) + .map(|key| String::from((*key).clone())) + .unwrap_or_else(|| format!("unkeyed-{}", input.results[0].id.0)); + let symbol = format!("__cuda_oxide_shared_{:x}", Sha256::digest(key.as_bytes())); + let ty = s.translate_type(element).map_err(|e| e.to_string())?; + let global_type = MlirType::dialect(format!( + "!llvm.array<{size} x {}>", + crate::mir::memory::type_spelling(&ty) + ))?; + let mut global = MlirOperation::new("llvm.mlir.global")?; + // Preserve source identity while deduplicating: signedness and enum + // identity may disappear from the lowered storage type. Removed below. + global.attributes.insert( + "__source_element_identity".into(), + MlirAttribute::String(format!("{element:?}")), + ); + global + .properties + .insert("sym_name".into(), MlirAttribute::String(symbol.clone())); + global.properties.insert( + "global_type".into(), + MlirAttribute::Type(global_type.clone()), + ); + global.properties.insert( + "linkage".into(), + MlirAttribute::dialect("#llvm.linkage")?, + ); + global.properties.insert( + "addr_space".into(), + MlirAttribute::Integer { + value: 3, + ty: MlirType::Integer(32), + }, + ); + global.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: alignment.into(), + ty: MlirType::Integer(64), + }, + ); + let id = s.fresh_value(); + let mut undef = MlirOperation::new("llvm.mlir.undef")?; + undef.results.push(MlirResult { + id, + ty: global_type.clone(), + }); + let mut ret = MlirOperation::new("llvm.return")?; + ret.operands.push(MlirValueUse { + id, + ty: global_type, + }); + global.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: s.fresh_block(), + arguments: vec![], + operations: vec![undef, ret], + }], + }); + global.location = input.location.clone(); + let mut address = MlirOperation::new("llvm.mlir.addressof")?; + address.results = input.results; + address + .properties + .insert("global_name".into(), MlirAttribute::SymbolRef(symbol)); + address.location = input.location; + // The profile hoists these declarations after translation; source IR + // is immutable and its local allocation operation still yields an SSA address. + Ok(vec![global, address]) + } +} + +/// Move shared declarations to module scope and merge only storage-compatible duplicates. +pub(crate) fn hoist_shared_globals(operations: &mut Vec) -> Result<(), String> { + /// Collect nested shared globals, strengthening dynamic-storage alignment and rejecting conflicting declarations. + fn gather( + operations: &mut Vec, + globals: &mut BTreeMap, + ) -> Result<(), String> { + let mut remaining = vec![]; + for mut operation in std::mem::take(operations) { + if operation.name == "llvm.mlir.global" { + let Some(MlirAttribute::String(symbol)) = + operation.properties.get("sym_name").cloned() + else { + return Err("global without symbol".into()); + }; + if let Some(previous) = globals.get_mut(&symbol) { + // All extern shared declarations refer to the same dynamic + // allocation. Its base must meet the strongest requested alignment. + if symbol == "__cuda_oxide_dynamic_shared" { + let alignment = match ( + previous.properties.get("alignment"), + operation.properties.get("alignment"), + ) { + ( + Some(MlirAttribute::Integer { value: a, .. }), + Some(MlirAttribute::Integer { value: b, .. }), + ) => (*a).max(*b), + _ => return Err("dynamic shared declaration lacks alignment".into()), + }; + let value = MlirAttribute::Integer { + value: alignment, + ty: MlirType::Integer(64), + }; + previous + .properties + .insert("alignment".into(), value.clone()); + operation.properties.insert("alignment".into(), value); + } + if previous.properties != operation.properties + || previous.attributes != operation.attributes + { + return Err(format!( + "incompatible declarations for shared allocation {symbol}" + )); + } + } else { + globals.insert(symbol.clone(), operation); + } + } else { + for region in &mut operation.regions { + for block in &mut region.blocks { + gather(&mut block.operations, globals)?; + } + } + remaining.push(operation); + } + } + *operations = remaining; + Ok(()) + } + let mut globals = BTreeMap::new(); + gather(operations, &mut globals)?; + operations.splice( + 0..0, + globals.into_values().map(|mut global| { + global.attributes.remove("__source_element_identity"); + global + }), + ); + Ok(()) +} + +/// Dynamic shared storage for the instruction-level profile. The CuTe profile +/// retains its compiler-native dynamic-storage operation. +pub(crate) fn register_dynamic_shared( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_operation::(ExternShared) +} +/// Address the one external dynamic shared allocation at the requested byte offset. +struct ExternShared; +impl OperationTranslation for ExternShared { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + use pliron::common_traits::Verify; + let source = dialect_mir::ops::MirExternSharedOp::new(source); + source.verify(ctx).map_err(|e| e.to_string())?; + input.attributes.remove("extern_byte_offset"); + input.attributes.remove("extern_alignment"); + if !input.attributes.is_empty() + || !input.operands.is_empty() + || input.results.len() != 1 + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("dynamic shared storage expects one shared pointer result".into()); + } + let alignment = source.get_alignment_value(ctx); + if alignment == 0 || !alignment.is_power_of_two() { + return Err("dynamic shared alignment must be a nonzero power of two".into()); + } + let offset = source.get_byte_offset_value(ctx); + if offset > i64::MAX as u64 { + return Err("dynamic shared byte offset exceeds i64".into()); + } + let symbol = "__cuda_oxide_dynamic_shared"; + let mut global = MlirOperation::new("llvm.mlir.global")?; + global + .properties + .insert("sym_name".into(), MlirAttribute::String(symbol.into())); + global.properties.insert( + "global_type".into(), + MlirAttribute::Type(MlirType::dialect("!llvm.array<0 x i8>")?), + ); + global.properties.insert( + "linkage".into(), + MlirAttribute::dialect("#llvm.linkage")?, + ); + global.properties.insert( + "addr_space".into(), + MlirAttribute::Integer { + value: 3, + ty: MlirType::Integer(32), + }, + ); + global.properties.insert( + "alignment".into(), + MlirAttribute::Integer { + value: alignment.into(), + ty: MlirType::Integer(64), + }, + ); + global.regions.push(MlirRegion { blocks: vec![] }); + global.location = input.location.clone(); + let mut address = MlirOperation::new("llvm.mlir.addressof")?; + address.properties.insert( + "global_name".into(), + MlirAttribute::SymbolRef(symbol.into()), + ); + address.location = input.location.clone(); + if offset == 0 { + address.results = input.results; + return Ok(vec![global, address]); + } + let base = MlirResult { + id: session.fresh_value(), + ty: input.results[0].ty.clone(), + }; + address.results.push(base.clone()); + let index = MlirResult { + id: session.fresh_value(), + ty: MlirType::Integer(64), + }; + let mut constant = MlirOperation::new("arith.constant")?; + constant.results.push(index.clone()); + constant.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: offset.into(), + ty: index.ty.clone(), + }, + ); + constant.location = input.location.clone(); + let mut gep = MlirOperation::new("llvm.getelementptr")?; + gep.operands = vec![ + MlirValueUse { + id: base.id, + ty: base.ty, + }, + MlirValueUse { + id: index.id, + ty: index.ty, + }, + ]; + gep.results = input.results; + gep.properties.insert( + "elem_type".into(), + MlirAttribute::Type(MlirType::Integer(8)), + ); + gep.properties.insert( + "rawConstantIndices".into(), + MlirAttribute::DenseI32Array(vec![i32::MIN]), + ); + gep.location = input.location; + Ok(vec![global, address, constant, gep]) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/storage.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/storage.rs new file mode 100644 index 0000000000..2e90f44ce5 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/storage.rs @@ -0,0 +1,300 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Byte-preserving storage for pointer-free unions and aggregate reinterpretation. +//! +//! Independent byte lanes preserve partially initialized union contents. Typed +//! scratch stores and loads reinterpret those bytes without numeric conversion. +//! Pointer provenance and bool's canonical storage require separate adapters and +//! remain unsupported here. + +use crate::{ + MlirOperation, MlirType, OperationInput, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use dialect_mir::{ + ops::{MirCastOp, MirExtractFieldOp, MirInsertFieldOp}, + types::{ + MirArrayType, MirDisjointSliceType, MirEnumType, MirFP16Type, MirPtrType, MirSliceType, + MirStructType, MirTupleType, MirUnionType, + }, +}; +use pliron::{ + builtin::types::{FP16Type, FP32Type, FP64Type, IntegerType, UnitType}, + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +use super::{ + checked::{load, op, slot, store}, + memory::{source_abi_align, source_stored_size}, +}; + +/// Register union types without changing the scalar or pointer representation. +pub(super) fn register(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_type::(UnionTypeTranslation) +} + +/// Reject values whose LLVM storage needs a canonicalization or provenance adapter. +fn check_byte_storage(ctx: &Context, source: TypeHandle) -> Result<(), String> { + let ty = source.deref(ctx); + if let Some(integer) = ty.downcast_ref::() { + return if matches!(integer.width(), 8 | 16 | 32 | 64 | 128) { + Ok(()) + } else { + Err("byte storage requires whole-byte scalar integers; bool needs an adapter".into()) + }; + } + if ty.is::() || ty.is::() || ty.is::() { + return Err("pointer-bearing byte storage requires a provenance-preserving adapter".into()); + } + if ty.is::() + || ty.is::() + || ty.is::() + || ty.is::() + || ty.is::() + { + return Ok(()); + } + if let Some(union) = ty.downcast_ref::() { + union_storage_type(ctx, union)?; + return Ok(()); + } + if let Some(array) = ty.downcast_ref::() { + return check_byte_storage(ctx, array.element_type()); + } + let fields = if let Some(structure) = ty.downcast_ref::() { + structure.field_types() + } else if let Some(tuple) = ty.downcast_ref::() { + tuple.get_types() + } else if let Some(enumeration) = ty.downcast_ref::() { + super::checked::enum_storage_type(enumeration)?; + &enumeration.all_field_types + } else { + return Err("unreviewed type in byte storage".into()); + }; + for &field in fields { + check_byte_storage(ctx, field)?; + } + Ok(()) +} + +/// Validate every union view before choosing independent bytes plus an alignment anchor. +pub(super) fn union_storage_type(ctx: &Context, union: &MirUnionType) -> Result { + union.verify(ctx).map_err(|error| error.to_string())?; + let size = union.total_size(); + let align = union.abi_align(); + if !matches!(align, 1 | 2 | 4 | 8 | 16) || !size.is_multiple_of(align) { + return Err("union storage needs a size divisible by alignment 1, 2, 4, 8 or 16".into()); + } + for &field in union.field_types() { + check_byte_storage(ctx, field)?; + let field_size = source_stored_size(ctx, field).ok_or("unknown union field size")?; + let field_align = source_abi_align(ctx, field).ok_or("unknown union field alignment")?; + if field_size > size || field_align > align { + return Err("union field exceeds its recorded storage size or alignment".into()); + } + } + MlirType::dialect(format!( + "!llvm.struct<(array<{size} x i8>, array<0 x i{}>)>", + align * 8 + )) +} + +/// Translate only union layouts whose bytes can be moved without adapters. +struct UnionTypeTranslation; + +impl TypeTranslation for UnionTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _: &TranslationRegistry, + ) -> Result { + union_storage_type( + ctx, + source + .deref(ctx) + .downcast_ref::() + .ok_or("expected a MIR union")?, + ) + } +} + +/// Read one declared field from byte zero while retaining its typed representation. +pub(crate) fn extract_union_field( + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + index: usize, +) -> Result, String> { + MirExtractFieldOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + check_shape(&input, 1)?; + let source_type = source.deref(ctx).get_operand(0).get_type(ctx); + let union_ref = source_type.deref(ctx); + let union = union_ref + .downcast_ref::() + .ok_or("expected a union")?; + let storage = union_storage_type(ctx, union)?; + let field = union + .get_field_type(index) + .ok_or("union field index is out of bounds")?; + let result = input.results[0].clone(); + if source_stored_size(ctx, field) == Some(0) { + return Ok(vec![op( + "ub.poison", + vec![], + vec![result], + &input.location, + )?]); + } + let mut operations = vec![]; + let pointer = slot( + storage, + union.abi_align(), + session, + &input.location, + &mut operations, + )?; + store( + input.operands[0].clone(), + pointer.clone(), + &input.location, + &mut operations, + )?; + load(pointer, result, &input.location, &mut operations)?; + Ok(operations) +} + +/// Overwrite one union view at byte zero, preserving all other bytes. +pub(crate) fn insert_union_field( + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + index: usize, +) -> Result, String> { + MirInsertFieldOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + check_shape(&input, 2)?; + let source_type = source.deref(ctx).get_operand(0).get_type(ctx); + let union_ref = source_type.deref(ctx); + let union = union_ref + .downcast_ref::() + .ok_or("expected a union")?; + let storage = union_storage_type(ctx, union)?; + let field = union + .get_field_type(index) + .ok_or("union field index is out of bounds")?; + let result = input.results[0].clone(); + if source_stored_size(ctx, field) == Some(0) { + return Ok(vec![op( + "builtin.unrealized_conversion_cast", + vec![input.operands[0].clone()], + vec![result], + &input.location, + )?]); + } + let mut operations = vec![]; + let pointer = slot( + storage, + union.abi_align(), + session, + &input.location, + &mut operations, + )?; + store( + input.operands[0].clone(), + pointer.clone(), + &input.location, + &mut operations, + )?; + store( + input.operands[1].clone(), + pointer.clone(), + &input.location, + &mut operations, + )?; + load(pointer, result, &input.location, &mut operations)?; + Ok(operations) +} + +/// Identify aggregate casts while leaving scalar and pointer rules with their existing mapping. +pub(super) fn is_aggregate(ctx: &Context, source: TypeHandle) -> bool { + let ty = source.deref(ctx); + ty.is::() + || ty.is::() + || ty.is::() + || ty.is::() + || ty.is::() +} + +/// Reinterpret equal-sized, byte-faithful aggregates using one private aligned scratch slot. +pub(super) fn transmute( + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, +) -> Result, String> { + MirCastOp::new(source) + .verify(ctx) + .map_err(|error| error.to_string())?; + check_shape(&input, 1)?; + let from = source.deref(ctx).get_operand(0).get_type(ctx); + let to = source.deref(ctx).get_result(0).get_type(ctx); + check_byte_storage(ctx, from)?; + check_byte_storage(ctx, to)?; + let size = source_stored_size(ctx, from).ok_or("unknown transmute source size")?; + if source_stored_size(ctx, to) != Some(size) { + return Err("aggregate transmute requires equal stored sizes".into()); + } + let result = input.results[0].clone(); + if size == 0 { + return Ok(vec![op( + "ub.poison", + vec![], + vec![result], + &input.location, + )?]); + } + let alignment = source_abi_align(ctx, from) + .ok_or("unknown transmute source alignment")? + .max(source_abi_align(ctx, to).ok_or("unknown transmute target alignment")?); + let mut operations = vec![]; + let pointer = slot( + input.operands[0].ty.clone(), + alignment, + session, + &input.location, + &mut operations, + )?; + store( + input.operands[0].clone(), + pointer.clone(), + &input.location, + &mut operations, + )?; + load(pointer, result, &input.location, &mut operations)?; + Ok(operations) +} + +/// Keep scratch-based translations limited to one result and ordinary SSA operands. +fn check_shape(input: &OperationInput, operands: usize) -> Result<(), String> { + if input.results.len() != 1 + || input.operands.len() != operands + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("invalid byte-storage operation shape".into()); + } + Ok(()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/attention.rs new file mode 100644 index 0000000000..450e3b8a68 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/attention.rs @@ -0,0 +1,216 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Asynchronous-copy control and reviewed scalar instruction mappings. +//! +//! Selectors and canonical intrinsic identities are checked before dropping +//! source attributes. Common numeric operations share their precise mappings. + +use crate::nvvm::sm120::{constant_integer, fresh, operation, property, require_input}; +use crate::{ + DropAttribute, MlirAttribute, MlirOperation, MlirType, OperationInput, OperationTranslation, + TranslationError, TranslationRegistry, TranslationSession, +}; +use dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Register the instruction variants used by the Ampere attention schedule. +pub fn register_nvvm_attention_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(MathTranslation)?; + registry.register_operation::(ShuffleTranslation)?; + register_nvvm_async_control_pack(registry) +} + +/// Register asynchronous-copy commit and compile-time wait-count operations. +pub(crate) fn register_nvvm_async_control_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_operation::(CommitTranslation)?; + registry.register_operation::(WaitTranslation)?; + Ok(()) +} + +/// Map reviewed attention math while preserving approximation, subnormal mode, and identity. +struct MathTranslation; +impl OperationTranslation for MathTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let math = ScalarMathOp::new(source); + math.verify(ctx).map_err(|e| e.to_string())?; + // Log2, reciprocal square root and sine use the common numeric mapping; + // exp2 and reciprocal retain the SM120 attention schedule's + // validated CUTLASS lowering below. + if matches!( + math.get_attr_nvvm_scalar_math_operation(ctx).as_deref(), + Some( + ScalarMathOperationAttr::Lg2 + | ScalarMathOperationAttr::Rsqrt + | ScalarMathOperationAttr::Sin + ) + ) { + return crate::nvvm::numeric::MathTranslation.translate(ctx, source, input, session); + } + let selectors = ( + math.get_attr_nvvm_scalar_math_format(ctx), + math.get_attr_nvvm_scalar_math_operation(ctx), + math.get_attr_nvvm_scalar_math_precision(ctx), + math.get_attr_nvvm_scalar_math_subnormal(ctx), + ); + let (target, marker) = + match ( + selectors.0.as_deref(), + selectors.1.as_deref(), + selectors.2.as_deref(), + selectors.3.as_deref(), + ) { + ( + Some(ScalarMathFormatAttr::F32), + Some(ScalarMathOperationAttr::Ex2), + Some(ScalarMathPrecisionAttr::Approx), + Some(ScalarMathSubnormalAttr::Preserve), + ) => ("llvm.inline_asm", "v1:i0819"), + ( + Some(ScalarMathFormatAttr::F32), + Some(ScalarMathOperationAttr::Ex2), + Some(ScalarMathPrecisionAttr::Approx), + Some(ScalarMathSubnormalAttr::Ftz), + ) => ("llvm.inline_asm", "v1:i0820"), + ( + Some(ScalarMathFormatAttr::F32), + Some(ScalarMathOperationAttr::Rcp), + Some(ScalarMathPrecisionAttr::Approx), + Some(ScalarMathSubnormalAttr::Ftz), + ) => ("nvvm.rcp.approx.ftz.f", "v1:i0788"), + _ => return Err( + "attention math mapping supports F32 approximate exp2 and FTZ reciprocal only" + .into(), + ), + }; + require_input(&mut input, marker, 1, 1)?; + let mut op = operation(target, input.operands, input.results, &input.location)?; + if matches!(marker, "v1:i0819" | "v1:i0820") { + // CUTLASS 4.7 parses nvvm.ex2 and llvm.nvvm.ex2.approx.f, but its + // bundled libNVVM rejects their lowering. Keep each typed + // intrinsic's exact subnormal mode through this compiler fallback. + // The FTZ form also matches CuTeDSL 4.7 math.exp2(fastmath=True). + // There are no memory effects, so the pure asm needs no sideeffect. + let instruction = if marker == "v1:i0820" { + "ex2.approx.ftz.f32 $0, $1;" + } else { + "ex2.approx.f32 $0, $1;" + }; + op.properties.insert( + "asm_string".into(), + MlirAttribute::String(instruction.into()), + ); + op.properties + .insert("constraints".into(), MlirAttribute::String("=f,f".into())); + } + Ok(vec![op]) + } +} + +/// Preserve butterfly-shuffle membership, lane offset, and warp clamp. +struct ShuffleTranslation; +impl OperationTranslation for ShuffleTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + ShflSyncBflyF32Op::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0055", 3, 1)?; + let (result, clamp) = fresh(session, MlirType::Integer(32)); + let mut constant = operation("arith.constant", vec![], vec![result], &input.location)?; + constant.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 31, + ty: MlirType::Integer(32), + }, + ); + input.operands.push(clamp); + let mut shuffle = operation( + "nvvm.shfl.sync", + input.operands, + input.results, + &input.location, + )?; + property(&mut shuffle, "kind", "#nvvm")?; + Ok(vec![constant, shuffle]) + } +} + +/// Commit preceding asynchronous copies into a hardware group. +struct CommitTranslation; +impl OperationTranslation for CommitTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + CpAsyncCommitGroupOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0094", 0, 0)?; + Ok(vec![operation( + "nvvm.cp.async.commit.group", + vec![], + vec![], + &input.location, + )?]) + } +} + +/// Wait for a legal compile-time number of outstanding asynchronous-copy groups. +struct WaitTranslation; +impl OperationTranslation for WaitTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + CpAsyncWaitGroupOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0096", 1, 0)?; + let count = constant_integer(ctx, source, 0)?; + if count > 7 { + return Err("cp.async.wait_group requires an immediate group count in 0..=7".into()); + } + let mut wait = operation("nvvm.cp.async.wait.group", vec![], vec![], &input.location)?; + wait.properties.insert( + "n".into(), + MlirAttribute::Integer { + value: count.into(), + ty: MlirType::Integer(32), + }, + ); + Ok(vec![wait]) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/cluster.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/cluster.rs new file mode 100644 index 0000000000..de5ae45285 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/cluster.rs @@ -0,0 +1,1315 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit cluster, mbarrier, and TMA mappings for the CUTLASS MLIR consumer. +//! +//! Native intrinsic markers remain part of the contract. Pointer conversions +//! preserve the native address spaces, and try-wait preserves its boolean +//! result rather than becoming the similarly named blocking NVVM operation. + +use crate::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; +use dialect_mir::ops::MirConstantOp; +use dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, +}; + +/// Register the native primitives used by persistent two-CTA TMA kernels. +pub fn register_nvvm_cluster_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + /// Register one source operation with its marker, selector recipe, and source verifier. + macro_rules! register { + ($op:ty, $marker:literal, $kind:expr) => { + registry.register_operation::<$op>(Recipe { + marker: $marker, + kind: $kind, + verify: |ctx, source| { + <$op>::new(source) + .verify(ctx) + .map_err(|error| error.to_string()) + }, + })?; + }; + } + use Kind::*; + register!(TrapOp, "v1:i0295", NoOperands("llvm.intr.trap")); + register!( + ReadPtxSregClusterCtarankOp, + "v1:i0275", + Sreg("nvvm.read.ptx.sreg.cluster.ctarank") + ); + register!( + ReadPtxSregClusterNctarankOp, + "v1:i0276", + Sreg("nvvm.read.ptx.sreg.cluster.nctarank") + ); + register!( + ReadPtxSregClusterCtaidXOp, + "v1:i0263", + Sreg("nvvm.read.ptx.sreg.cluster.ctaid.x") + ); + register!( + ReadPtxSregClusterCtaidYOp, + "v1:i0264", + Sreg("nvvm.read.ptx.sreg.cluster.ctaid.y") + ); + register!( + ReadPtxSregClusterCtaidZOp, + "v1:i0265", + Sreg("nvvm.read.ptx.sreg.cluster.ctaid.z") + ); + register!( + ReadPtxSregClusterNctaidXOp, + "v1:i0266", + Sreg("nvvm.read.ptx.sreg.cluster.nctaid.x") + ); + register!( + ReadPtxSregClusterNctaidYOp, + "v1:i0267", + Sreg("nvvm.read.ptx.sreg.cluster.nctaid.y") + ); + register!( + ReadPtxSregClusterNctaidZOp, + "v1:i0268", + Sreg("nvvm.read.ptx.sreg.cluster.nctaid.z") + ); + register!( + ReadPtxSregClusterIdXOp, + "v1:i0269", + Sreg("nvvm.read.ptx.sreg.clusterid.x") + ); + register!( + ReadPtxSregClusterIdYOp, + "v1:i0270", + Sreg("nvvm.read.ptx.sreg.clusterid.y") + ); + register!( + ReadPtxSregClusterIdZOp, + "v1:i0271", + Sreg("nvvm.read.ptx.sreg.clusterid.z") + ); + register!( + ReadPtxSregNclusterIdXOp, + "v1:i0272", + Sreg("nvvm.read.ptx.sreg.nclusterid.x") + ); + register!( + ReadPtxSregNclusterIdYOp, + "v1:i0273", + Sreg("nvvm.read.ptx.sreg.nclusterid.y") + ); + register!( + ReadPtxSregNclusterIdZOp, + "v1:i0274", + Sreg("nvvm.read.ptx.sreg.nclusterid.z") + ); + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(ClusterBarrierTranslation)?; + register!( + FenceMbarrierInitReleaseClusterOp, + "v1:i0313", + NoOperands("nvvm.fence.mbarrier.init") + ); + register!(FenceProxyAsyncSharedCtaOp, "v1:i0312", AsyncFence); + register!( + FenceAcquireSharedClusterClusterOp, + "v1:i1042", + SharedAcquireFence + ); + register!( + BarrierCtaSyncCountOp, + "v1:i0883", + CountedBarrier { + arrive: false, + aligned: false + } + ); + register!( + BarrierCtaSyncAlignedCountOp, + "v1:i0884", + CountedBarrier { + arrive: false, + aligned: true + } + ); + register!( + BarrierCtaArriveCountOp, + "v1:i0885", + CountedBarrier { + arrive: true, + aligned: false + } + ); + register!( + BarrierCtaArriveAlignedCountOp, + "v1:i0886", + CountedBarrier { + arrive: true, + aligned: true + } + ); + register!(MbarrierInitSharedOp, "v1:i0097", BarrierInit); + register!(MbarrierInvalSharedOp, "v1:i0100", BarrierInval); + register!(MbarrierArriveSharedOp, "v1:i0098", BarrierArrive); + register!( + MbarrierArriveExpectTxSharedOp, + "v1:i0306", + ExpectTx { cluster: false } + ); + register!( + MbarrierArriveExpectTxClusterOp, + "v1:i0307", + ExpectTx { cluster: true } + ); + register!( + MbarrierTryWaitParitySharedOp, + "v1:i0310", + WaitParity { + cluster: false, + relaxed: false + } + ); + register!( + MbarrierTryWaitParityClusterOp, + "v1:i0311", + WaitParity { + cluster: true, + relaxed: false + } + ); + register!( + MbarrierArriveClusterOp, + "v1:i0308", + RemoteArrive { relaxed: false } + ); + register!( + MbarrierArriveClusterRelaxedOp, + "v1:i1040", + RemoteArrive { relaxed: true } + ); + register!( + MbarrierTryWaitParityClusterRelaxedOp, + "v1:i1041", + WaitParity { + cluster: true, + relaxed: true + } + ); + register!(MapaSharedClusterOp, "v1:i0320", MapShared); + register!(ElectSyncOp, "v1:i0367", Elect); + register!(PrefetchTensorMapOp, "v1:i0887", Prefetch); + register!( + CpAsyncBulkTensorG2sTile2dMulticastCg2Op, + "v1:i0331", + TmaLoadCg2 + ); + register!(CpAsyncBulkTensorS2gTile2dOp, "v1:i0336", TmaStore); + register!( + CpAsyncBulkCommitGroupOp, + "v1:i0340", + NoOperands("nvvm.cp.async.bulk.commit.group") + ); + register!( + CpAsyncBulkWaitGroupOp, + "v1:i0341", + WaitGroup { read: false } + ); + register!( + CpAsyncBulkWaitGroupReadOp, + "v1:i0342", + WaitGroup { read: true } + ); + Ok(()) +} + +/// Select one reviewed cluster instruction contract and its structural selectors. +enum Kind { + /// Read the named cluster special register. + Sreg(#[doc = "Target NVVM operation spelling for the selected cluster register."] &'static str), + /// Emit the named result-free synchronization instruction. + NoOperands(#[doc = "Target NVVM synchronization operation spelling."] &'static str), + /// Publish generic shared accesses to the asynchronous proxy. + AsyncFence, + /// Acquire ordering restricted to cluster shared-memory accesses. + SharedAcquireFence, + /// Emit arrival or synchronization on a named CTA barrier with explicit participation. + CountedBarrier { + #[doc = "Emit nonblocking arrival instead of blocking synchronization."] + arrive: bool, + #[doc = "Preserve the source aligned-participation instruction variant."] + aligned: bool, + }, + /// Initialize an mbarrier with its expected arrival count. + BarrierInit, + /// Invalidate an mbarrier without adding synchronization. + BarrierInval, + /// Arrive on a local mbarrier and preserve its result contract. + BarrierArrive, + /// Announce asynchronous transaction bytes with the selected barrier scope. + ExpectTx { + #[doc = "Use cluster scope instead of CTA scope."] + cluster: bool, + }, + /// Probe barrier phase completion and preserve its boolean predicate result. + WaitParity { + #[doc = "Use cluster scope instead of CTA scope."] + cluster: bool, + #[doc = "Select relaxed memory ordering rather than acquire/release ordering."] + relaxed: bool, + }, + /// Arrive on a cluster-mapped barrier using the selected memory ordering. + RemoteArrive { + #[doc = "Select relaxed memory ordering rather than acquire/release ordering."] + relaxed: bool, + }, + /// Preserve CTA rank when mapping a shared address across the cluster. + MapShared, + /// Return both the elected lane and its participation predicate. + Elect, + /// Prefetch the tensor-map descriptor without changing its address space contract. + Prefetch, + /// Issue a two-CTA tensor load with remote-barrier completion routing. + TmaLoadCg2, + /// Issue shared-to-global tensor stores in NVVM operand order. + TmaStore, + /// Wait for a legal compile-time number of pending bulk groups. + WaitGroup { + #[doc = "Wait only until shared source reads complete when set."] + read: bool, + }, +} + +/// Combine intrinsic identity, source verification, and one concrete cluster lowering. +struct Recipe { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// Concrete cluster instruction and selector contract for this recipe. + kind: Kind, + /// Check the source operation's selectors and operand contract. + verify: fn(&Context, Ptr) -> Result<(), String>, +} + +impl OperationTranslation for Recipe { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + check_marker(&mut input, self.marker)?; + let location = &input.location; + let mut operations = vec![]; + use TypeSpec::*; + let mut target = match self.kind { + Kind::Sreg(name) => { + signature(&input, &[], &[Int(32)])?; + operation(name, input.results, input.operands, location)? + } + Kind::NoOperands(name) => { + signature(&input, &[], &[])?; + operation(name, vec![], vec![], location)? + } + Kind::SharedAcquireFence => { + signature(&input, &[], &[])?; + let mut op = operation("nvvm.fence.sync_restrict", vec![], vec![], location)?; + dialect_property(&mut op, "order", "#nvvm.mem_order")?; + op + } + Kind::AsyncFence => { + signature(&input, &[], &[])?; + let mut op = operation("nvvm.fence.proxy", vec![], vec![], location)?; + dialect_property(&mut op, "kind", "#nvvm.proxy_kind")?; + dialect_property(&mut op, "space", "#nvvm.shared_space")?; + op + } + Kind::CountedBarrier { arrive, aligned } => { + signature(&input, &[Int(32), Int(32)], &[])?; + let name = if arrive { + "nvvm.barrier.cta.arrive" + } else { + "nvvm.barrier.cta.sync" + }; + let mut op = operation(name, vec![], input.operands, location)?; + op.properties + .insert("aligned".into(), MlirAttribute::Bool(aligned)); + op + } + Kind::BarrierInit => { + signature(&input, &[Pointer(&[0, 3]), Int(32)], &[])?; + operation("nvvm.mbarrier.init", vec![], input.operands, location)? + } + Kind::BarrierInval => { + signature(&input, &[Pointer(&[0, 3])], &[])?; + operation("nvvm.mbarrier.inval", vec![], input.operands, location)? + } + Kind::BarrierArrive => { + signature(&input, &[Pointer(&[0, 3])], &[Int(64)])?; + operation( + "nvvm.mbarrier.arrive", + input.results, + input.operands, + location, + )? + } + Kind::ExpectTx { cluster } => { + // The native semantic op has already discarded the compatibility + // tx_count argument. PTX's arrival decrement is always one. + signature(&input, &[Pointer(&[0, 3]), Int(32)], &[Int(64)])?; + let mut op = operation( + "nvvm.mbarrier.arrive.expect_tx", + input.results, + input.operands, + location, + )?; + if cluster { + dialect_property(&mut op, "scope", "#nvvm.mem_scope")?; + op.properties + .insert("relaxed".into(), MlirAttribute::Bool(true)); + } + op + } + Kind::WaitParity { cluster, relaxed } => { + signature(&input, &[Pointer(&[0, 3]), Int(32)], &[Int(1)])?; + let mut op = operation( + "nvvm.mbarrier.wait.parity", + input.results, + input.operands, + location, + )?; + dialect_property(&mut op, "kind", "#nvvm.mbar_wait")?; + if cluster { + dialect_property(&mut op, "scope", "#nvvm.mbar_scope")?; + dialect_property( + &mut op, + "order", + if relaxed { + "#nvvm.mem_order" + } else { + "#nvvm.mem_order" + }, + )?; + } + op + } + Kind::RemoteArrive { relaxed } => { + signature(&input, &[Int(64)], &[])?; + let pointer = cast( + session, + &mut operations, + "llvm.inttoptr", + input.operands[0].clone(), + pointer_type(7)?, + location, + )?; + let mut op = operation("nvvm.mbarrier.arrive", vec![], vec![pointer], location)?; + dialect_property(&mut op, "scope", "#nvvm.mem_scope")?; + if relaxed { + op.properties + .insert("relaxed".into(), MlirAttribute::Bool(true)); + } + op + } + Kind::MapShared => { + signature(&input, &[Pointer(&[0, 3]), Int(32)], &[Pointer(&[7])])?; + let shared = pointer_cast( + session, + &mut operations, + input.operands[0].clone(), + 3, + location, + )?; + operation( + "nvvm.mapa", + input.results, + vec![shared, input.operands[1].clone()], + location, + )? + } + Kind::Elect => { + signature(&input, &[Int(32)], &[Int(32), Int(1)])?; + // NVVM dialect elect.sync only exposes the predicate. The LLVM + // intrinsic preserves both native results, including lane ID. + let ty = MlirType::dialect("!llvm.struct<(i32, i1)>")?; + let (result, value) = fresh(session, ty); + let mut call = operation( + "llvm.call_intrinsic", + vec![result], + input.operands, + location, + )?; + call.properties.insert( + "intrin".into(), + MlirAttribute::String("llvm.nvvm.elect.sync".into()), + ); + call.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + call.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 0]), + ); + operations.push(call); + for (index, result) in input.results.into_iter().enumerate() { + let mut extract = operation( + "llvm.extractvalue", + vec![result], + vec![value.clone()], + location, + )?; + extract.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + operations.push(extract); + } + return Ok(operations); + } + Kind::Prefetch => { + signature(&input, &[Pointer(&[0, 1, 4])], &[])?; + let descriptor = pointer_cast( + session, + &mut operations, + input.operands[0].clone(), + 0, + location, + )?; + let mut op = operation("nvvm.prefetch", vec![], vec![descriptor], location)?; + op.properties + .insert("tensormap".into(), MlirAttribute::Unit); + op + } + Kind::TmaLoadCg2 => { + // Importer order follows LLVM: dst, barrier, descriptor, x, y, + // multicast mask, zero cache hint. This differs from the API. + signature( + &input, + &[ + Pointer(&[0, 3, 7]), + Pointer(&[0, 3, 7]), + Pointer(&[0, 1, 4]), + Int(32), + Int(32), + Int(16), + Int(64), + ], + &[], + )?; + if constant_integer(ctx, source, 6)? != 0 { + return Err( + "cg2 TMA native mapping requires the importer's zero cache hint".into(), + ); + } + // Cluster TMA's LLVM intrinsic requires addrspace(7), including + // when the Rust API began with the local CTA shared pointer. + let dst = pointer_cast( + session, + &mut operations, + input.operands[0].clone(), + 7, + location, + )?; + let descriptor = pointer_cast( + session, + &mut operations, + input.operands[2].clone(), + 0, + location, + )?; + let barrier = if input.operands[1].ty == pointer_type(7)? { + // mapa(rank=0) carries a cluster-space address. The CG2 + // intrinsic encodes its barrier as a shared-space pointer; + // retype the raw address without changing its rank bits. + let address = cast( + session, + &mut operations, + "llvm.ptrtoint", + input.operands[1].clone(), + MlirType::Integer(32), + location, + )?; + cast( + session, + &mut operations, + "llvm.inttoptr", + address, + pointer_type(3)?, + location, + )? + } else { + pointer_cast( + session, + &mut operations, + input.operands[1].clone(), + 3, + location, + )? + }; + let mut op = operation( + "nvvm.cp.async.bulk.tensor.shared.cluster.global", + vec![], + vec![ + dst, + descriptor, + input.operands[3].clone(), + input.operands[4].clone(), + barrier, + input.operands[5].clone(), + ], + location, + )?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 1, 0, 1, 0, 0]), + ); + dialect_property(&mut op, "group", "#nvvm.cta_group")?; + // In pinned CUTLASS 4.7, false selects LLVM's g2s.tile intrinsic + // (addrspace 7). True selects the older CUDA gmem.to.smem + // intrinsic, whose destination contract is addrspace 3. + op.properties + .insert("useIntrinsic".into(), MlirAttribute::Bool(false)); + op + } + Kind::TmaStore => { + signature( + &input, + &[Pointer(&[0, 3]), Pointer(&[0, 1, 4]), Int(32), Int(32)], + &[], + )?; + let src = pointer_cast( + session, + &mut operations, + input.operands[0].clone(), + 3, + location, + )?; + let descriptor = pointer_cast( + session, + &mut operations, + input.operands[1].clone(), + 0, + location, + )?; + // Native API is source-first; NVVM's operation is descriptor-first. + let mut op = operation( + "nvvm.cp.async.bulk.tensor.global.shared.cta", + vec![], + vec![ + descriptor, + src, + input.operands[2].clone(), + input.operands[3].clone(), + ], + location, + )?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 0, 0]), + ); + op + } + Kind::WaitGroup { read } => { + signature(&input, &[Int(32)], &[])?; + let group = constant_group(ctx, source)?; + let mut op = operation("nvvm.cp.async.bulk.wait_group", vec![], vec![], location)?; + op.properties.insert( + "group".into(), + MlirAttribute::Integer { + value: group.into(), + ty: MlirType::Integer(32), + }, + ); + if read { + op.properties.insert("read".into(), MlirAttribute::Unit); + } + op + } + }; + target.location = input.location; + operations.push(target); + Ok(operations) + } +} + +/// Preserve source cluster-barrier mode and native intrinsic identity. +struct ClusterBarrierTranslation; + +impl OperationTranslation for ClusterBarrierTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let op = ClusterBarrierOp::new(source); + op.verify(ctx).map_err(|error| error.to_string())?; + let mode = op + .get_attr_nvvm_cluster_barrier_mode(ctx) + .ok_or("missing cluster barrier mode")?; + use ClusterBarrierModeAttr::*; + let (marker, name, aligned) = match *mode { + Arrive => ("v1:i0277", "nvvm.cluster.arrive", false), + ArriveAligned => ("v1:i0278", "nvvm.cluster.arrive", true), + ArriveRelaxed => ("v1:i0279", "nvvm.cluster.arrive.relaxed", false), + ArriveRelaxedAligned => ("v1:i0280", "nvvm.cluster.arrive.relaxed", true), + Wait => ("v1:i0281", "nvvm.cluster.wait", false), + WaitAligned => ("v1:i0282", "nvvm.cluster.wait", true), + }; + check_marker(&mut input, marker)?; + signature(&input, &[], &[])?; + let mut target = operation(name, vec![], vec![], &input.location)?; + if aligned { + target + .properties + .insert("aligned".into(), MlirAttribute::Unit); + } + Ok(vec![target]) + } +} + +/// Consume and validate the canonical intrinsic marker before target emission. +fn check_marker(input: &mut OperationInput, expected: &str) -> Result<(), String> { + let marker = input.attributes.remove("cuda_oxide_intrinsic_marker"); + if marker != Some(MlirAttribute::String(expected.into())) { + return Err(format!( + "expected intrinsic marker {expected}, got {marker:?}" + )); + } + if !input.attributes.is_empty() || !input.regions.is_empty() || !input.successors.is_empty() { + return Err( + "native NVVM primitive has unexpected attributes, regions, or successors".into(), + ); + } + Ok(()) +} + +/// Operand/result type constraint used by explicit NVVM signatures. +#[derive(Clone, Copy, Debug)] +enum TypeSpec { + /// An integer scalar with the required bit width. + Int(#[doc = "Required integer width in bits."] u32), + /// An opaque pointer in one of the accepted numeric address spaces. + Pointer(#[doc = "Accepted numeric LLVM pointer address spaces."] &'static [u32]), +} + +impl TypeSpec { + /// Check a translated operand type against this signature constraint. + fn matches(self, ty: &MlirType) -> bool { + match self { + Self::Int(bits) => *ty == MlirType::Integer(bits), + Self::Pointer(spaces) => spaces + .iter() + .any(|space| pointer_type(*space).as_ref() == Ok(ty)), + } + } +} + +/// Check operand/result counts and each translated type against the reviewed signature. +fn signature( + input: &OperationInput, + operands: &[TypeSpec], + results: &[TypeSpec], +) -> Result<(), String> { + if input.operands.len() != operands.len() || input.results.len() != results.len() { + return Err(format!( + "expected {} operands and {} results, got {} and {}", + operands.len(), + results.len(), + input.operands.len(), + input.results.len() + )); + } + for (index, (expected, actual)) in operands.iter().zip(&input.operands).enumerate() { + if !expected.matches(&actual.ty) { + return Err(format!( + "operand {index} expected {expected:?}, got {:?}", + actual.ty + )); + } + } + for (index, (expected, actual)) in results.iter().zip(&input.results).enumerate() { + if !expected.matches(&actual.ty) { + return Err(format!( + "result {index} expected {expected:?}, got {:?}", + actual.ty + )); + } + } + Ok(()) +} + +/// Require a legal compile-time bulk wait-group count from the source operand. +fn constant_group(ctx: &Context, source: Ptr) -> Result { + let parsed = constant_integer(ctx, source, 0) + .map_err(|_| "bulk wait group must be a constant in 0..=7")?; + if parsed > 7 { + return Err("bulk wait group must be a constant in 0..=7".into()); + } + Ok(parsed as u32) +} + +/// Read an unsigned source integer constant or reject runtime/noninteger values. +fn constant_integer(ctx: &Context, source: Ptr, index: usize) -> Result { + let operand = source.deref(ctx).get_operand(index); + let definition = operand + .defining_op() + .ok_or("expected a constant integer operand")?; + if Operation::get_opid(definition, ctx) != MirConstantOp::get_opid_static() { + return Err("expected a constant integer operand".into()); + } + let value = MirConstantOp::new(definition) + .get_attr_value(ctx) + .ok_or("constant has no value")?; + value + .value() + .to_string_decimal(false) + .parse::() + .map_err(|_| "constant integer exceeds u64".into()) +} + +/// Spell an opaque LLVM pointer for a required NVVM address space. +fn pointer_type(space: u32) -> Result { + MlirType::dialect(if space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{space}>") + }) +} + +/// Build a target operation with the supplied operands, results, and source location. +fn operation( + name: &str, + results: Vec, + operands: Vec, + location: &MlirLocation, +) -> Result { + let mut op = MlirOperation::new(name)?; + op.results = results; + op.operands = operands; + op.location = location.clone(); + Ok(op) +} + +/// Attach a parsed dialect attribute as a target operation property. +fn dialect_property(op: &mut MlirOperation, name: &str, value: &str) -> Result<(), String> { + op.properties + .insert(name.into(), MlirAttribute::dialect(value)?); + Ok(()) +} + +/// Allocate one SSA identifier and return matching result and operand handles. +fn fresh(session: &mut TranslationSession<'_>, ty: MlirType) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +/// Append a one-operand conversion with a fresh result of the requested type. +fn cast( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + name: &str, + operand: MlirValueUse, + ty: MlirType, + location: &MlirLocation, +) -> Result { + let (result, value) = fresh(session, ty); + operations.push(operation(name, vec![result], vec![operand], location)?); + Ok(value) +} + +/// Convert pointer space explicitly while retaining any required integer-address bridge. +fn pointer_cast( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + operand: MlirValueUse, + space: u32, + location: &MlirLocation, +) -> Result { + let ty = pointer_type(space)?; + if operand.ty == ty { + Ok(operand) + } else { + cast( + session, + operations, + "llvm.addrspacecast", + operand, + ty, + location, + ) + } +} + +#[cfg(test)] +mod tests { + //! Contract and rejection tests for native cluster, barrier, and TMA mappings. + + use super::*; + use crate::{ + CutlassFullCuteMlir22, MlirConsumerProfile, + profile::render_mapping_module_without_cutlass_envelope, + }; + use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::MirPtrType, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + identifier::Identifier, + r#type::{TypeHandle, TypedHandle}, + utils::apint::APInt, + value::Value, + }; + use std::num::NonZeroUsize; + + /// Typed source function exposing descriptor, shared-address, parity, and coordinate test operands. + struct Fixture { + /// Source IR context owning the fixture's registered types and operations. + ctx: Context, + /// Source module containing the operations under test. + module: ModuleOp, + /// Entry block where the fixture appends native source instructions. + entry: Ptr, + /// Argument carrier types used to construct well-typed source instructions. + types: Vec, + /// Unsigned 32-bit type used for narrow immediate operands. + u32: TypedHandle, + /// Unsigned 64-bit type used for wide immediate operands. + u64: TypedHandle, + } + + impl Fixture { + /// Build the canonical cluster test function with distinct address-space and scalar arguments. + fn new() -> Self { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let u8 = IntegerType::get(&ctx, 8, Signedness::Unsigned); + let u32 = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let u64 = IntegerType::get(&ctx, 64, Signedness::Unsigned); + // Arguments: descriptor, shared, remote shared, parity, coordinate, + // multicast mask, raw remote address, predicate. + let types = vec![ + MirPtrType::get(&mut ctx, u8.into(), true, 0).into(), + MirPtrType::get(&mut ctx, u8.into(), true, 3).into(), + MirPtrType::get(&mut ctx, u8.into(), true, 7).into(), + u32.into(), + IntegerType::get(&ctx, 32, Signedness::Signed).into(), + IntegerType::get(&ctx, 16, Signedness::Unsigned).into(), + u64.into(), + IntegerType::get(&ctx, 1, Signedness::Signless).into(), + ]; + let module = ModuleOp::new(&mut ctx, Identifier::try_from("cluster_test").unwrap()); + let function_type = FunctionType::get(&ctx, types.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = + MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("kernel_primitives").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let region = function_op.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, types.clone()); + entry.insert_at_back(region, &ctx); + Self { + ctx, + module, + entry, + types, + u32, + u64, + } + } + + /// Return the indexed source function argument for constructing a test operation. + fn arg(&self, index: usize) -> Value { + self.entry.deref(&self.ctx).get_argument(index) + } + + /// Append a typed source operation to the fixture block and return its operation handle. + fn append( + &mut self, + marker: &str, + operands: Vec, + result_types: Vec, + ) -> Ptr { + let op = Operation::new( + &mut self.ctx, + O::get_concrete_op_info(), + result_types, + operands, + vec![], + 0, + ); + op.deref_mut(&self.ctx).attributes.set( + Identifier::try_from("cuda_oxide_intrinsic_marker").unwrap(), + StringAttr::new(marker.into()), + ); + op.insert_at_back(self.entry, &self.ctx); + op + } + + /// Append a u32 or u64 source constant according to the requested width. + fn constant(&mut self, value: u64, wide: bool) -> Value { + let ty = if wide { self.u64 } else { self.u32 }; + let width = if wide { 64 } else { 32 }; + let op = Operation::new( + &mut self.ctx, + MirConstantOp::get_concrete_op_info(), + vec![ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(op).set_attr_value( + &self.ctx, + IntegerAttr::new( + ty, + APInt::from_u64(value, NonZeroUsize::new(width).unwrap()), + ), + ); + op.insert_at_back(self.entry, &self.ctx); + op.deref(&self.ctx).get_result(0) + } + + /// Terminate, export, and render the fixture without the separately tested GPU envelope. + fn render(mut self) -> Result { + let ret = Operation::new( + &mut self.ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(self.entry, &self.ctx); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + profile + .translate_module(&self.ctx, &self.module) + .map(|target| { + render_mapping_module_without_cutlass_envelope(&target, "cluster_test") + }) + .map_err(|error| error.to_string()) + } + } + + /// Check native descriptor prefetch accepts generic storage and rejects a shared descriptor pointer. + #[test] + fn descriptor_prefetch_is_native_and_rejects_shared_storage() { + let mut fixture = Fixture::new(); + fixture.append::("v1:i0887", vec![fixture.arg(0)], vec![]); + let text = fixture.render().unwrap(); + assert!(text.contains("\"nvvm.prefetch\"(%v0)"), "{text}"); + assert!(text.contains("tensormap = unit"), "{text}"); + assert!( + !text.contains("inline_ptx") && !text.contains("inline_asm"), + "{text}" + ); + for argument in [1, 2] { + let mut fixture = Fixture::new(); + fixture.append::("v1:i0887", vec![fixture.arg(argument)], vec![]); + assert!(fixture.render().is_err()); + } + } + + /// Check invalidation retains its barrier address without adding an unrelated barrier wait. + #[test] + fn barrier_invalidation_preserves_pointer_and_has_no_hidden_sync() { + for argument in [0, 1] { + let mut fixture = Fixture::new(); + fixture.append::( + "v1:i0100", + vec![fixture.arg(argument)], + vec![], + ); + let text = fixture.render().unwrap(); + assert_eq!(text.matches("\"nvvm.mbarrier.inval\"").count(), 1, "{text}"); + assert!( + !text.contains("nvvm.barrier.cta") && !text.contains("nvvm.fence"), + "{text}" + ); + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } + let mut fixture = Fixture::new(); + fixture.append::("v1:i0097", vec![fixture.arg(1)], vec![]); + assert!( + fixture + .render() + .unwrap_err() + .contains("expected intrinsic marker v1:i0100") + ); + let mut fixture = Fixture::new(); + fixture.append::("v1:i0100", vec![fixture.arg(2)], vec![]); + assert!(fixture.render().is_err()); + } + + /// Check counted arrival/sync forms retain their aligned-participation selectors. + #[test] + fn counted_barrier_forms_preserve_instruction_and_alignment() { + let mut fixture = Fixture::new(); + let id = fixture.constant(3, false); + let count = fixture.constant(128, false); + fixture.append::("v1:i0883", vec![id, count], vec![]); + fixture.append::("v1:i0884", vec![id, count], vec![]); + fixture.append::("v1:i0885", vec![id, count], vec![]); + fixture.append::("v1:i0886", vec![id, count], vec![]); + let text = fixture.render().unwrap(); + assert_eq!( + text.matches("\"nvvm.barrier.cta.sync\"").count(), + 2, + "{text}" + ); + assert_eq!( + text.matches("\"nvvm.barrier.cta.arrive\"").count(), + 2, + "{text}" + ); + assert_eq!(text.matches("aligned = true").count(), 2, "{text}"); + for line in text + .lines() + .filter(|line| line.contains("\"nvvm.barrier.cta.")) + { + assert!(line.contains("(%v8, %v9)"), "{line}"); + } + } + + /// Check cluster synchronization modes and catalog identities are validated and preserved. + #[test] + fn cluster_sync_preserves_mode_and_intrinsic_identity() { + for (mode, marker, name) in [ + ( + ClusterBarrierModeAttr::ArriveRelaxed, + "v1:i0279", + "nvvm.cluster.arrive.relaxed", + ), + ( + ClusterBarrierModeAttr::Wait, + "v1:i0281", + "nvvm.cluster.wait", + ), + ] { + let mut fixture = Fixture::new(); + let op = fixture.append::(marker, vec![], vec![]); + ClusterBarrierOp::new(op).set_attr_nvvm_cluster_barrier_mode(&fixture.ctx, mode); + let text = fixture.render().unwrap(); + assert!(text.contains(name), "{text}"); + } + let mut fixture = Fixture::new(); + let op = fixture.append::("v1:i0281", vec![], vec![]); + ClusterBarrierOp::new(op).set_attr_nvvm_cluster_barrier_mode( + &fixture.ctx, + ClusterBarrierModeAttr::ArriveRelaxed, + ); + assert!( + fixture + .render() + .unwrap_err() + .contains("expected intrinsic marker v1:i0279") + ); + } + + /// Check cluster try-wait produces the completion predicate and retains acquire ordering. + #[test] + fn barrier_wait_returns_predicate_and_preserves_cluster_acquire() { + let mut fixture = Fixture::new(); + fixture.append::( + "v1:i0311", + vec![fixture.arg(1), fixture.arg(3)], + vec![fixture.types[7]], + ); + fixture.append::( + "v1:i0307", + vec![fixture.arg(1), fixture.arg(3)], + vec![fixture.types[6]], + ); + fixture.append::("v1:i0308", vec![fixture.arg(6)], vec![]); + let text = fixture.render().unwrap(); + assert!(text.contains("\"nvvm.mbarrier.wait.parity\""), "{text}"); + assert!(text.contains("kind = #nvvm.mbar_wait"), "{text}"); + assert!(text.contains("order = #nvvm.mem_order"), "{text}"); + assert!(text.contains("scope = #nvvm.mbar_scope"), "{text}"); + assert!(text.contains("relaxed = true"), "{text}"); + assert!( + text.contains("\"llvm.inttoptr\"(%v6) : (i64) -> (!llvm.ptr<7>)"), + "{text}" + ); + assert!(!text.contains("\"nvvm.mbarrier.try_wait.parity\"")); + } + + /// Check relaxed barrier instructions retain cluster scope without gaining stronger ordering. + #[test] + fn relaxed_cluster_arrival_and_wait_keep_scope_without_acquire_release() { + let mut fixture = Fixture::new(); + fixture.append::("v1:i1040", vec![fixture.arg(6)], vec![]); + fixture.append::( + "v1:i1041", + vec![fixture.arg(1), fixture.arg(3)], + vec![fixture.types[7]], + ); + let text = fixture.render().unwrap(); + assert!( + text.contains("\"llvm.inttoptr\"(%v6) : (i64) -> (!llvm.ptr<7>)"), + "{text}" + ); + assert!(text.contains("\"nvvm.mbarrier.arrive\""), "{text}"); + assert!(text.contains("scope = #nvvm.mem_scope"), "{text}"); + assert!(text.contains("relaxed = true"), "{text}"); + assert!(text.contains("order = #nvvm.mem_order"), "{text}"); + assert!(text.contains("scope = #nvvm.mbar_scope"), "{text}"); + assert!(!text.contains("mem_order"), "{text}"); + assert!(!text.contains("inline_ptx"), "{text}"); + } + + /// Check restricted acquire fences do not become unrestricted memory-order operations. + #[test] + fn restricted_acquire_fence_retains_shared_only_ordering() { + let mut fixture = Fixture::new(); + fixture.append::("v1:i1042", vec![], vec![]); + let text = fixture.render().unwrap(); + assert!(text.contains("\"nvvm.fence.sync_restrict\""), "{text}"); + assert!(text.contains("order = #nvvm.mem_order"), "{text}"); + assert!(!text.contains("inline_ptx"), "{text}"); + assert!(!text.contains("fence.proxy"), "{text}"); + } + + /// Check lane election retains both its lane ID and elected-thread predicate. + #[test] + fn elect_preserves_both_lane_and_predicate_results() { + let mut fixture = Fixture::new(); + fixture.append::( + "v1:i0367", + vec![fixture.arg(3)], + vec![fixture.types[3], fixture.types[7]], + ); + let text = fixture.render().unwrap(); + assert!(text.contains("intrin = \"llvm.nvvm.elect.sync\""), "{text}"); + assert!(text.contains("position = array"), "{text}"); + assert!(text.contains("position = array"), "{text}"); + assert!(text.contains("!llvm.struct<(i32, i1)>"), "{text}"); + } + + /// Check tensor loads use NVVM operand order while retaining remote-barrier rank bits. + #[test] + fn tma_reorders_operands_and_preserves_remote_barrier_address() { + let mut fixture = Fixture::new(); + let zero = fixture.constant(0, true); + fixture.append::( + "v1:i0331", + vec![ + fixture.arg(1), + fixture.arg(2), + fixture.arg(0), + fixture.arg(4), + fixture.arg(4), + fixture.arg(5), + zero, + ], + vec![], + ); + fixture.append::( + "v1:i0336", + vec![ + fixture.arg(1), + fixture.arg(0), + fixture.arg(4), + fixture.arg(4), + ], + vec![], + ); + let text = fixture.render().unwrap(); + assert!(text.contains("group = #nvvm.cta_group"), "{text}"); + assert!(text.contains("useIntrinsic = false"), "{text}"); + assert!( + text.contains("operandSegmentSizes = array"), + "{text}" + ); + assert!( + text.contains("\"llvm.ptrtoint\"(%v2) : (!llvm.ptr<7>) -> (i32)"), + "{text}" + ); + assert!( + text.contains("\"nvvm.cp.async.bulk.tensor.global.shared.cta\"(%v0, %v1, %v4, %v4)"), + "{text}" + ); + } + + /// Reject malformed tensor transfer signatures and unsupported cache selectors. + #[test] + fn tma_rejects_incorrect_shape_and_hidden_cache_policy() { + for cache_hint in [None, Some(1)] { + let mut fixture = Fixture::new(); + let zero = fixture.constant(cache_hint.unwrap_or(0), true); + let mut operands = vec![ + fixture.arg(1), + fixture.arg(2), + fixture.arg(0), + fixture.arg(4), + fixture.arg(4), + fixture.arg(5), + zero, + ]; + if cache_hint.is_none() { + operands.pop(); + } + fixture.append::( + "v1:i0331", + operands, + vec![], + ); + let error = fixture.render().unwrap_err(); + assert!( + error.contains(if cache_hint.is_none() { + "Expected 7 operand(s), but found 6" + } else { + "zero cache hint" + }), + "{error}" + ); + } + } + + /// Reject runtime or out-of-range bulk wait counts while preserving legal immediates. + #[test] + fn bulk_wait_requires_legal_compile_time_group_count() { + for group in [0, 1, 7, 8] { + let mut fixture = Fixture::new(); + let count = fixture.constant(group, false); + fixture.append::("v1:i0342", vec![count], vec![]); + let result = fixture.render(); + if group <= 7 { + let text = result.unwrap(); + assert!(text.contains(&format!("group = {group} : i32")), "{text}"); + assert!(text.contains("read = unit"), "{text}"); + } else { + assert!(result.unwrap_err().contains("constant in 0..=7")); + } + } + let mut fixture = Fixture::new(); + fixture.append::("v1:i0342", vec![fixture.arg(3)], vec![]); + assert!(fixture.render().unwrap_err().contains("constant in 0..=7")); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/mod.rs new file mode 100644 index 0000000000..a926e70267 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/mod.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Private NVVM primitive mappings grouped by instruction capability. + +pub(crate) mod attention; +pub(crate) mod cluster; +pub(crate) mod numeric; +pub(crate) mod sm120; +pub(crate) mod sregs; +pub(crate) mod sync_restrict; +pub(crate) mod tcgen05; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs new file mode 100644 index 0000000000..a9e4e9f120 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/numeric.rs @@ -0,0 +1,547 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Shared scalar and packed arithmetic for CuTe and instruction-only profiles. +//! Precision, rounding, lane masks, and intrinsic identity survive translation. +use crate::nvvm::tcgen05::{fresh, operation, require_input}; +use crate::{ + DropAttribute, MlirAttribute, MlirFloatType, MlirOperation, MlirType, OperationInput, + OperationTranslation, TranslationError, TranslationRegistry, TranslationSession, +}; +use dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Source-operation verifier run before discarding its intrinsic metadata. +type VerifyOperation = fn(&Context, Ptr) -> Result<(), String>; + +/// Register reviewed scalar, packed, and warp arithmetic with exact native intrinsic identities. +pub fn register_nvvm_numeric_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(MathTranslation)?; + registry.register_operation::(IntrinsicTranslation { + marker: "v1:i0055", + name: "llvm.nvvm.shfl.sync.bfly.f32", + operands: 3, + shuffle: true, + verify: |ctx, source| { + ShflSyncBflyF32Op::new(source) + .verify(ctx) + .map_err(|e| e.to_string()) + }, + })?; + register_nvvm_numeric_common_pack(registry) +} + +/// Preserve each architecture's validated CUTLASS math lowering while sharing +/// the remaining arithmetic and asynchronous-copy control registrations. +pub(crate) fn register_nvvm_profile_numeric_pack( + registry: &mut TranslationRegistry, + gpu_arch: &str, +) -> Result<(), TranslationError> { + if gpu_arch == "sm_120a" { + register_nvvm_numeric_common_pack(registry)?; + crate::register_nvvm_attention_pack(registry) + } else { + register_nvvm_numeric_pack(registry)?; + crate::nvvm::attention::register_nvvm_async_control_pack(registry) + } +} + +/// Register arithmetic and synchronization mappings shared by both architecture profiles. +fn register_nvvm_numeric_common_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_operation::(PackTranslation)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(ArithmeticTranslation)?; + /// Register one packed FP32 instruction with its exact identity and source verifier. + macro_rules! packed { + ($op:ty, $marker:literal, $name:literal, $args:literal) => { + registry.register_operation::<$op>(PackedArithmeticTranslation { + marker: $marker, + name: $name, + operands: $args, + verify: |ctx, source| <$op>::new(source).verify(ctx).map_err(|e| e.to_string()), + })?; + }; + } + packed!(AddF32x2Op, "v1:i0995", "nvvm.add.packed.f32x2", 2); + packed!(MulF32x2Op, "v1:i0999", "nvvm.mul.packed.f32x2", 2); + packed!(FmaF32x2Op, "v1:i1001", "nvvm.fma.packed.f32x2", 3); + + /// Register one direct intrinsic call with its operand count and shuffle contract. + macro_rules! direct { + ($op:ty, $marker:literal, $name:literal, $args:literal, $shuffle:literal) => { + registry.register_operation::<$op>(IntrinsicTranslation { + marker: $marker, + name: $name, + operands: $args, + shuffle: $shuffle, + verify: |ctx, source| <$op>::new(source).verify(ctx).map_err(|e| e.to_string()), + })?; + }; + } + registry.register_operation::(SyncTranslation(0))?; + registry.register_operation::(SyncTranslation(1))?; + registry.register_operation::(SyncTranslation(2))?; + registry.register_operation::(SyncTranslation(3))?; + registry.register_operation::(SyncTranslation(4))?; + direct!( + ShflSyncDownF32Op, + "v1:i0056", + "llvm.nvvm.shfl.sync.down.f32", + 3, + true + ); + direct!( + ShflSyncDownI32Op, + "v1:i0052", + "llvm.nvvm.shfl.sync.down.i32", + 3, + true + ); + direct!( + ShflSyncIdxF32Op, + "v1:i0054", + "llvm.nvvm.shfl.sync.idx.f32", + 3, + true + ); + direct!( + ShflSyncIdxI32Op, + "v1:i0050", + "llvm.nvvm.shfl.sync.idx.i32", + 3, + true + ); + direct!( + ShflSyncBflyI32Op, + "v1:i0051", + "llvm.nvvm.shfl.sync.bfly.i32", + 3, + true + ); + registry.register_operation::(RegisterTranslation(true))?; + registry.register_operation::(RegisterTranslation(false))?; + Ok(()) +} + +/// Build an LLVM intrinsic call retaining translated operands/results and an empty bundle list. +fn intrinsic(input: OperationInput, name: &str) -> Result { + let count = input.operands.len(); + let mut call = operation( + "llvm.call_intrinsic", + input.operands, + input.results, + &input.location, + )?; + call.properties + .insert("intrin".into(), MlirAttribute::String(name.into())); + call.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + call.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![count as i32, 0]), + ); + Ok(call) +} + +/// Validate one catalog identity and emit its named intrinsic with reviewed operands. +struct IntrinsicTranslation { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// LLVM intrinsic spelling selected for this reviewed source operation. + name: &'static str, + /// Expected number of runtime operands at the source boundary. + operands: usize, + /// Add the full-warp clamp required by the canonical butterfly-shuffle contract. + shuffle: bool, + /// Check the source operation's selectors and operand contract. + verify: VerifyOperation, +} +impl OperationTranslation for IntrinsicTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + require_input(&mut input, self.marker, self.operands, 1)?; + let mut ops = vec![]; + if self.shuffle { + // The public warp shuffle operates across the full 32-lane warp; + // the membership mask remains a distinct caller-supplied operand. + let (result, clamp) = fresh(session, MlirType::Integer(32)); + let mut constant = operation("arith.constant", vec![], vec![result], &input.location)?; + constant.properties.insert( + "value".into(), + MlirAttribute::Integer { + value: 31, + ty: MlirType::Integer(32), + }, + ); + ops.push(constant); + input.operands.push(clamp); + } + ops.push(intrinsic(input, self.name)?); + Ok(ops) + } +} + +/// Map reviewed scalar math to exact native intrinsics with preserved precision and subnormal mode. +pub(crate) struct MathTranslation; +impl OperationTranslation for MathTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let op = ScalarMathOp::new(source); + op.verify(ctx).map_err(|e| e.to_string())?; + if op.get_attr_nvvm_scalar_math_format(ctx).as_deref() != Some(&ScalarMathFormatAttr::F32) + || op.get_attr_nvvm_scalar_math_precision(ctx).as_deref() + != Some(&ScalarMathPrecisionAttr::Approx) + { + return Err( + "numeric pack supports approximate FP32 exp2/log2/reciprocal/rsqrt/sine only" + .into(), + ); + } + let pair = ( + op.get_attr_nvvm_scalar_math_operation(ctx), + op.get_attr_nvvm_scalar_math_subnormal(ctx), + ); + let (marker, name) = match (pair.0.as_deref(), pair.1.as_deref()) { + (Some(ScalarMathOperationAttr::Ex2), Some(ScalarMathSubnormalAttr::Preserve)) => { + ("v1:i0819", "llvm.nvvm.ex2.approx.f32") + } + (Some(ScalarMathOperationAttr::Ex2), Some(ScalarMathSubnormalAttr::Ftz)) => { + ("v1:i0820", "llvm.nvvm.ex2.approx.ftz.f32") + } + (Some(ScalarMathOperationAttr::Lg2), Some(ScalarMathSubnormalAttr::Preserve)) => { + ("v1:i0786", "llvm.nvvm.lg2.approx.f") + } + (Some(ScalarMathOperationAttr::Lg2), Some(ScalarMathSubnormalAttr::Ftz)) => { + ("v1:i0787", "llvm.nvvm.lg2.approx.ftz.f") + } + (Some(ScalarMathOperationAttr::Rcp), Some(ScalarMathSubnormalAttr::Ftz)) => { + ("v1:i0788", "llvm.nvvm.rcp.approx.ftz.f") + } + (Some(ScalarMathOperationAttr::Sin), Some(ScalarMathSubnormalAttr::Preserve)) => { + ("v1:i0782", "llvm.nvvm.sin.approx.f") + } + (Some(ScalarMathOperationAttr::Rsqrt), Some(ScalarMathSubnormalAttr::Preserve)) => { + ("v1:i0802", "llvm.nvvm.rsqrt.approx.f") + } + _ => { + return Err( + "numeric pack does not support this scalar math selector combination".into(), + ); + } + }; + require_input(&mut input, marker, 1, 1)?; + Ok(vec![intrinsic(input, name)?]) + } +} + +/// Preserve scalar arithmetic format, rounding, subnormal, and saturation selectors. +struct ArithmeticTranslation; +impl OperationTranslation for ArithmeticTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let op = ScalarArithmeticOp::new(source); + op.verify(ctx).map_err(|e| e.to_string())?; + if op.get_attr_nvvm_scalar_arithmetic_format(ctx).as_deref() + != Some(&ScalarArithmeticFormatAttr::F32) + || op.get_attr_nvvm_scalar_arithmetic_operation(ctx).as_deref() + != Some(&ScalarArithmeticOperationAttr::Fma) + || op.get_attr_nvvm_scalar_arithmetic_rounding(ctx).as_deref() + != Some(&ScalarArithmeticRoundingAttr::Rn) + || op + .get_attr_nvvm_scalar_arithmetic_saturation(ctx) + .as_deref() + != Some(&ScalarArithmeticSaturationAttr::None) + { + return Err("numeric pack supports unsaturated FP32 round-nearest FMA only".into()); + } + let (marker, name) = match op.get_attr_nvvm_scalar_arithmetic_subnormal(ctx).as_deref() { + Some(ScalarArithmeticSubnormalAttr::Preserve) => ("v1:i0418", "llvm.nvvm.fma.rn.f"), + Some(ScalarArithmeticSubnormalAttr::Ftz) => ("v1:i0419", "llvm.nvvm.fma.rn.ftz.f"), + _ => return Err("FP32 FMA requires a subnormal policy".into()), + }; + require_input(&mut input, marker, 3, 1)?; + Ok(vec![intrinsic(input, name)?]) + } +} + +/// The Rust carrier holds the low FP32 lane in bits 0..31. Bitcasting to an +/// LLVM vector preserves both lanes and lets NVVM select packed instructions; +/// scalar arithmetic would lose the source instruction-family contract. +struct PackedArithmeticTranslation { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// LLVM intrinsic spelling selected for this reviewed source operation. + name: &'static str, + /// Expected number of runtime operands at the source boundary. + operands: usize, + /// Check the source operation's selectors and operand contract. + verify: VerifyOperation, +} +impl OperationTranslation for PackedArithmeticTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + require_input(&mut input, self.marker, self.operands, 1)?; + let pair_type = MlirType::Vector { + shape: vec![2], + element: Box::new(MlirType::Float(MlirFloatType::F32)), + }; + let mut ops = Vec::with_capacity(self.operands + 2); + let mut pairs = Vec::with_capacity(self.operands); + for operand in input.operands { + let (result, pair) = fresh(session, pair_type.clone()); + ops.push(operation( + "llvm.bitcast", + vec![operand], + vec![result], + &input.location, + )?); + pairs.push(pair); + } + let (result, pair) = fresh(session, pair_type); + let mut packed = operation(self.name, pairs, vec![result], &input.location)?; + packed.properties.insert( + "rnd".into(), + MlirAttribute::dialect("#nvvm.fp_rnd_mode")?, + ); + packed + .properties + .insert("ftz".into(), MlirAttribute::Bool(false)); + ops.push(packed); + ops.push(operation( + "llvm.bitcast", + vec![pair], + input.results, + &input.location, + )?); + Ok(ops) + } +} + +/// Emit the selected register-budget increase or decrease instruction. +struct RegisterTranslation(#[doc = "True increases the register budget; false decreases it."] bool); +impl OperationTranslation for RegisterTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let count = if self.0 { + let op = SetMaxNRegIncOp::new(source); + op.verify(ctx).map_err(|e| e.to_string())?; + op.register_count(ctx) + } else { + let op = SetMaxNRegDecOp::new(source); + op.verify(ctx).map_err(|e| e.to_string())?; + op.register_count(ctx) + } + .ok_or("missing register count")?; + input.attributes.remove("register_count"); + require_input( + &mut input, + if self.0 { "v1:i0915" } else { "v1:i0916" }, + 0, + 0, + )?; + let mut op = operation("nvvm.setmaxregister", vec![], vec![], &input.location)?; + op.properties.insert( + "regCount".into(), + MlirAttribute::Integer { + value: count as i128, + ty: MlirType::Integer(32), + }, + ); + op.properties.insert( + "action".into(), + MlirAttribute::dialect(if self.0 { + "#nvvm" + } else { + "#nvvm" + })?, + ); + Ok(vec![op]) + } +} +/// Convert low-first FP32 operands into the exact packed FP16 register representation. +struct PackTranslation; +impl OperationTranslation for PackTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CvtF16x2F32Op::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0081", 2, 1)?; + let (result, packed) = fresh( + session, + MlirType::Vector { + shape: vec![2], + element: Box::new(MlirType::Float(MlirFloatType::F16)), + }, + ); + // Canonical Rust API is (low, high); LLVM's intrinsic is (high, low). + input.operands.reverse(); + let mut call = operation( + "llvm.call_intrinsic", + input.operands, + vec![result], + &input.location, + )?; + call.properties.insert( + "intrin".into(), + MlirAttribute::String("llvm.nvvm.ff2f16x2.rn".into()), + ); + call.properties.insert( + "op_bundle_sizes".into(), + MlirAttribute::DenseI32Array(vec![]), + ); + call.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![2, 0]), + ); + Ok(vec![ + call, + operation("llvm.bitcast", vec![packed], input.results, &input.location)?, + ]) + } +} + +/// Native MLIR operations retain the compiler's PTX lowering for these instructions. +struct SyncTranslation( + #[doc = "Reviewed selector: 0 reduction, 1 copy, 2 barrier arrival, 3 vote, or 4 warp sync."] + u8, +); +impl OperationTranslation for SyncTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let (marker, args, results) = match self.0 { + 0 => { + ReduxSyncFmaxOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ("v1:i1007", 2, 1) + } + 1 => { + CpAsyncCa4Op::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ("v1:i0086", 2, 0) + } + 2 => { + CpAsyncMbarrierArriveNoIncOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ("v1:i0103", 1, 0) + } + 3 => { + VoteSyncAllOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ("v1:i0040", 2, 1) + } + 4 => { + BarWarpSyncOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + ("v1:i0049", 1, 0) + } + _ => unreachable!(), + }; + require_input(&mut input, marker, args, results)?; + let mut e = crate::cute::Emit { + session, + location: input.location, + ops: vec![], + }; + let (name, operands) = match self.0 { + 0 => ( + "nvvm.redux.sync", + vec![input.operands[1].clone(), input.operands[0].clone()], + ), + 1 => { + let dst = e.llvm_ptr(input.operands[0].clone(), 3)?; + let src = e.llvm_ptr(input.operands[1].clone(), 1)?; + ("nvvm.cp.async.shared.global", vec![dst, src]) + } + 2 => { + let bar = e.llvm_ptr(input.operands[0].clone(), 3)?; + ("nvvm.cp.async.mbarrier.arrive", vec![bar]) + } + 3 => ("nvvm.vote.sync", input.operands), + 4 => ("nvvm.bar.warp.sync", input.operands), + _ => unreachable!(), + }; + let mut op = e.operation(name, operands)?; + op.results = input.results; + match self.0 { + 0 => { + crate::cute::property(&mut op, "kind", "#nvvm")?; + } + 1 => { + crate::cute::integer_property(&mut op, "size", 32, 4); + crate::cute::property(&mut op, "modifier", "#nvvm")?; + } + 2 => { + op.properties + .insert("noinc".into(), MlirAttribute::Bool(true)); + } + 3 => { + crate::cute::property(&mut op, "kind", "#nvvm")?; + } + _ => {} + } + e.ops.push(op); + Ok(e.ops) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs new file mode 100644 index 0000000000..553918b703 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs @@ -0,0 +1,586 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Register-MMA and shared-memory primitives for SM120, using standard NVVM. +//! Source intrinsic identities and structural selectors are checked before any +//! source-only attributes are removed. No CuTe operations or types are involved. + +use crate::{ + DropAttribute, MlirAttribute, MlirFloatType, MlirLocation, MlirOperation, MlirResult, MlirType, + MlirValueUse, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; +use dialect_mir::ops::MirConstantOp; +use dialect_nvvm::ops::*; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, +}; + +/// The exact register-MMA and matrix-load variants used by the supported kernels. +pub fn register_nvvm_sm120_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(LoadTranslation)?; + registry.register_operation::(StoreTranslation)?; + registry.register_operation::(TmaTranslation)?; + registry.register_operation::(TmaCtaTranslation)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(MmaTranslation)?; + Ok(()) +} + +/// Reorder reviewed register MMA operands and preserve packed lanes and scale selectors. +struct MmaTranslation; +impl OperationTranslation for MmaTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mma = RegisterMmaOp::new(source); + mma.verify(ctx).map_err(|e| e.to_string())?; + // Both formats use packed u32 source carriers, but NVVM deliberately + // consumes BF16 pairs as i32 and FP16 pairs as vector<2xf16>. + let a_element = mma.get_attr_nvvm_register_mma_a_element(ctx); + let b_element = mma.get_attr_nvvm_register_mma_b_element(ctx); + let is_bf16 = a_element.as_deref() == Some(&RegisterMmaElementAttr::Bf16); + let matching_half_operands = a_element.as_deref() == b_element.as_deref() + && matches!( + a_element.as_deref(), + Some(RegisterMmaElementAttr::F16 | RegisterMmaElementAttr::Bf16) + ); + if mma.get_attr_nvvm_register_mma_shape(ctx).as_deref() + == Some(&RegisterMmaShapeAttr::M16n8k16) + && mma.get_attr_nvvm_register_mma_kind(ctx).as_deref() + == Some(&RegisterMmaKindAttr::Standard) + && mma.get_attr_nvvm_register_mma_operation(ctx).as_deref() + == Some(&RegisterMmaOperationAttr::Multiply) + && mma.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() + == Some(&RegisterMmaAccumulatorAttr::F32) + && matching_half_operands + && mma.get_attr_nvvm_register_mma_a_layout(ctx).as_deref() + == Some(&RegisterMmaLayoutAttr::Row) + && mma.get_attr_nvvm_register_mma_b_layout(ctx).as_deref() + == Some(&RegisterMmaLayoutAttr::Col) + && mma.get_attr_nvvm_register_mma_overflow(ctx).as_deref() + == Some(&RegisterMmaOverflowAttr::NotApplicable) + { + let marker = if is_bf16 { "v1:i0105" } else { "v1:i0106" }; + require_input(&mut input, marker, 10, 4)?; + let mut operations = vec![]; + let mut operands = vec![]; + // Source operands are C[4], packed A[4], packed B[2]. NVVM requires + // A[4], B[2], C[4], preserving each packed pair's original bits. + for packed in &input.operands[4..] { + if is_bf16 { + operands.push(packed.clone()); + continue; + } + let (result, value) = fresh( + session, + MlirType::Vector { + shape: vec![2], + element: Box::new(MlirType::Float(MlirFloatType::F16)), + }, + ); + operations.push(operation( + "llvm.bitcast", + vec![packed.clone()], + vec![result], + &input.location, + )?); + operands.push(value); + } + operands.extend_from_slice(&input.operands[..4]); + let (result, aggregate) = fresh( + session, + MlirType::dialect("!llvm.struct<(f32, f32, f32, f32)>")?, + ); + let mut op = operation("nvvm.mma.sync", operands, vec![result], &input.location)?; + property(&mut op, "shape", "#nvvm.shape")?; + property(&mut op, "layoutA", "#nvvm.mma_layout")?; + property(&mut op, "layoutB", "#nvvm.mma_layout")?; + let element = if is_bf16 { + "#nvvm.mma_type" + } else { + "#nvvm.mma_type" + }; + property(&mut op, "multiplicandAPtxType", element)?; + property(&mut op, "multiplicandBPtxType", element)?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![4, 2, 4]), + ); + operations.push(op); + extract_results(&mut operations, aggregate, input.results, &input.location)?; + return Ok(operations); + } + if mma.get_attr_nvvm_register_mma_shape(ctx).as_deref() + != Some(&RegisterMmaShapeAttr::M16n8k64) + || mma.get_attr_nvvm_register_mma_kind(ctx).as_deref() + != Some(&RegisterMmaKindAttr::Mxf4) + || mma.get_attr_nvvm_register_mma_operation(ctx).as_deref() + != Some(&RegisterMmaOperationAttr::Multiply) + || mma.get_attr_nvvm_register_mma_accumulator(ctx).as_deref() + != Some(&RegisterMmaAccumulatorAttr::F32) + || mma.get_attr_nvvm_register_mma_a_element(ctx).as_deref() + != Some(&RegisterMmaElementAttr::E2m1) + || mma.get_attr_nvvm_register_mma_b_element(ctx).as_deref() + != Some(&RegisterMmaElementAttr::E2m1) + || mma.get_attr_nvvm_register_mma_a_layout(ctx).as_deref() + != Some(&RegisterMmaLayoutAttr::Row) + || mma.get_attr_nvvm_register_mma_b_layout(ctx).as_deref() + != Some(&RegisterMmaLayoutAttr::Col) + || mma.get_attr_nvvm_register_mma_overflow(ctx).as_deref() + != Some(&RegisterMmaOverflowAttr::NotApplicable) + { + return Err("SM120 register MMA mapping supports m16n8k64 MXF4 row/col E2M1 with F32 accumulation and UE8M0 scale_vec::2X, or m16n8k16 standard row/col matching F16 or BF16 with F32 accumulation only".into()); + } + require_input(&mut input, "v1:i1039", 16, 4)?; + // Canonical source order is C, A, B, scales; NVVM uses A, B, C, scales. + let operands = input.operands[4..10] + .iter() + .chain(&input.operands[..4]) + .chain(&input.operands[10..]) + .cloned() + .collect(); + let (result, aggregate) = fresh( + session, + MlirType::dialect("!llvm.struct<(f32, f32, f32, f32)>")?, + ); + let mut op = operation( + "nvvm.mma.block_scale", + operands, + vec![result], + &input.location, + )?; + property(&mut op, "shape", "#nvvm.shape")?; + property(&mut op, "multiplicandAPtxType", "#nvvm.mma_type")?; + property(&mut op, "multiplicandBPtxType", "#nvvm.mma_type")?; + property(&mut op, "scaleVecSize", "#nvvm.scale_vec_size")?; + property( + &mut op, + "blockScaleFormat", + "#nvvm.block_scale_format", + )?; + property(&mut op, "kind", "#nvvm.block_scale_kind")?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![4, 2, 4, 1, 1, 1, 1, 1, 1]), + ); + let mut operations = vec![op]; + extract_results(&mut operations, aggregate, input.results, &input.location)?; + Ok(operations) + } +} + +/// Emit the exact matrix-load transpose, shape, packing, and result-register contract. +struct LoadTranslation; +impl OperationTranslation for LoadTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let load = LdmatrixOp::new(source); + load.verify(ctx).map_err(|e| e.to_string())?; + if load.get_attr_nvvm_ldmatrix_shape(ctx).as_deref() != Some(&LdmatrixShapeAttr::M8n8) + || load.get_attr_nvvm_ldmatrix_element(ctx).as_deref() + != Some(&LdmatrixElementAttr::B16) + || load.get_attr_nvvm_ldmatrix_state_space(ctx).as_deref() + != Some(&LdmatrixStateSpaceAttr::Shared) + { + return Err( + "SM120 matrix load mapping supports m8n8.x2/x4.b16 normal or transposed shared loads only" + .into(), + ); + } + let (count, marker, layout) = match ( + load.get_attr_nvvm_ldmatrix_multiplicity(ctx).as_deref(), + load.get_attr_nvvm_ldmatrix_layout(ctx).as_deref(), + ) { + (Some(LdmatrixMultiplicityAttr::X2), Some(LdmatrixLayoutAttr::Normal)) => (2, "v1:i0027", "#nvvm.mma_layout"), + (Some(LdmatrixMultiplicityAttr::X2), Some(LdmatrixLayoutAttr::Transposed)) => (2, "v1:i0028", "#nvvm.mma_layout"), + (Some(LdmatrixMultiplicityAttr::X4), Some(LdmatrixLayoutAttr::Normal)) => (4, "v1:i0013", "#nvvm.mma_layout"), + (Some(LdmatrixMultiplicityAttr::X4), Some(LdmatrixLayoutAttr::Transposed)) => (4, "v1:i0029", "#nvvm.mma_layout"), + _ => return Err("SM120 matrix load mapping supports m8n8.x2/x4.b16 normal or transposed shared loads only".into()), + }; + require_input(&mut input, marker, 1, count)?; + let mut operations = vec![]; + let pointer = pointer( + session, + &mut operations, + input.operands[0].clone(), + 3, + &input.location, + )?; + let (result, aggregate) = fresh( + session, + MlirType::dialect(if count == 2 { + "!llvm.struct<(i32, i32)>" + } else { + "!llvm.struct<(i32, i32, i32, i32)>" + })?, + ); + let mut op = operation( + "nvvm.ldmatrix", + vec![pointer], + vec![result], + &input.location, + )?; + property(&mut op, "layout", layout)?; + op.properties.insert( + "num".into(), + MlirAttribute::Integer { + value: count as i128, + ty: MlirType::Integer(32), + }, + ); + operations.push(op); + extract_results(&mut operations, aggregate, input.results, &input.location)?; + Ok(operations) + } +} + +/// Emit the verified matrix store without reordering its scalar register operands. +struct StoreTranslation; +impl OperationTranslation for StoreTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + StmatrixM8n8X2Op::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0301", 3, 0)?; + let mut operations = vec![]; + input.operands[0] = pointer( + session, + &mut operations, + input.operands[0].clone(), + 3, + &input.location, + )?; + let mut op = operation("nvvm.stmatrix", input.operands, vec![], &input.location)?; + property(&mut op, "layout", "#nvvm.mma_layout")?; + operations.push(op); + Ok(operations) + } +} + +/// Emit local shared-barrier TMA loads while rejecting hidden multicast or cache-policy changes. +struct TmaTranslation; +impl OperationTranslation for TmaTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CpAsyncBulkTensorG2sTile2dOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i0329", 7, 0)?; + if input.operands[..3] + .iter() + .any(|operand| matches!(operand.ty, MlirType::Integer(_))) + { + return Err( + "TMA requires pointer carriers for destination, barrier, and descriptor".into(), + ); + } + if constant_integer(ctx, source, 5)? != 0 || constant_integer(ctx, source, 6)? != 0 { + return Err( + "ordinary 2D TMA requires the importer's zero multicast mask and cache hint".into(), + ); + } + for index in [3, 4] { + if input.operands[index].ty != MlirType::Integer(32) { + return Err("2D TMA coordinates must be i32".into()); + } + } + if input.operands[5].ty != MlirType::Integer(16) + || input.operands[6].ty != MlirType::Integer(64) + { + return Err("2D TMA mask/cache carriers must be i16/i64".into()); + } + let mut operations = vec![]; + // The canonical LLVM g2s intrinsic encodes the destination in p7, + // including for a non-multicast copy into the executing CTA's SMEM. + let dst = pointer( + session, + &mut operations, + input.operands[0].clone(), + 7, + &input.location, + )?; + let barrier = pointer( + session, + &mut operations, + input.operands[1].clone(), + 3, + &input.location, + )?; + let descriptor = pointer( + session, + &mut operations, + input.operands[2].clone(), + 0, + &input.location, + )?; + let mut op = operation( + "nvvm.cp.async.bulk.tensor.shared.cluster.global", + vec![ + dst, + descriptor, + input.operands[3].clone(), + input.operands[4].clone(), + barrier, + ], + vec![], + &input.location, + )?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 1, 0, 0, 0, 0]), + ); + op.properties + .insert("useIntrinsic".into(), MlirAttribute::Bool(false)); + operations.push(op); + Ok(operations) + } +} + +/// Lower a CTA-local 2D TMA load (`cp.async.bulk.tensor.2d.shared::cta.global.tile`). +struct TmaCtaTranslation; +impl OperationTranslation for TmaCtaTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + CpAsyncBulkTensorG2sCtaTile2dOp::new(source) + .verify(ctx) + .map_err(|e| e.to_string())?; + require_input(&mut input, "v1:i1035", 6, 0)?; + if input.operands[..3] + .iter() + .any(|operand| matches!(operand.ty, MlirType::Integer(_))) + { + return Err( + "TMA requires pointer carriers for destination, barrier, and descriptor".into(), + ); + } + if constant_integer(ctx, source, 5)? != 0 { + return Err("ordinary 2D CTA TMA requires the importer's zero cache hint".into()); + } + for index in [3, 4] { + if input.operands[index].ty != MlirType::Integer(32) { + return Err("2D TMA coordinates must be i32".into()); + } + } + if input.operands[5].ty != MlirType::Integer(64) { + return Err("2D CTA TMA cache carrier must be i64".into()); + } + let mut operations = vec![]; + let dst = pointer( + session, + &mut operations, + input.operands[0].clone(), + 7, + &input.location, + )?; + let barrier = pointer( + session, + &mut operations, + input.operands[1].clone(), + 3, + &input.location, + )?; + let descriptor = pointer( + session, + &mut operations, + input.operands[2].clone(), + 0, + &input.location, + )?; + let mut op = operation( + "nvvm.cp.async.bulk.tensor.shared.cluster.global", + vec![ + dst, + descriptor, + input.operands[3].clone(), + input.operands[4].clone(), + barrier, + ], + vec![], + &input.location, + )?; + op.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 2, 1, 0, 0, 0, 0]), + ); + op.properties + .insert("useIntrinsic".into(), MlirAttribute::Bool(false)); + operations.push(op); + Ok(operations) + } +} + +/// Check operand/result counts and canonical identity, then reject leftover source structure. +pub(crate) fn require_input( + input: &mut OperationInput, + marker: &str, + operands: usize, + results: usize, +) -> Result<(), String> { + if input.attributes.remove("cuda_oxide_intrinsic_marker") + != Some(MlirAttribute::String(marker.into())) + { + return Err(format!( + "SM120 operation expected intrinsic marker {marker}" + )); + } + if input.operands.len() != operands + || input.results.len() != results + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err(format!( + "SM120 operation expected {operands} operands, {results} results, and no residual attributes or control flow" + )); + } + Ok(()) +} +/// Require an unsigned compile-time integer selector at the specified source operand. +pub(crate) fn constant_integer( + ctx: &Context, + source: Ptr, + index: usize, +) -> Result { + let definition = source + .deref(ctx) + .get_operand(index) + .defining_op() + .ok_or("expected a constant integer operand")?; + if Operation::get_opid(definition, ctx) != MirConstantOp::get_opid_static() { + return Err("expected a constant integer operand".into()); + } + MirConstantOp::new(definition) + .get_attr_value(ctx) + .ok_or("constant has no value")? + .value() + .to_string_decimal(false) + .parse() + .map_err(|_| "constant integer exceeds u64".into()) +} +/// Preserve the required NVVM pointer space, including explicit integer-offset conversions. +fn pointer( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + value: MlirValueUse, + space: u32, + location: &MlirLocation, +) -> Result { + let target = MlirType::dialect(if space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{space}>") + })?; + if value.ty == target { + return Ok(value); + } + let shared_source = value.ty == MlirType::dialect("!llvm.ptr")? + || value.ty == MlirType::dialect("!llvm.ptr<3>")?; + let descriptor_source = value.ty == MlirType::dialect("!llvm.ptr<1>")? + || value.ty == MlirType::dialect("!llvm.ptr<4>")?; + let cast = if space == 3 && value.ty == MlirType::Integer(32) { + "llvm.inttoptr" + } else if (space != 0 && shared_source) || (space == 0 && descriptor_source) { + "llvm.addrspacecast" + } else { + return Err(format!( + "unsupported pointer {:?} for SM120 address space {space}", + value.ty + )); + }; + let (result, converted) = fresh(session, target); + operations.push(operation(cast, vec![value], vec![result], location)?); + Ok(converted) +} +/// Attach a parsed dialect attribute as a target operation property. +pub(crate) fn property(op: &mut MlirOperation, name: &str, value: &str) -> Result<(), String> { + op.properties + .insert(name.into(), MlirAttribute::dialect(value)?); + Ok(()) +} +/// Unpack a native aggregate result into source SSA results in their original register order. +fn extract_results( + operations: &mut Vec, + aggregate: MlirValueUse, + results: Vec, + location: &MlirLocation, +) -> Result<(), String> { + for (index, result) in results.into_iter().enumerate() { + let mut extract = operation( + "llvm.extractvalue", + vec![aggregate.clone()], + vec![result], + location, + )?; + extract.properties.insert( + "position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + operations.push(extract); + } + Ok(()) +} +/// Allocate one SSA identifier and return matching result and operand handles. +pub(crate) fn fresh( + session: &mut TranslationSession<'_>, + ty: MlirType, +) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} +/// Build a target operation with the supplied operands, results, and source location. +pub(crate) fn operation( + name: &str, + operands: Vec, + results: Vec, + location: &MlirLocation, +) -> Result { + let mut operation = MlirOperation::new(name)?; + operation.operands = operands; + operation.results = results; + operation.location = location.clone(); + Ok(operation) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sregs.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sregs.rs new file mode 100644 index 0000000000..a65e974cef --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sregs.rs @@ -0,0 +1,269 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CUDA thread and grid coordinates in MLIR's NVVM dialect. +//! +//! These operations already have a one-to-one target operation. The pack is +//! intentionally narrow: another NVVM operation needs its own reviewed recipe +//! instead of being renamed by a catch-all rule. + +use crate::{ + MlirOperation, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; +use dialect_nvvm::ops::{ + ReadPtxSregCtaidXOp, ReadPtxSregCtaidYOp, ReadPtxSregCtaidZOp, ReadPtxSregNctaidXOp, + ReadPtxSregNctaidYOp, ReadPtxSregNctaidZOp, ReadPtxSregNtidXOp, ReadPtxSregNtidYOp, + ReadPtxSregNtidZOp, ReadPtxSregTidXOp, ReadPtxSregTidYOp, ReadPtxSregTidZOp, +}; +use pliron::{ + context::{Context, Ptr}, + operation::Operation, +}; + +/// Register the zero-operand thread, block, and grid dimension reads. +pub fn register_nvvm_sreg_pack(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.tid.x", + marker: "v1:i0001", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.tid.y", + marker: "v1:i0005", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.tid.z", + marker: "v1:i0009", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ntid.x", + marker: "v1:i0003", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ntid.y", + marker: "v1:i0007", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ntid.z", + marker: "v1:i0011", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ctaid.x", + marker: "v1:i0002", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ctaid.y", + marker: "v1:i0006", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.ctaid.z", + marker: "v1:i0010", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.nctaid.x", + marker: "v1:i0004", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.nctaid.y", + marker: "v1:i0008", + })?; + registry.register_operation::(SregTranslation { + target: "nvvm.read.ptx.sreg.nctaid.z", + marker: "v1:i0012", + })?; + Ok(()) +} + +/// Translate one reviewed zero-operand coordinate read without widening the registered subset. +struct SregTranslation { + /// NVVM special-register operation spelling emitted by this recipe. + target: &'static str, + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, +} + +impl OperationTranslation for SregTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let marker = input + .attributes + .remove("cuda_oxide_intrinsic_marker") + .ok_or_else(|| { + format!( + "{} has no cuda_oxide_intrinsic_marker; expected {}", + self.target, self.marker + ) + })?; + if marker != crate::MlirAttribute::String(self.marker.into()) { + return Err(format!( + "{} expected intrinsic marker {}, got {marker:?}", + self.target, self.marker + )); + } + if !input.operands.is_empty() + || input.results.len() != 1 + || !input.successors.is_empty() + || !input.regions.is_empty() + || !input.attributes.is_empty() + { + return Err(format!( + "{} expected no operands, one result, no successors, no regions, and no attributes", + self.target + )); + } + + let mut target = MlirOperation::new(self.target)?; + target.results = input.results; + target.location = input.location; + Ok(vec![target]) + } +} + +#[cfg(test)] +mod tests { + //! One-to-one source/target checks for the reviewed coordinate-register subset. + + use super::*; + use crate::{ + CutlassFullCuteMlir22, MlirConsumerProfile, + profile::render_mapping_module_without_cutlass_envelope, + }; + use dialect_mir::ops::{MirFuncOp, MirReturnOp}; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + identifier::Identifier, + op::Op, + r#type::TypedHandle, + }; + + /// Append a chosen coordinate-read operation and return its source result. + fn append_sreg( + ctx: &mut Context, + block: Ptr, + result_type: TypedHandle, + marker: &str, + ) -> Ptr { + let operation = Operation::new( + ctx, + O::get_concrete_op_info(), + vec![result_type.into()], + vec![], + vec![], + 0, + ); + operation.deref_mut(ctx).attributes.set( + Identifier::try_from("cuda_oxide_intrinsic_marker").unwrap(), + StringAttr::new(marker.into()), + ); + operation.insert_at_back(block, ctx); + operation + } + + /// Check thread, block, and grid dimension reads retain their coordinate axis and target spelling. + #[test] + fn thread_block_and_grid_dimensions_map_one_to_one() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, Identifier::try_from("sregs").unwrap()); + let function_type = FunctionType::get(&ctx, vec![], vec![]); + let function_operation = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new( + &mut ctx, + function_operation, + TypeAttr::new(function_type.into()), + ); + function.set_symbol_name(&mut ctx, Identifier::try_from("coordinates").unwrap()); + module.append_operation(&mut ctx, function_operation, 0); + + let region = function_operation.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + let i32_type = IntegerType::get(&ctx, 32, Signedness::Unsigned); + + let tid_x = append_sreg::(&mut ctx, entry, i32_type, "v1:i0001"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0005"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0009"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0003"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0007"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0011"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0002"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0006"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0010"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0004"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0008"); + append_sreg::(&mut ctx, entry, i32_type, "v1:i0012"); + + let return_operation = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + return_operation.insert_at_back(entry, &ctx); + + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_mapping_module_without_cutlass_envelope(&target, "sregs"); + let expected = r#""builtin.module"() <{sym_name = "sregs"}> ({ + ^bb0: + "func.func"() <{function_type = () -> (), sym_name = "coordinates"}> ({ + ^bb1: + %v0 = "nvvm.read.ptx.sreg.tid.x"() : () -> (i32) + %v1 = "nvvm.read.ptx.sreg.tid.y"() : () -> (i32) + %v2 = "nvvm.read.ptx.sreg.tid.z"() : () -> (i32) + %v3 = "nvvm.read.ptx.sreg.ntid.x"() : () -> (i32) + %v4 = "nvvm.read.ptx.sreg.ntid.y"() : () -> (i32) + %v5 = "nvvm.read.ptx.sreg.ntid.z"() : () -> (i32) + %v6 = "nvvm.read.ptx.sreg.ctaid.x"() : () -> (i32) + %v7 = "nvvm.read.ptx.sreg.ctaid.y"() : () -> (i32) + %v8 = "nvvm.read.ptx.sreg.ctaid.z"() : () -> (i32) + %v9 = "nvvm.read.ptx.sreg.nctaid.x"() : () -> (i32) + %v10 = "nvvm.read.ptx.sreg.nctaid.y"() : () -> (i32) + %v11 = "nvvm.read.ptx.sreg.nctaid.z"() : () -> (i32) + "func.return"() : () -> () + }) {llvm.linkage = #llvm.linkage} : () -> () +}) : () -> () +"#; + assert_eq!( + crate::profile::normalize_mlir_identifiers(&text), + crate::profile::normalize_mlir_identifiers(expected) + ); + + tid_x.deref_mut(&ctx).attributes.set( + Identifier::try_from("cuda_oxide_intrinsic_marker").unwrap(), + StringAttr::new("v1:i0002".into()), + ); + let error = profile + .translate_module(&ctx, &module) + .unwrap_err() + .to_string(); + assert!( + error.contains("expected intrinsic marker v1:i0001"), + "{error}" + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sync_restrict.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sync_restrict.rs new file mode 100644 index 0000000000..8a6f435fe3 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sync_restrict.rs @@ -0,0 +1,107 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Native proxy fences retaining the catalog's restricted shared-memory ordering. + +use crate::{ + MlirAttribute, MlirOperation, OperationInput, OperationTranslation, TranslationError, + TranslationRegistry, TranslationSession, +}; +use dialect_nvvm::ops::{ + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceProxyAsyncGenericReleaseSharedCtaClusterOp, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Register cluster-scoped acquire/release ordering between generic and async proxies. +pub fn register_nvvm_sync_restrict_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + /// Register a proxy-fence operation with its canonical identity and memory order. + macro_rules! register { + ($op:ty, $marker:literal, $order:literal) => { + registry.register_operation::<$op>(ProxyFenceTranslation { + marker: $marker, + order: $order, + verify: |ctx, source| { + <$op>::new(source) + .verify(ctx) + .map_err(|error| error.to_string()) + }, + })?; + }; + } + register!( + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + "v1:i0315", + "acquire" + ); + register!( + FenceProxyAsyncGenericReleaseSharedCtaClusterOp, + "v1:i0314", + "release" + ); + Ok(()) +} + +/// Keep proxy direction and restricted shared-memory ordering on a native fence. +struct ProxyFenceTranslation { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// Acquire or release ordering spelling retained by the proxy fence. + order: &'static str, + /// Check the source operation's selectors and operand contract. + verify: fn(&Context, Ptr) -> Result<(), String>, +} + +impl OperationTranslation for ProxyFenceTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + if input.attributes.remove("cuda_oxide_intrinsic_marker") + != Some(MlirAttribute::String(self.marker.into())) + { + return Err(format!( + "sync-restrict proxy fence expected intrinsic marker {}", + self.marker + )); + } + if !input.operands.is_empty() + || !input.results.is_empty() + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err( + "sync-restrict proxy fence requires no operands, results, residual attributes, or control flow" + .into(), + ); + } + let mut op = MlirOperation::new("nvvm.fence.proxy.sync_restrict")?; + op.location = input.location; + op.properties.insert( + "order".into(), + MlirAttribute::dialect(format!("#nvvm.mem_order<{}>", self.order))?, + ); + // The native operation uses the same proxy properties for both orders; + // acquire restricts shared::cluster and release restricts shared::cta. + op.properties.insert( + "fromProxy".into(), + MlirAttribute::dialect("#nvvm.proxy_kind")?, + ); + op.properties.insert( + "toProxy".into(), + MlirAttribute::dialect("#nvvm.proxy_kind")?, + ); + Ok(vec![op]) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs new file mode 100644 index 0000000000..da916628ec --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/tcgen05.rs @@ -0,0 +1,1114 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Direct TCGen05 primitive mappings for the pinned CUTLASS MLIR consumer. +//! +//! CUDA Oxide represents a TMEM address as a packed `u32` token and a load as +//! scalar register results. MLIR uses an address-space-6 pointer and a vector. +//! Those carrier conversions must not rescale the packed TMEM row/column bits. +//! This pack does not lower through PTX or expand high-level CuTe operations. + +use dialect_nvvm::ops::*; + +use crate::{ + DropAttribute, MlirAttribute, MlirLocation, MlirOperation, MlirResult, MlirType, MlirValueUse, + OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, +}; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + operation::Operation, +}; + +/// Register the two-CTA FP16/TMEM primitives used by the split-N GEMM. +/// +/// Other MMA forms deliberately remain unsupported until their operand and +/// selector contracts have a reviewed mapping. Target capability validation +/// remains the caller's responsibility, as for the other intrinsic packs. +pub fn register_nvvm_tcgen05_pack( + registry: &mut TranslationRegistry, +) -> Result<(), TranslationError> { + // Read the verified source selectors directly rather than converting + // source enum discriminants into unrelated target enum discriminants. + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_attribute::(DropAttribute)?; + registry.register_operation::(MmaTranslation)?; + registry.register_operation::(SharedOffsetTranslation(false))?; + registry + .register_operation::(SharedOffsetTranslation(true))?; + /// Register a tensor-memory lifecycle instruction with group, arity, identity, and verifier. + macro_rules! primitive { + ($op:ty, $marker:literal, $kind:expr, $group:literal, $operands:literal) => { + registry.register_operation::<$op>(PrimitiveTranslation { + marker: $marker, + kind: $kind, + group: $group, + operands: $operands, + verify: |ctx, source| { + <$op>::new(source) + .verify(ctx) + .map_err(|error| error.to_string()) + }, + })?; + }; + } + primitive!(Tcgen05AllocOp, "v1:i0343", Primitive::Alloc, 1, 2); + primitive!(Tcgen05AllocCg2Op, "v1:i0359", Primitive::Alloc, 2, 2); + primitive!(Tcgen05DeallocOp, "v1:i0344", Primitive::Dealloc, 1, 2); + primitive!(Tcgen05DeallocCg2Op, "v1:i0360", Primitive::Dealloc, 2, 2); + primitive!(Tcgen05CommitOp, "v1:i0348", Primitive::Commit, 1, 1); + primitive!(Tcgen05CommitCg2Op, "v1:i0363", Primitive::Commit, 2, 1); + primitive!( + Tcgen05CommitMulticastOp, + "v1:i0760", + Primitive::Commit, + 1, + 2 + ); + primitive!( + Tcgen05CommitMulticastCg2Op, + "v1:i0365", + Primitive::Commit, + 2, + 2 + ); + primitive!( + Tcgen05RelinquishAllocPermitOp, + "v1:i0345", + Primitive::Relinquish, + 1, + 0 + ); + primitive!( + Tcgen05RelinquishAllocPermitCg2Op, + "v1:i0361", + Primitive::Relinquish, + 2, + 0 + ); + primitive!(Tcgen05LoadWaitOp, "v1:i0357", Primitive::LoadWait, 1, 0); + primitive!(Tcgen05StoreWaitOp, "v1:i0358", Primitive::StoreWait, 1, 0); + primitive!( + Tcgen05FenceBeforeThreadSyncOp, + "v1:i0346", + Primitive::FenceBefore, + 1, + 0 + ); + primitive!( + Tcgen05FenceAfterThreadSyncOp, + "v1:i0347", + Primitive::FenceAfter, + 1, + 0 + ); + /// Register a tensor-memory transfer with exact scalar register count and packing mode. + macro_rules! memory { + ($op:ty, $marker:literal, $width:literal, $store:literal, $pack:literal) => { + registry.register_operation::<$op>(MemoryTranslation { + marker: $marker, + width: $width, + store: $store, + pack: $pack, + verify: |ctx, source| { + <$op>::new(source) + .verify(ctx) + .map_err(|error| error.to_string()) + }, + })?; + }; + } + memory!(Tcgen05Ld32x32bX1RawOp, "v1:i0654", 1, false, false); + memory!(Tcgen05Ld32x32bX1Pack16Op, "v1:i0655", 1, false, true); + memory!(Tcgen05St32x32bX1RawOp, "v1:i0712", 1, true, false); + memory!(Tcgen05St32x32bX1Unpack16Op, "v1:i0713", 1, true, true); + memory!(Tcgen05Ld32x32bX2RawOp, "v1:i0656", 2, false, false); + memory!(Tcgen05Ld32x32bX2Pack16Op, "v1:i0657", 2, false, true); + memory!(Tcgen05St32x32bX2RawOp, "v1:i0714", 2, true, false); + memory!(Tcgen05St32x32bX2Unpack16Op, "v1:i0715", 2, true, true); + memory!(Tcgen05Ld32x32bX4RawOp, "v1:i0658", 4, false, false); + memory!(Tcgen05Ld32x32bX4Pack16Op, "v1:i0659", 4, false, true); + memory!(Tcgen05St32x32bX4RawOp, "v1:i0716", 4, true, false); + memory!(Tcgen05St32x32bX4Unpack16Op, "v1:i0717", 4, true, true); + memory!(Tcgen05Ld32x32bX8RawOp, "v1:i0660", 8, false, false); + memory!(Tcgen05Ld32x32bX8Pack16Op, "v1:i0661", 8, false, true); + memory!(Tcgen05St32x32bX8RawOp, "v1:i0718", 8, true, false); + memory!(Tcgen05St32x32bX8Unpack16Op, "v1:i0719", 8, true, true); + memory!(Tcgen05Ld32x32bX16RawOp, "v1:i0662", 16, false, false); + memory!(Tcgen05Ld32x32bX16Pack16Op, "v1:i0663", 16, false, true); + memory!(Tcgen05St32x32bX16RawOp, "v1:i0720", 16, true, false); + memory!(Tcgen05St32x32bX16Unpack16Op, "v1:i0721", 16, true, true); + memory!(Tcgen05Ld32x32bX32RawOp, "v1:i0664", 32, false, false); + memory!(Tcgen05Ld32x32bX32Pack16Op, "v1:i0665", 32, false, true); + memory!(Tcgen05St32x32bX32RawOp, "v1:i0722", 32, true, false); + memory!(Tcgen05St32x32bX32Unpack16Op, "v1:i0723", 32, true, true); + memory!(Tcgen05Ld32x32bX64RawOp, "v1:i0666", 64, false, false); + memory!(Tcgen05Ld32x32bX64Pack16Op, "v1:i0667", 64, false, true); + memory!(Tcgen05St32x32bX64RawOp, "v1:i0724", 64, true, false); + memory!(Tcgen05St32x32bX64Unpack16Op, "v1:i0725", 64, true, true); + memory!(Tcgen05Ld32x32bX128RawOp, "v1:i0668", 128, false, false); + memory!(Tcgen05Ld32x32bX128Pack16Op, "v1:i0669", 128, false, true); + memory!(Tcgen05St32x32bX128RawOp, "v1:i0726", 128, true, false); + memory!(Tcgen05St32x32bX128Unpack16Op, "v1:i0727", 128, true, true); + Ok(()) +} + +/// Preserve reviewed TCGen05 MMA operands, CTA group, shape, and A-collector selectors. +struct MmaTranslation; + +/// Convert a pointer to an explicit CTA- or cluster-shared integer offset. +struct SharedOffsetTranslation( + #[doc = "True selects cluster-shared AS7 offsets; false selects CTA-shared AS3 offsets."] bool, +); + +impl OperationTranslation for SharedOffsetTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (if self.0 { + CvtaGenericToClusterSharedOffsetOp::new(source).verify(ctx) + } else { + CvtaGenericToSharedOffsetOp::new(source).verify(ctx) + }) + .map_err(|error| error.to_string())?; + if input.operands.len() != 1 + || input.results.len() != 1 + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err("shared offset conversion expects one pointer, one result, and no attributes or control flow".into()); + } + let mut operations = vec![]; + let pointer = if self.0 && input.operands[0].ty != MlirType::dialect("!llvm.ptr<7>")? { + let (result, pointer) = fresh(session, MlirType::dialect("!llvm.ptr<7>")?); + operations.push(operation( + "llvm.addrspacecast", + input.operands.clone(), + vec![result.clone()], + &input.location, + )?); + pointer + } else if input.operands[0].ty == MlirType::dialect("!llvm.ptr<7>")? { + input.operands[0].clone() + } else { + shared_pointer( + session, + &mut operations, + input.operands[0].clone(), + &input.location, + )? + }; + operations.push(operation( + "llvm.ptrtoint", + vec![pointer], + input.results, + &input.location, + )?); + Ok(operations) + } +} + +impl OperationTranslation for MmaTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mma = Tcgen05MmaOp::new(source); + mma.verify(ctx).map_err(|error| error.to_string())?; + let tensor_a = match mma.get_attr_nvvm_tcgen05_mma_form(ctx).as_deref() { + Some(Tcgen05MmaFormAttr::Shared) => false, + Some(Tcgen05MmaFormAttr::Tensor) => true, + _ => { + return Err( + "CUTLASS TCGen05 mapping supports dense shared-A or tensor-A FP16 MMA only" + .into(), + ); + } + }; + if mma.get_attr_nvvm_tcgen05_mma_kind(ctx).as_deref() != Some(&Tcgen05MmaKindAttr::F16) { + return Err( + "CUTLASS TCGen05 mapping supports dense shared-A or tensor-A FP16 MMA only".into(), + ); + } + let group = match mma.get_attr_nvvm_tcgen05_mma_cta_group(ctx).as_deref() { + Some(Tcgen05MmaCtaGroupAttr::Cg1) => 1, + Some(Tcgen05MmaCtaGroupAttr::Cg2) => 2, + None => return Err("TCGen05 MMA requires a CTA group".into()), + }; + let collector = match mma.get_attr_nvvm_tcgen05_mma_collector_a(ctx).as_deref() { + Some(Tcgen05MmaCollectorAAttr::Discard) => "discard", + Some(Tcgen05MmaCollectorAAttr::LastUse) => "lastuse", + Some(Tcgen05MmaCollectorAAttr::Fill) => "fill", + Some(Tcgen05MmaCollectorAAttr::Use) => "use", + None => return Err("TCGen05 MMA requires a collector-A selector".into()), + }; + require_input( + &mut input, + if tensor_a { "v1:i0764" } else { "v1:i0763" }, + 5, + 0, + )?; + let (cast, tmem) = tmem_pointer(session, input.operands[0].clone(), &input.location)?; + input.operands[0] = tmem; + let mut operations = vec![cast]; + if tensor_a { + let (cast_a, pointer_a) = + tmem_pointer(session, input.operands[1].clone(), &input.location)?; + operations.push(cast_a); + input.operands[1] = pointer_a; + } + let mut target = operation("nvvm.tcgen05.mma", input.operands, vec![], &input.location)?; + target.properties.insert( + "mmaKind".into(), + MlirAttribute::dialect("#nvvm.tcgen05_mma_kind")?, + ); + target + .properties + .insert("ctaGroup".into(), cta_group_value(group)?); + target.properties.insert( + "collectorOp".into(), + MlirAttribute::dialect(format!("#nvvm.tcgen05_mma_collectorop<{collector}>"))?, + ); + target.properties.insert( + "operandSegmentSizes".into(), + MlirAttribute::DenseI32Array(vec![1, 1, 1, 1, 1, 0, 0]), + ); + operations.push(target); + Ok(operations) + } +} + +/// Select one tensor-memory lifecycle or ordering instruction. +#[derive(Clone, Copy)] +enum Primitive { + /// Allocate the requested tensor-memory columns with the specified CTA group. + Alloc, + /// Release the requested tensor-memory column allocation. + Dealloc, + /// Signal completion of preceding tensor-core operations to a barrier. + Commit, + /// Relinquish the tensor-memory allocation permit. + Relinquish, + /// Wait for previously issued tensor-memory loads. + LoadWait, + /// Wait for previously issued tensor-memory stores. + StoreWait, + /// Order preceding generic accesses before tensor-core accesses. + FenceBefore, + /// Order following generic accesses after tensor-core accesses. + FenceAfter, +} + +/// Source-operation verifier run before discarding its intrinsic metadata. +type VerifyOperation = fn(&Context, Ptr) -> Result<(), String>; + +/// Validate a reviewed lifecycle instruction and emit its explicit CTA-group contract. +struct PrimitiveTranslation { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// Tensor-memory lifecycle or ordering action to emit. + kind: Primitive, + /// Number of CTAs participating in the native instruction. + group: u32, + /// Expected number of source runtime operands. + operands: usize, + /// Check the source operation's selectors and operand contract. + verify: VerifyOperation, +} + +impl OperationTranslation for PrimitiveTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + require_input(&mut input, self.marker, self.operands, 0)?; + let mut operations = vec![]; + if matches!(self.kind, Primitive::Dealloc) { + let (cast, pointer) = + tmem_pointer(session, input.operands[0].clone(), &input.location)?; + operations.push(cast); + input.operands[0] = pointer; + } + if matches!(self.kind, Primitive::Alloc | Primitive::Commit) { + input.operands[0] = shared_pointer( + session, + &mut operations, + input.operands[0].clone(), + &input.location, + )?; + } + let (name, kind) = match self.kind { + Primitive::Alloc => ("nvvm.tcgen05.alloc", None), + Primitive::Dealloc => ("nvvm.tcgen05.dealloc", None), + Primitive::Commit => ("nvvm.tcgen05.commit", None), + Primitive::Relinquish => ("nvvm.tcgen05.relinquish_alloc_permit", None), + Primitive::LoadWait => ("nvvm.tcgen05.wait", Some("#nvvm.tcgen05_wait")), + Primitive::StoreWait => ("nvvm.tcgen05.wait", Some("#nvvm.tcgen05_wait")), + Primitive::FenceBefore => ("nvvm.tcgen05.fence", Some("#nvvm.tcgen05_fence")), + Primitive::FenceAfter => ("nvvm.tcgen05.fence", Some("#nvvm.tcgen05_fence")), + }; + let mut target = operation(name, input.operands, vec![], &input.location)?; + if let Some(kind) = kind { + target + .properties + .insert("kind".into(), MlirAttribute::dialect(kind)?); + } else { + target + .properties + .insert("group".into(), cta_group_value(self.group)?); + } + operations.push(target); + Ok(operations) + } +} + +/// Scalar carriers retain the native register order. The native instructions +/// are asynchronous; only explicit wait operations introduce completion waits. +struct MemoryTranslation { + /// Canonical catalog identity required before emitting this intrinsic. + marker: &'static str, + /// Number of scalar register results or inputs in the memory instruction. + width: usize, + /// Select a tensor-memory store rather than a load. + store: bool, + /// Select the native packed memory-instruction variant. + pack: bool, + /// Check the source operation's selectors and operand contract. + verify: VerifyOperation, +} + +impl OperationTranslation for MemoryTranslation { + fn translate( + &self, + ctx: &Context, + source: Ptr, + mut input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + (self.verify)(ctx, source)?; + require_input( + &mut input, + self.marker, + if self.store { 1 + self.width } else { 1 }, + if self.store { 0 } else { self.width }, + )?; + let (cast, tmem) = tmem_pointer(session, input.operands[0].clone(), &input.location)?; + let mut operations = vec![cast]; + // The pinned NVVM lowering uses scalar i32 for x1; vector<1xi32> + // would create invalid extractelement/insertelement LLVM IR. + if self.width == 1 { + let mut target = if self.store { + operation( + "nvvm.tcgen05.st", + vec![tmem, input.operands[1].clone()], + vec![], + &input.location, + )? + } else { + operation( + "nvvm.tcgen05.ld", + vec![tmem], + input.results, + &input.location, + )? + }; + target.properties.insert( + "shape".into(), + MlirAttribute::dialect("#nvvm.tcgen05_ldst_shape")?, + ); + if self.pack { + target.properties.insert( + if self.store { "unpack" } else { "pack" }.into(), + MlirAttribute::Unit, + ); + } + operations.push(target); + return Ok(operations); + } + let vector_type = MlirType::Vector { + shape: vec![self.width as u64], + element: Box::new(MlirType::Integer(32)), + }; + let mut target; + if self.store { + let (undef, mut vector) = fresh(session, vector_type.clone()); + operations.push(operation( + "llvm.mlir.undef", + vec![], + vec![undef], + &input.location, + )?); + for (index, scalar) in input.operands.into_iter().skip(1).enumerate() { + let (result, next) = fresh(session, vector_type.clone()); + let mut insert = operation( + "vector.insert", + vec![scalar, vector], + vec![result], + &input.location, + )?; + insert.properties.insert( + "static_position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + operations.push(insert); + vector = next; + } + target = operation( + "nvvm.tcgen05.st", + vec![tmem, vector], + vec![], + &input.location, + )?; + } else { + let (result, vector) = fresh(session, vector_type); + target = operation("nvvm.tcgen05.ld", vec![tmem], vec![result], &input.location)?; + for (index, result) in input.results.into_iter().enumerate() { + let mut extract = operation( + "vector.extract", + vec![vector.clone()], + vec![result], + &input.location, + )?; + extract.properties.insert( + "static_position".into(), + MlirAttribute::DenseI64Array(vec![index as i64]), + ); + operations.push(extract); + } + } + target.properties.insert( + "shape".into(), + MlirAttribute::dialect("#nvvm.tcgen05_ldst_shape")?, + ); + if self.pack { + target.properties.insert( + if self.store { "unpack" } else { "pack" }.into(), + MlirAttribute::Unit, + ); + } + if self.store { + operations.push(target); + } else { + operations.insert(1, target); + } + Ok(operations) + } +} + +/// Check canonical intrinsic identity, operand/result counts, and absence of leftover source structure. +pub(crate) fn require_input( + input: &mut OperationInput, + marker: &str, + operands: usize, + results: usize, +) -> Result<(), String> { + if input.attributes.remove("cuda_oxide_intrinsic_marker") + != Some(MlirAttribute::String(marker.into())) + { + return Err(format!( + "TCGen05 operation expected intrinsic marker {marker}" + )); + } + if input.operands.len() != operands + || input.results.len() != results + || !input.attributes.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err(format!( + "TCGen05 operation expected {operands} operands, {results} results, and no residual attributes or control flow" + )); + } + Ok(()) +} + +/// Build the reviewed NVVM CTA-group attribute from the source group count. +fn cta_group_value(group: u32) -> Result { + MlirAttribute::dialect(format!("#nvvm.cta_group")) +} + +/// Retype the packed u32 TMEM token as an AS6 pointer without scaling its row/column bits. +fn tmem_pointer( + session: &mut TranslationSession<'_>, + token: MlirValueUse, + location: &MlirLocation, +) -> Result<(MlirOperation, MlirValueUse), String> { + if token.ty != MlirType::Integer(32) { + return Err("TMEM address must be a packed 32-bit token".into()); + } + let (result, pointer) = fresh(session, MlirType::dialect("!llvm.ptr<6>")?); + Ok(( + operation("llvm.inttoptr", vec![token], vec![result], location)?, + pointer, + )) +} + +/// Convert the source address to CTA shared space before forming its native offset. +fn shared_pointer( + session: &mut TranslationSession<'_>, + operations: &mut Vec, + pointer: MlirValueUse, + location: &MlirLocation, +) -> Result { + let shared = MlirType::dialect("!llvm.ptr<3>")?; + if pointer.ty == shared { + return Ok(pointer); + } + if pointer.ty != MlirType::dialect("!llvm.ptr")? { + return Err("TCGen05 allocation/commit needs a shared or generic pointer".into()); + } + let (result, converted) = fresh(session, shared); + operations.push(operation( + "llvm.addrspacecast", + vec![pointer], + vec![result], + location, + )?); + Ok(converted) +} + +/// Allocate one SSA identifier and return matching result and operand handles. +pub(crate) fn fresh( + session: &mut TranslationSession<'_>, + ty: MlirType, +) -> (MlirResult, MlirValueUse) { + let id = session.fresh_value(); + (MlirResult { id, ty: ty.clone() }, MlirValueUse { id, ty }) +} + +/// Build a target operation with the supplied operands, results, and source location. +pub(crate) fn operation( + name: &str, + operands: Vec, + results: Vec, + location: &MlirLocation, +) -> Result { + let mut operation = MlirOperation::new(name)?; + operation.operands = operands; + operation.results = results; + operation.location = location.clone(); + Ok(operation) +} + +#[cfg(test)] +mod tests { + //! Selector, identity, and register-order tests for tensor-memory instructions. + + use super::*; + use crate::render_module; + use crate::{CutlassFullCuteMlir22, MlirConsumerProfile}; + use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirPtrType, address_space}, + }; + use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + identifier::Identifier, + op::Op, + r#type::TypeHandle, + value::Value, + }; + + /// Append a typed source operation to the fixture block and return its operation handle. + fn append( + ctx: &mut Context, + block: Ptr, + operands: Vec, + results: Vec, + marker: Option<&str>, + ) -> Ptr { + let op = Operation::new(ctx, O::get_concrete_op_info(), results, operands, vec![], 0); + if let Some(marker) = marker { + op.deref_mut(ctx).attributes.set( + Identifier::try_from("cuda_oxide_intrinsic_marker").unwrap(), + StringAttr::new(marker.into()), + ); + } + op.insert_at_back(block, ctx); + op + } + + /// Build a split-N MMA source function with distinct descriptor, allocation, and register operands. + fn fixture() -> (Context, ModuleOp, Ptr) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "tcgen05".try_into().unwrap()); + let i1: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let i16: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let i64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get(&mut ctx, i64, true, address_space::SHARED).into(); + let generic: TypeHandle = + MirPtrType::get(&mut ctx, i64, true, address_space::GENERIC).into(); + // Arguments are deliberately distinct: do not alias D, descriptor A, + // descriptor B, instruction descriptor, columns, and enable-D. + let arguments = vec![shared, generic, i32, i64, i64, i32, i32, i1, i16]; + let fty = FunctionType::get(&ctx, arguments.clone(), vec![]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, f, TypeAttr::new(fty.into())); + function.set_symbol_name(&mut ctx, "tcgen05_probe".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let block = BasicBlock::new(&mut ctx, None, arguments); + block.insert_at_back(f.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + append::(&mut ctx, block, vec![a[1]], vec![i64], None); + append::(&mut ctx, block, vec![a[1], a[6]], vec![], Some("v1:i0359")); + append::( + &mut ctx, + block, + vec![], + vec![], + Some("v1:i0361"), + ); + let mut first = None; + for collector in [ + Tcgen05MmaCollectorAAttr::Fill, + Tcgen05MmaCollectorAAttr::LastUse, + ] { + let op = append::( + &mut ctx, + block, + vec![a[2], a[3], a[4], a[5], a[7]], + vec![], + Some("v1:i0763"), + ); + let mma = Tcgen05MmaOp::new(op); + mma.set_attr_nvvm_tcgen05_mma_form(&ctx, Tcgen05MmaFormAttr::Shared); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::F16); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&ctx, Tcgen05MmaCtaGroupAttr::Cg2); + mma.set_attr_nvvm_tcgen05_mma_collector_a(&ctx, collector); + first.get_or_insert(op); + } + append::( + &mut ctx, + block, + vec![a[0], a[8]], + vec![], + Some("v1:i0365"), + ); + append::( + &mut ctx, + block, + vec![a[2]], + vec![i32; 32], + Some("v1:i0664"), + ); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0357")); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0346")); + append::(&mut ctx, block, vec![], vec![], Some("v1:i0347")); + append::(&mut ctx, block, vec![a[2], a[6]], vec![], Some("v1:i0360")); + append::(&mut ctx, block, vec![], vec![], None); + (ctx, module, first.unwrap()) + } + + /// Export the source fixture and render mapping MLIR without its executable envelope. + fn translate(ctx: &Context, module: &ModuleOp) -> Result { + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile + .translate_module(ctx, module) + .map_err(|error| error.to_string())?; + render_module(&target).map_err(|error| error.to_string()) + } + + /// Check paired MMA retains collector fill/reuse/last-use order and TMEM scalar register order. + #[test] + fn split_n_collector_chain_and_tmem_register_order_survive_export() { + let (ctx, module, _) = fixture(); + let text = translate(&ctx, &module).unwrap(); + assert_eq!(text.matches("\"nvvm.tcgen05.mma\"").count(), 2, "{text}"); + assert!( + text.contains("#nvvm.tcgen05_mma_collectorop"), + "{text}" + ); + assert!( + text.contains("#nvvm.tcgen05_mma_collectorop"), + "{text}" + ); + assert_eq!(text.matches("\"llvm.inttoptr\"").count(), 4, "{text}"); + assert!( + text.contains("(!llvm.ptr<6>, i64, i64, i32, i1) -> ()"), + "{text}" + ); + assert!(text.contains("-> (vector<32xi32>)"), "{text}"); + assert_eq!(text.matches("\"vector.extract\"").count(), 32, "{text}"); + for index in 0..32 { + assert!( + text.contains(&format!("static_position = array")), + "{text}" + ); + } + assert_eq!(text.matches("\"nvvm.tcgen05.wait\"").count(), 1); + assert!( + text.find("\"nvvm.tcgen05.ld\"").unwrap() < text.find("\"nvvm.tcgen05.wait\"").unwrap() + ); + assert!(text.contains("#nvvm.tcgen05_fence")); + assert!(text.contains("#nvvm.tcgen05_fence")); + // A generic SMEM pointer must be converted before its integer offset + // is formed, otherwise descriptors encode a generic virtual address. + assert!( + text.find("\"llvm.addrspacecast\"").unwrap() < text.find("\"llvm.ptrtoint\"").unwrap() + ); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); + } + + /// Reject MMA variants whose source selector contracts have no reviewed mapping. + #[test] + fn unsupported_mma_forms_fail_instead_of_losing_their_selectors() { + let (ctx, module, first) = fixture(); + let mma = Tcgen05MmaOp::new(first); + mma.set_attr_nvvm_tcgen05_mma_kind(&ctx, Tcgen05MmaKindAttr::Tf32); + let error = translate(&ctx, &module).unwrap_err(); + assert!( + error.contains("dense shared-A or tensor-A FP16 MMA only"), + "{error}" + ); + } + + /// Check single-CTA attention MMA retains its tensor-memory A source. + #[test] + fn attention_mma_preserves_group_one_and_tensor_source_a() { + for tensor in [false, true] { + let (ctx, module, first) = fixture(); + let mma = Tcgen05MmaOp::new(first); + mma.set_attr_nvvm_tcgen05_mma_cta_group(&ctx, Tcgen05MmaCtaGroupAttr::Cg1); + if tensor { + mma.set_attr_nvvm_tcgen05_mma_form(&ctx, Tcgen05MmaFormAttr::Tensor); + let address = first.deref(&ctx).get_operand(0); + Operation::replace_operand(first, &ctx, 1, address); + first.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0764".into()), + ); + } + let text = translate(&ctx, &module).unwrap(); + assert!(text.contains("#nvvm.cta_group"), "{text}"); + if tensor { + assert!( + text.contains("(!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1) -> ()"), + "{text}" + ); + } + } + } + + /// Check attention loads/stores retain register counts, packing modes, and caller-specified waits. + #[test] + fn attention_tmem_memory_preserves_width_pack_and_explicit_waits() { + let (mut ctx, module, first) = fixture(); + let block = first.deref(&ctx).get_parent_block().unwrap(); + let a = block.deref(&ctx).arguments().collect::>(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + /// Append a marked tensor-memory operation to the attention memory test fixture. + macro_rules! insert { + ($ty:ty, $args:expr, $results:expr, $marker:literal) => {{ + let op = Operation::new( + &mut ctx, + <$ty>::get_concrete_op_info(), + $results, + $args, + vec![], + 0, + ); + op.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new($marker.into()), + ); + op.insert_before(&ctx, first); + op + }}; + } + insert!( + dialect_nvvm::ops::Tcgen05AllocOp, + vec![a[0], a[2]], + vec![], + "v1:i0343" + ); + insert!( + dialect_nvvm::ops::Tcgen05RelinquishAllocPermitOp, + vec![], + vec![], + "v1:i0345" + ); + insert!( + dialect_nvvm::ops::Tcgen05CommitOp, + vec![a[0]], + vec![], + "v1:i0348" + ); + insert!( + dialect_nvvm::ops::Tcgen05CommitCg2Op, + vec![a[0]], + vec![], + "v1:i0363" + ); + insert!( + dialect_nvvm::ops::Tcgen05CommitMulticastOp, + vec![a[0], a[8]], + vec![], + "v1:i0760" + ); + insert!( + dialect_nvvm::ops::Tcgen05DeallocOp, + vec![a[2], a[2]], + vec![], + "v1:i0344" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX1RawOp, + vec![a[2]], + vec![i32; 1], + "v1:i0654" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX1Pack16Op, + vec![a[2]], + vec![i32; 1], + "v1:i0655" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX1RawOp, + vec![a[2]; 2], + vec![], + "v1:i0712" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX1Unpack16Op, + vec![a[2]; 2], + vec![], + "v1:i0713" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX2RawOp, + vec![a[2]], + vec![i32; 2], + "v1:i0656" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX2Pack16Op, + vec![a[2]], + vec![i32; 2], + "v1:i0657" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX2RawOp, + vec![a[2]; 3], + vec![], + "v1:i0714" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX2Unpack16Op, + vec![a[2]; 3], + vec![], + "v1:i0715" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX4RawOp, + vec![a[2]], + vec![i32; 4], + "v1:i0658" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX4Pack16Op, + vec![a[2]], + vec![i32; 4], + "v1:i0659" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX4RawOp, + vec![a[2]; 5], + vec![], + "v1:i0716" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX4Unpack16Op, + vec![a[2]; 5], + vec![], + "v1:i0717" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX8RawOp, + vec![a[2]], + vec![i32; 8], + "v1:i0660" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX8Pack16Op, + vec![a[2]], + vec![i32; 8], + "v1:i0661" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX8RawOp, + vec![a[2]; 9], + vec![], + "v1:i0718" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX8Unpack16Op, + vec![a[2]; 9], + vec![], + "v1:i0719" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX16RawOp, + vec![a[2]], + vec![i32; 16], + "v1:i0662" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX16Pack16Op, + vec![a[2]], + vec![i32; 16], + "v1:i0663" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX16RawOp, + vec![a[2]; 17], + vec![], + "v1:i0720" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX16Unpack16Op, + vec![a[2]; 17], + vec![], + "v1:i0721" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX32RawOp, + vec![a[2]], + vec![i32; 32], + "v1:i0664" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX32Pack16Op, + vec![a[2]], + vec![i32; 32], + "v1:i0665" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX32RawOp, + vec![a[2]; 33], + vec![], + "v1:i0722" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX32Unpack16Op, + vec![a[2]; 33], + vec![], + "v1:i0723" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX64RawOp, + vec![a[2]], + vec![i32; 64], + "v1:i0666" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX64Pack16Op, + vec![a[2]], + vec![i32; 64], + "v1:i0667" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX64RawOp, + vec![a[2]; 65], + vec![], + "v1:i0724" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX64Unpack16Op, + vec![a[2]; 65], + vec![], + "v1:i0725" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX128RawOp, + vec![a[2]], + vec![i32; 128], + "v1:i0668" + ); + insert!( + dialect_nvvm::ops::Tcgen05Ld32x32bX128Pack16Op, + vec![a[2]], + vec![i32; 128], + "v1:i0669" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX128RawOp, + vec![a[2]; 129], + vec![], + "v1:i0726" + ); + insert!( + dialect_nvvm::ops::Tcgen05St32x32bX128Unpack16Op, + vec![a[2]; 129], + vec![], + "v1:i0727" + ); + insert!( + dialect_nvvm::ops::Tcgen05StoreWaitOp, + vec![], + vec![], + "v1:i0358" + ); + let text = translate(&ctx, &module).unwrap(); + assert_eq!(text.matches("\"nvvm.tcgen05.ld\"").count(), 17); + assert_eq!(text.matches("\"nvvm.tcgen05.st\"").count(), 16); + assert_eq!(text.matches("pack = unit").count(), 16); + assert_eq!(text.matches("\"nvvm.tcgen05.wait\"").count(), 2); + assert!(text.contains("#nvvm.tcgen05_wait")); + assert!(text.contains("#nvvm.cta_group")); + for width in [2, 4, 8, 16, 32, 64, 128] { + assert!(text.contains(&format!("vector<{width}xi32>")), "{text}"); + } + } + + /// Reject a source instruction carrying a forged native catalog marker. + #[test] + fn intrinsic_identity_is_checked_before_export() { + let (ctx, module, first) = fixture(); + first.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0001".into()), + ); + let error = translate(&ctx, &module).unwrap_err(); + assert!( + error.contains("expected intrinsic marker v1:i0763"), + "{error}" + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs new file mode 100644 index 0000000000..9d849258e5 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/packs.rs @@ -0,0 +1,367 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builtin scalar types, constants, and attributes shared by every export profile. + +use crate::{ + AttributeTranslation, DropAttribute, FixedType, MlirAttribute, MlirFloatType, MlirOperation, + MlirType, OperationInput, OperationTranslation, TranslationError, TranslationRegistry, + TranslationSession, TypeTranslation, +}; +use pliron::{ + attribute::AttrObj, + builtin::{ + attr_interfaces::TypedAttrInterface, + attributes::{ + BoolAttr, FPDoubleAttr, FPHalfAttr, FPSingleAttr, IdentifierAttr, IntegerAttr, + OperandSegmentSizesAttr, StringAttr, TypeAttr, UnitAttr, + }, + ops::{ConstantOp, ModuleOp}, + type_interfaces::FunctionTypeInterface, + types::{FP16Type, FP32Type, FP64Type, FunctionType, IntegerType, UnitType}, + }, + context::{Context, Ptr}, + operation::Operation, + r#type::{TypeHandle, Typed, type_cast}, + utils::apfloat::Float, +}; + +/// Register Pliron builtin items used by every CUDA mapping profile. +pub fn register_builtin_pack(registry: &mut TranslationRegistry) -> Result<(), TranslationError> { + registry.register_operation::(BuiltinModuleTranslation)?; + registry.register_operation::(BuiltinConstantTranslation)?; + + registry.register_type::(IntegerTypeTranslation)?; + registry.register_type::(FixedType(MlirType::Float(MlirFloatType::F16)))?; + registry.register_type::(FixedType(MlirType::Float(MlirFloatType::F32)))?; + registry.register_type::(FixedType(MlirType::Float(MlirFloatType::F64)))?; + registry.register_type::(FunctionTypeTranslation)?; + registry.register_type::(UnitTypeTranslation)?; + + registry.register_attribute::(IdentifierAttributeTranslation)?; + registry.register_attribute::(StringAttributeTranslation)?; + registry.register_attribute::(BoolAttributeTranslation)?; + registry.register_attribute::(IntegerAttributeTranslation)?; + registry.register_attribute::(FloatAttributeTranslation::F16)?; + registry.register_attribute::(FloatAttributeTranslation::F32)?; + registry.register_attribute::(FloatAttributeTranslation::F64)?; + registry.register_attribute::(UnitAttributeTranslation)?; + registry.register_attribute::(TypeAttributeTranslation)?; + // This describes how Pliron flattened a variadic operand list. Recipes + // such as `mir.cond_br` rebuild the corresponding target property from + // the source op instead of copying a Pliron-only attribute name. + registry.register_attribute::(DropAttribute)?; + Ok(()) +} + +/// Translate a builtin module and preserve its nested regions and optional symbol name. +struct BuiltinModuleTranslation; + +impl OperationTranslation for BuiltinModuleTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mut target = MlirOperation::new("builtin.module")?; + target.results = input.results; + target.operands = input.operands; + target.successors = input.successors; + target.regions = input.regions; + target.attributes = input.attributes; + target.location = input.location; + if let Some(name) = target + .attributes + .remove(pliron::builtin::op_interfaces::ATTR_KEY_SYM_NAME.as_ref()) + { + target.properties.insert("sym_name".into(), name); + } + Ok(vec![target]) + } +} + +// Scalar constants created by canonicalization and loop-unrolling passes use +// builtin.constant rather than the source-facing mir.constant operation. +/// Map canonicalized builtin scalar constants to arith.constant. +struct BuiltinConstantTranslation; +impl OperationTranslation for BuiltinConstantTranslation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + mut input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + if input.results.len() != 1 + || !input.operands.is_empty() + || !input.regions.is_empty() + || !input.successors.is_empty() + { + return Err( + "builtin.constant requires one result and no operands or control flow".into(), + ); + } + let value = input + .attributes + .remove("builtin_constant_value") + .ok_or("builtin.constant requires a typed value")?; + let ty = match &value { + MlirAttribute::Integer { ty, .. } | MlirAttribute::Float { ty, .. } => ty, + _ => { + return Err( + "builtin.constant currently supports scalar integer and floating values".into(), + ); + } + }; + if *ty != input.results[0].ty || !input.attributes.is_empty() { + return Err( + "builtin.constant value type or attributes violate the scalar contract".into(), + ); + } + let mut target = MlirOperation::new("arith.constant")?; + target.results = input.results; + target.location = input.location; + target.properties.insert("value".into(), value); + Ok(vec![target]) + } +} + +/// Retain integer width while erasing source signedness from the MLIR type spelling. +struct IntegerTypeTranslation; + +impl TypeTranslation for IntegerTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let integer = source_ref + .downcast_ref::() + .ok_or_else(|| "expected builtin.integer".to_owned())?; + Ok(MlirType::Integer(integer.width())) + } +} + +/// Recursively translate function arguments and results through the active registry. +struct FunctionTypeTranslation; + +impl TypeTranslation for FunctionTypeTranslation { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result { + let source_ref = source.deref(ctx); + let function = type_cast::(&*source_ref) + .ok_or_else(|| "expected builtin function type".to_owned())?; + let inputs = function + .arg_types() + .into_iter() + .map(|ty| translate_nested_type(ctx, registry, ty)) + .collect::>()?; + let results = function + .res_types() + .into_iter() + .map(|ty| translate_nested_type(ctx, registry, ty)) + .collect::>()?; + Ok(MlirType::Function { inputs, results }) + } +} + +/// Represent Rust unit as an empty aggregate rather than inventing a scalar value. +struct UnitTypeTranslation; + +impl TypeTranslation for UnitTypeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + // Unit is only legal as an empty aggregate carrier. Mapping it to an + // empty tuple keeps that fact explicit instead of inventing i1/i8. + Ok(MlirType::Tuple(vec![])) + } +} + +/// Translate a nested source type and preserve any registry diagnostic. +fn translate_nested_type( + ctx: &Context, + registry: &TranslationRegistry, + source: TypeHandle, +) -> Result { + registry.translate_type(ctx, source) +} + +/// Preserve the bytes of a builtin string attribute. +struct StringAttributeTranslation; + +impl AttributeTranslation for StringAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.string".to_owned())?; + Ok(Some(MlirAttribute::String(value.as_str().into()))) + } +} + +/// Represent a source identifier as the target string spelling. +struct IdentifierAttributeTranslation; + +impl AttributeTranslation for IdentifierAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.identifier".to_owned())?; + Ok(Some(MlirAttribute::String(value.as_ref().to_string()))) + } +} + +/// Preserve a builtin boolean attribute's truth value. +struct BoolAttributeTranslation; + +impl AttributeTranslation for BoolAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.bool".to_owned())? + .clone(); + Ok(Some(MlirAttribute::Bool(value.into()))) + } +} + +/// Preserve integer width and all bits in the target signless attribute. +struct IntegerAttributeTranslation; + +impl AttributeTranslation for IntegerAttributeTranslation { + fn translate( + &self, + ctx: &Context, + source: &AttrObj, + registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.integer".to_owned())?; + let source_type = Typed::get_type(value, ctx); + let width = value.get_type().deref(ctx).width(); + if width > 128 { + return Err(format!( + "integer attributes wider than 128 bits are not supported yet (got i{width})" + )); + } + Ok(Some(MlirAttribute::Integer { + // MLIR integer types are signless. Rendering the same two's- + // complement bits as an i128 preserves both signed and unsigned + // Pliron constants, including u64::MAX -> -1 : i64. + value: value.value().to_i128(), + ty: registry.translate_type(ctx, source_type)?, + })) + } +} + +/// Select the source floating-point width before preserving its value in MLIR. +enum FloatAttributeTranslation { + /// Translate a half-precision floating-point attribute. + F16, + /// Translate a single-precision floating-point attribute. + F32, + /// Translate a double-precision floating-point attribute. + F64, +} + +impl AttributeTranslation for FloatAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + let (value, ty) = match self { + Self::F16 => { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.half".to_owned())?; + ( + format!("0x{:04X}", value.0.to_bits()), + MlirType::Float(MlirFloatType::F16), + ) + } + Self::F32 => { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.single".to_owned())?; + ( + format!("0x{:08X}", value.0.to_bits()), + MlirType::Float(MlirFloatType::F32), + ) + } + Self::F64 => { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.double".to_owned())?; + ( + format!("0x{:016X}", value.0.to_bits()), + MlirType::Float(MlirFloatType::F64), + ) + } + }; + Ok(Some(MlirAttribute::Float { value, ty })) + } +} + +/// Translate a source presence marker to an MLIR unit attribute. +struct UnitAttributeTranslation; + +impl AttributeTranslation for UnitAttributeTranslation { + fn translate( + &self, + _ctx: &Context, + _source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + Ok(Some(MlirAttribute::Unit)) + } +} + +/// Recursively translate the type carried by a source type attribute. +struct TypeAttributeTranslation; + +impl AttributeTranslation for TypeAttributeTranslation { + fn translate( + &self, + ctx: &Context, + source: &AttrObj, + registry: &TranslationRegistry, + ) -> Result, String> { + let value = source + .downcast_ref::() + .ok_or_else(|| "expected builtin.type".to_owned())?; + let ty = TypedAttrInterface::get_type(value, ctx); + Ok(Some(MlirAttribute::Type(translate_nested_type( + ctx, registry, ty, + )?))) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs new file mode 100644 index 0000000000..1492e7325f --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/primitives.rs @@ -0,0 +1,250 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Pinned instruction-level consumer, independent of CuTe source semantics. + +use crate::{CutlassFullCuteMlir22, ExportManifest, MlirConsumerProfile, ProfileError}; +use crate::{MlirModule, TranslationRegistry, translate_module}; +use dialect_mir::types::*; +use pliron::{ + attribute::{Attribute, attr_type}, + builtin::{ + attributes::{DictAttr, VecAttr}, + ops::ModuleOp, + type_interfaces::FunctionTypeInterface, + }, + context::Context, + dialect::DialectName, + linked_list::ContainsLinkedList, + location::Location, + op::Op, + operation::{Operation, verify_operation}, + printable::Printable, + r#type::{TypeHandle, Typed, type_cast}, +}; +use std::collections::HashSet; + +/// CUTLASS 4.7 / MLIR 22 primitives for SM100a and SM120a. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutlassPrimitivesMlir22 { + /// Validated CUDA target architecture, including its architecture-specific suffix. + gpu_arch: String, +} + +impl CutlassPrimitivesMlir22 { + /// Stable profile identity included in manifests and build cache keys. + pub const NAME: &'static str = "cutlass-primitives-mlir22"; + + /// Accept only the reviewed sm_100a and sm_120a primitive targets. + pub fn new(gpu_arch: impl Into) -> Result { + let gpu_arch = gpu_arch.into(); + if !matches!(gpu_arch.as_str(), "sm_100a" | "sm_120a") { + return Err(ProfileError::UnsupportedArchitecture(gpu_arch)); + } + Ok(Self { gpu_arch }) + } +} + +impl MlirConsumerProfile for CutlassPrimitivesMlir22 { + fn name(&self) -> &'static str { + Self::NAME + } + fn manifest(&self) -> ExportManifest { + ExportManifest { + profile: Self::NAME.into(), + exporter_version: env!("CARGO_PKG_VERSION").into(), + pliron_revision: CutlassFullCuteMlir22::PLIRON_REVISION.into(), + interop_revision: crate::INTEROP_REVISION.into(), + consumer_revision: CutlassFullCuteMlir22::CONSUMER_REVISION.into(), + llvm_major: 22, + gpu_arch: self.gpu_arch.clone(), + pointer_model: "opaque".into(), + index_bits: 64, + } + } + fn build_registry(&self) -> Result { + let mut r = TranslationRegistry::new(); + crate::register_builtin_pack(&mut r)?; + crate::register_mir_core_pack(&mut r)?; + crate::mir::shared::register_dynamic_shared(&mut r)?; + crate::register_nvvm_sreg_pack(&mut r)?; + crate::register_nvvm_tcgen05_pack(&mut r)?; + crate::register_nvvm_cluster_pack(&mut r)?; + crate::register_nvvm_sync_restrict_pack(&mut r)?; + crate::nvvm::numeric::register_nvvm_profile_numeric_pack(&mut r, &self.gpu_arch)?; + crate::register_nvvm_sm120_pack(&mut r)?; + r.seal(); + Ok(r) + } + fn verify_source_module(&self, ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError> { + verify_primitives_source(ctx, module)?; + crate::profile::verify_nvvm_architecture(ctx, module, &self.gpu_arch) + } + fn translate_module( + &self, + ctx: &Context, + module: &ModuleOp, + ) -> Result { + self.verify_source_module(ctx, module)?; + let mut target = translate_module( + ctx, + module, + &self.build_registry()?, + &self.translation_config(), + )?; + crate::profile::add_cutlass_executable_envelope(&mut target) + .map_err(ProfileError::ExecutableEnvelope)?; + Ok(target) + } +} + +/// Check before preparation can erase ghost types or opaque pointees. +pub fn verify_primitives_source(ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError> { + let mut scan = SourceScan { + ctx, + cute: DialectName::try_new("cute").unwrap(), + seen: HashSet::new(), + }; + let mut pending = vec![module.get_operation()]; + while let Some(ptr) = pending.pop() { + let id = Operation::get_opid(ptr, ctx); + if id.dialect == scan.cute { + return Err(rejected(format!("CuTe operation `{id}`"))); + } + let op = ptr.deref(ctx); + for value in op.operands().chain(op.results()) { + scan.ty(value.get_type(ctx))?; + } + for attr in op.attributes.0.values() { + scan.attr(&**attr)?; + } + scan.location(op.loc())?; + for region in op.regions() { + for block in region.deref(ctx).iter(ctx) { + let block = block.deref(ctx); + for arg in block.arguments() { + scan.ty(arg.get_type(ctx))?; + } + for attr in block.attributes.0.values() { + scan.attr(&**attr)?; + } + scan.location(block.loc())?; + pending.extend(block.iter(ctx)); + } + } + } + dialect_mir::verification::verify_pointer_kind_producers(ctx, module.get_operation()) + .map_err(|e| ProfileError::SourceValidation(e.disp(ctx).to_string()))?; + verify_operation(module.get_operation(), ctx) + .map_err(|e| ProfileError::SourceValidation(e.disp(ctx).to_string())) +} + +/// Build a source-validation error identifying a forbidden CuTe item. +fn rejected(item: String) -> ProfileError { + ProfileError::SourceValidation(format!("{item} is unsupported by the primitives profile")) +} + +/// Recursively reject CuTe semantics hidden inside source metadata or nested types. +struct SourceScan<'a> { + /// Immutable source context used to inspect registered type and attribute objects. + ctx: &'a Context, + /// Dialect identity that the primitives-only profile must reject. + cute: DialectName, + /// Type handles already checked, preventing recursive metadata from being scanned repeatedly. + seen: HashSet, +} +impl SourceScan<'_> { + /// Check a source type and recursively inspect its pointees, members, and function signature. + fn ty(&mut self, handle: TypeHandle) -> Result<(), ProfileError> { + if !self.seen.insert(handle) { + return Ok(()); + } + let ty = handle.deref(self.ctx); + let id = ty.get_type_id(); + if id.dialect == self.cute { + return Err(rejected(format!("CuTe type `{id}`"))); + } + let children: Vec = + if let Some(f) = type_cast::(&*ty) { + f.arg_types().into_iter().chain(f.res_types()).collect() + } else if let Some(t) = ty.downcast_ref::() { + vec![t.pointee] + } else if let Some(t) = ty.downcast_ref::() { + t.types.clone() + } else if let Some(t) = ty.downcast_ref::() { + t.field_types.clone() + } else if let Some(t) = ty.downcast_ref::() { + t.field_types.clone() + } else if let Some(t) = ty.downcast_ref::() { + vec![t.element_ty] + } else if let Some(t) = ty.downcast_ref::() { + vec![t.element_ty] + } else if let Some(t) = ty.downcast_ref::() { + std::iter::once(t.element_ty) + .chain(t.space_tys.iter().copied()) + .collect() + } else if let Some(t) = ty.downcast_ref::() { + std::iter::once(t.discriminant_ty) + .chain(t.all_field_types.iter().copied()) + .collect() + } else if let Some(t) = ty.downcast_ref::() { + vec![t.elem_type()] + } else if let Some(t) = ty.downcast_ref::() { + vec![t.elem_type()] + } else if let Some(t) = ty.downcast_ref::() { + t.fields().collect() + } else { + vec![] + }; + for child in children { + self.ty(child)?; + } + Ok(()) + } + /// Check a source attribute and recursively inspect nested dictionaries, vectors, and types. + fn attr(&mut self, attr: &dyn Attribute) -> Result<(), ProfileError> { + let id = attr.get_attr_id(); + if id.dialect == self.cute { + return Err(rejected(format!("CuTe attribute `{id}`"))); + } + if let Some(ty) = attr_type(attr, self.ctx) { + self.ty(ty)?; + } + if let Some(v) = attr.downcast_ref::() { + for a in &v.0 { + self.attr(&**a)?; + } + } + if let Some(d) = attr.downcast_ref::() { + for a in d.0.0.values() { + self.attr(&**a)?; + } + } + Ok(()) + } + /// Inspect location metadata for nested attributes carrying forbidden CuTe semantics. + fn location(&mut self, loc: &Location) -> Result<(), ProfileError> { + match loc { + Location::Fused { + metadata, + locations, + } => { + if let Some(a) = metadata { + self.attr(&**a)?; + } + for loc in locations { + self.location(loc)?; + } + } + Location::Named { child_loc, .. } => self.location(child_loc)?, + Location::CallSite { callee, caller } => { + self.location(callee)?; + self.location(caller)?; + } + _ => {} + } + Ok(()) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs new file mode 100644 index 0000000000..a202692cc3 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/profile.rs @@ -0,0 +1,442 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Pinned consumer contracts, source validation, and executable GPU module packaging. + +use crate::{ + MlirAttribute, MlirBlock, MlirBlockId, MlirModule, MlirOperation, MlirRegion, + TranslationConfig, TranslationError, TranslationRegistry, + translate_module as translate_pliron_module, +}; +use pliron::{ + builtin::ops::ModuleOp, context::Context, linked_list::ContainsLinkedList, op::Op, + operation::Operation, +}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::{ + register_builtin_pack, register_cute_cutlass_profile_pack, register_cute_gemm_pack, + register_cute_sm100_pack, register_mir_core_pack, register_nvvm_cluster_pack, + register_nvvm_sreg_pack, register_nvvm_tcgen05_pack, +}; + +/// Everything that makes one textual MLIR contract reproducible. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExportManifest { + /// Stable consumer-profile identity used to distinguish mapping contracts. + pub profile: String, + /// Version of the cuda-oxide MLIR exporter that produced the artifact. + pub exporter_version: String, + /// Exact Pliron revision whose source IR contract was translated. + pub pliron_revision: String, + /// Published upstream printer version and local compatibility adapter revision. + pub interop_revision: String, + /// Pinned CUTLASS compiler release expected to accept this target MLIR. + pub consumer_revision: String, + /// LLVM/MLIR major version defining the target dialect syntax. + pub llvm_major: u32, + /// Validated GPU architecture used for source capability checks and compilation. + pub gpu_arch: String, + /// Target pointer representation, currently opaque LLVM pointers. + pub pointer_model: String, + /// Bit width used for target index values. + pub index_bits: u32, +} + +impl ExportManifest { + /// Serialize the complete reproducibility manifest as indented JSON. + pub fn to_pretty_json(&self) -> Result { + serde_json::to_string_pretty(self) + } +} + +/// Failure to validate, register, or package a source module for the selected consumer. +#[derive(Debug, Error)] +pub enum ProfileError { + /// A mapping registration or generic translation failure. + #[error(transparent)] + Registry( + #[doc = "Underlying registration or generic translation diagnostic."] + #[from] + TranslationError, + ), + /// The shared high-level CuTe verifier rejected the source module. + #[error("CuTe source validation failed before MLIR translation: {0}")] + CuteSemantics(#[doc = "Diagnostic from the shared CuTe source verifier."] String), + /// The selected profile rejected a source instruction or metadata contract. + #[error("primitives source validation failed: {0}")] + SourceValidation(#[doc = "Reason the source violates the selected profile's contract."] String), + /// The translated module could not be packaged for CUTLASS binary emission. + #[error("could not build CUTLASS's executable GPU module envelope: {0}")] + ExecutableEnvelope( + #[doc = "Reason a valid executable module envelope could not be constructed."] String, + ), + /// The requested target architecture has no reviewed profile. + #[error("unsupported GPU architecture `{0}` for this MLIR profile")] + UnsupportedArchitecture( + #[doc = "Unrecognized or unreviewed target architecture spelling."] String, + ), +} + +/// A pinned target contract plus the mapping packs that implement it. +pub trait MlirConsumerProfile { + /// Return the stable profile identity used in diagnostics and cache manifests. + fn name(&self) -> &'static str; + /// Describe the exact exporter, source IR, consumer, and target contract. + fn manifest(&self) -> ExportManifest; + /// Register the profile's reviewed mappings exactly once and seal the resulting registry. + fn build_registry(&self) -> Result; + + /// Check profile-specific source rules before any target text is built. + fn verify_source_module(&self, ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError>; + + /// Configure generic translation diagnostics with this profile's identity. + fn translation_config(&self) -> TranslationConfig { + TranslationConfig::new(self.name()) + } + + /// Validate the shared high-level module, then translate it. + fn translate_module( + &self, + ctx: &Context, + module: &ModuleOp, + ) -> Result { + self.verify_source_module(ctx, module)?; + let registry = self.build_registry()?; + Ok(translate_pliron_module( + ctx, + module, + ®istry, + &self.translation_config(), + )?) + } +} + +/// CuTe mappings supported by the pinned CUTLASS 4.7 compiler. +/// +/// Each CUTLASS release needs its own profile until its syntax, generated code, +/// and runtime behavior have passed validation. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutlassFullCuteMlir22 { + /// Validated GPU architecture used for source capability checks and compilation. + gpu_arch: String, +} + +impl CutlassFullCuteMlir22 { + /// Exact CUTLASS compiler release defining this profile's accepted dialects. + pub const CONSUMER_REVISION: &'static str = "cutlass-compiler-v4.7.0"; + /// Exact Pliron revision defining this profile's source IR contract. + pub const PLIRON_REVISION: &'static str = "435536c8611af34c48c8f73d78e77894a05d103d"; + + /// Accept only sm_100a or sm_120a, whose mappings have been validated for this consumer. + pub fn new(gpu_arch: impl Into) -> Result { + let gpu_arch = gpu_arch.into(); + if !matches!(gpu_arch.as_str(), "sm_100a" | "sm_120a") { + return Err(ProfileError::UnsupportedArchitecture(gpu_arch)); + } + Ok(Self { gpu_arch }) + } +} + +impl MlirConsumerProfile for CutlassFullCuteMlir22 { + fn name(&self) -> &'static str { + "cutlass-full-cute-mlir22" + } + + fn manifest(&self) -> ExportManifest { + ExportManifest { + profile: self.name().into(), + exporter_version: env!("CARGO_PKG_VERSION").into(), + pliron_revision: Self::PLIRON_REVISION.into(), + interop_revision: crate::INTEROP_REVISION.into(), + consumer_revision: Self::CONSUMER_REVISION.into(), + llvm_major: 22, + gpu_arch: self.gpu_arch.clone(), + pointer_model: "opaque".into(), + index_bits: 64, + } + } + + fn build_registry(&self) -> Result { + let mut registry = TranslationRegistry::new(); + register_builtin_pack(&mut registry)?; + register_mir_core_pack(&mut registry)?; + register_nvvm_sreg_pack(&mut registry)?; + register_nvvm_tcgen05_pack(&mut registry)?; + register_nvvm_cluster_pack(&mut registry)?; + crate::register_nvvm_sm120_pack(&mut registry)?; + crate::register_nvvm_sync_restrict_pack(&mut registry)?; + crate::nvvm::numeric::register_nvvm_profile_numeric_pack(&mut registry, &self.gpu_arch)?; + register_cute_cutlass_profile_pack(&mut registry)?; + register_cute_gemm_pack(&mut registry)?; + register_cute_sm100_pack(&mut registry)?; + crate::register_cute_attention_pack(&mut registry)?; + crate::cute::register_cute_sm100_attention_pack(&mut registry)?; + registry.seal(); + Ok(registry) + } + + fn verify_source_module(&self, ctx: &Context, module: &ModuleOp) -> Result<(), ProfileError> { + // Both backends check the same CuTe operation graph and tensor origins + // before generating code. Verification leaves the module unchanged. + dialect_cute::verify::verify_cute_semantics(ctx, module.get_operation()) + .map_err(|error| ProfileError::CuteSemantics(error.to_string()))?; + verify_nvvm_architecture(ctx, module, &self.gpu_arch) + } + + fn translate_module( + &self, + ctx: &Context, + module: &ModuleOp, + ) -> Result { + self.verify_source_module(ctx, module)?; + let registry = self.build_registry()?; + let mut target = + translate_pliron_module(ctx, module, ®istry, &self.translation_config())?; + add_cutlass_executable_envelope(&mut target).map_err(ProfileError::ExecutableEnvelope)?; + Ok(target) + } +} + +/// Raw NVVM intrinsics bypass CuTe atom selection, so both profiles must check +/// their architecture requirements before translating any operation. +pub(crate) fn verify_nvvm_architecture( + ctx: &Context, + module: &ModuleOp, + gpu_arch: &str, +) -> Result<(), ProfileError> { + let mut pending = vec![module.get_operation()]; + while let Some(ptr) = pending.pop() { + let id = Operation::get_opid(ptr, ctx); + if gpu_arch == "sm_120a" && id.to_string().starts_with("nvvm.tcgen05") { + return Err(ProfileError::SourceValidation( + "TCGen05 is unsupported on sm_120a".into(), + )); + } + if gpu_arch != "sm_120a" + && id == dialect_nvvm::ops::RegisterMmaOp::get_opid_static() + && dialect_nvvm::ops::RegisterMmaOp::new(ptr) + .get_attr_nvvm_register_mma_kind(ctx) + .as_deref() + == Some(&dialect_nvvm::ops::RegisterMmaKindAttr::Mxf4) + { + return Err(ProfileError::SourceValidation( + "MXF4 register MMA requires sm_120a".into(), + )); + } + for region in ptr.deref(ctx).regions() { + for block in region.deref(ctx).iter(ctx) { + pending.extend(block.deref(ctx).iter(ctx)); + } + } + } + Ok(()) +} + +/// CUTLASS's binary pass only serializes `gpu.module` operations. Keeping device +/// functions at the top level is syntactically valid but makes the compiler +/// return success without producing PTX or a cubin, so the profile always +/// builds the executable envelope itself. +pub(crate) fn add_cutlass_executable_envelope(module: &mut MlirModule) -> Result<(), String> { + if module.root.name != "builtin.module" { + return Err(format!( + "expected builtin.module root, got {}", + module.root.name + )); + } + if module.root.regions.len() != 1 || module.root.regions[0].blocks.len() != 1 { + return Err("builtin.module must contain exactly one top-level block".into()); + } + + let gpu_block_id = next_block_id(&module.root)?; + let root_block = &mut module.root.regions[0].blocks[0]; + let mut device_operations = std::mem::take(&mut root_block.operations); + crate::mir::shared::hoist_shared_globals(&mut device_operations)?; + + let mut gpu_module = MlirOperation::new("gpu.module")?; + gpu_module + .properties + .insert("sym_name".into(), MlirAttribute::String("kernels".into())); + gpu_module.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: gpu_block_id, + arguments: vec![], + operations: device_operations, + }], + }); + gpu_module.location = module.root.location.clone(); + + // The source module name is useful before translation, while CUTLASS's runtime + // contract names the serialized device module `kernels`. + module.root.properties.remove("sym_name"); + module + .root + .attributes + .insert("gpu.container_module".into(), MlirAttribute::Unit); + root_block.operations.push(gpu_module); + Ok(()) +} + +/// Choose a fresh block identifier beyond every block already present in the module. +fn next_block_id(operation: &MlirOperation) -> Result { + /// Walk nested regions to update the largest existing block identifier. + fn visit(operation: &MlirOperation, maximum: &mut u64) { + for region in &operation.regions { + for block in ®ion.blocks { + *maximum = (*maximum).max(block.id.0); + for operation in &block.operations { + visit(operation, maximum); + } + } + } + } + + let mut maximum = 0; + visit(operation, &mut maximum); + maximum + .checked_add(1) + .map(MlirBlockId) + .ok_or_else(|| "MLIR block id space is exhausted".into()) +} + +/// Keep the operation-mapping goldens focused on their mapping pack. The +/// executable envelope has its own structural tests below and in the +/// production-shaped integration fixture. +#[cfg(test)] +pub(crate) fn render_mapping_module_without_cutlass_envelope( + module: &MlirModule, + source_name: &str, +) -> String { + let mut unwrapped = module.clone(); + let gpu_module = &module.root.regions[0].blocks[0].operations[0]; + assert_eq!(gpu_module.name, "gpu.module"); + unwrapped.root.regions[0].blocks[0].operations = + gpu_module.regions[0].blocks[0].operations.clone(); + unwrapped.root.attributes.remove("gpu.container_module"); + unwrapped + .root + .properties + .insert("sym_name".into(), MlirAttribute::String(source_name.into())); + crate::render_module(&unwrapped).unwrap() +} + +/// Compare mapping goldens modulo SSA and block spelling, retaining each binding's identity. +#[cfg(test)] +pub(crate) fn normalize_mlir_identifiers(text: &str) -> String { + use std::collections::HashMap; + + let mut values = HashMap::new(); + let mut blocks = HashMap::new(); + let mut normalized = String::new(); + let mut characters = text.chars().peekable(); + let mut quoted = false; + let mut escaped = false; + while let Some(character) = characters.next() { + if quoted { + normalized.push(character); + if escaped { + escaped = false; + } else if character == '\\' { + escaped = true; + } else if character == '"' { + quoted = false; + } + } else if character == '"' { + quoted = true; + normalized.push(character); + } else if matches!(character, '%' | '^') { + let mut name = String::new(); + while let Some(next) = characters.peek().copied() { + if !next.is_ascii_alphanumeric() && !matches!(next, '_' | '.' | '$' | '-') { + break; + } + name.push(characters.next().unwrap()); + } + if name.is_empty() { + normalized.push(character); + continue; + } + let names = if character == '%' { + &mut values + } else { + &mut blocks + }; + let next = names.len(); + let index = names.entry(name).or_insert(next); + normalized.push(character); + normalized.push_str(&index.to_string()); + } else { + normalized.push(character); + } + } + normalized +} + +#[cfg(test)] +mod tests { + //! Profile manifest, architecture rejection, and immutable-source verification checks. + + use super::*; + use pliron::identifier::Identifier; + + /// Alpha-renaming must not erase changed operands, block references, or string contents. + #[test] + fn golden_normalization_preserves_binding_identity_and_strings() { + let original = "%x = op(%a, %a) [^x, ^next] {note = \"%x ^next\"}"; + let renamed = "%r = op(%input, %input) [^entry, ^exit] {note = \"%x ^next\"}"; + assert_eq!( + normalize_mlir_identifiers(original), + normalize_mlir_identifiers(renamed) + ); + for changed in [ + "%r = op(%input, %other) [^entry, ^exit] {note = \"%x ^next\"}", + "%r = op(%input, %input) [^entry, ^entry] {note = \"%x ^next\"}", + "%r = op(%input, %input) [^entry, ^exit] {note = \"%r ^exit\"}", + ] { + assert_ne!( + normalize_mlir_identifiers(original), + normalize_mlir_identifiers(changed) + ); + } + } + + /// Check manifest JSON records the pinned consumer/LLVM contract and the registry is sealed. + #[test] + fn profile_is_exact_and_machine_readable() { + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let json = profile.manifest().to_pretty_json().unwrap(); + + assert!(json.contains(CutlassFullCuteMlir22::CONSUMER_REVISION)); + assert!(json.contains(CutlassFullCuteMlir22::PLIRON_REVISION)); + assert!(json.contains(crate::INTEROP_REVISION)); + assert!(json.contains("\"llvm_major\": 22")); + assert!(profile.build_registry().unwrap().is_sealed()); + } + + /// Reject a target without a reviewed CUTLASS mapping profile. + #[test] + fn profile_rejects_an_unproved_architecture() { + assert!(matches!( + CutlassFullCuteMlir22::new("sm_90a"), + Err(ProfileError::UnsupportedArchitecture(arch)) if arch == "sm_90a" + )); + } + + /// Check whole-module verification accepts immutable access without rewriting source IR. + #[test] + fn profile_verifies_the_live_module_through_an_immutable_context() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("immutable_verify").unwrap()); + + let immutable_ctx = &ctx; + CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .verify_source_module(immutable_ctx, &module) + .unwrap(); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/ast.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/ast.rs new file mode 100644 index 0000000000..8b181f5c19 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/ast.rs @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::collections::BTreeMap; + +/// A deterministic SSA value number in the target module. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MlirValueId(pub u64); + +/// A deterministic target block number. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MlirBlockId(pub u64); + +/// Builtin floating-point types understood by the generic renderer. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum MlirFloatType { + F16, + BF16, + F32, + F64, + /// A target-defined builtin spelling such as `f8E4M3FN`. + Other(String), +} + +/// A target MLIR type. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum MlirType { + Index, + Integer(u32), + Float(MlirFloatType), + Vector { + shape: Vec, + element: Box, + }, + Function { + inputs: Vec, + results: Vec, + }, + Tuple(Vec), + /// A dialect type, including the leading `!`, held as one checked token. + Dialect(String), +} + +impl MlirType { + /// Create a dialect type while rejecting text that could escape its token. + pub fn dialect(spelling: impl Into) -> Result { + let spelling = spelling.into(); + if !spelling.starts_with('!') || has_forbidden_text(&spelling) { + return Err(format!( + "dialect type must start with '!' and stay on one line: {spelling:?}" + )); + } + Ok(Self::Dialect(spelling)) + } +} + +/// A target MLIR attribute. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MlirAttribute { + Unit, + Bool(bool), + Integer { + value: i128, + ty: MlirType, + }, + Float { + value: String, + ty: MlirType, + }, + String(String), + Type(MlirType), + Array(Vec), + /// MLIR's compact `array` form used by operation properties + /// such as `operandSegmentSizes`. + DenseI32Array(Vec), + /// MLIR's compact `array` form used by aggregate index + /// properties such as LLVM dialect `position`. + DenseI64Array(Vec), + Dictionary(BTreeMap), + SymbolRef(String), + /// A dialect attribute, including the leading `#`, held as one checked token. + Dialect(String), +} + +impl MlirAttribute { + /// Create a dialect attribute while rejecting text that could escape it. + pub fn dialect(spelling: impl Into) -> Result { + let spelling = spelling.into(); + if !spelling.starts_with('#') || has_forbidden_text(&spelling) { + return Err(format!( + "dialect attribute must start with '#' and stay on one line: {spelling:?}" + )); + } + Ok(Self::Dialect(spelling)) + } +} + +/// A source location carried into MLIR. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MlirLocation { + Unknown, + FileLineCol { + file: String, + line: u32, + column: u32, + }, + Fused { + metadata: Option>, + locations: Vec, + }, + Named { + name: String, + child: Box, + }, + CallSite { + callee: Box, + caller: Box, + }, +} + +/// One typed use of an SSA value. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirValueUse { + pub id: MlirValueId, + pub ty: MlirType, +} + +/// One typed SSA result. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirResult { + pub id: MlirValueId, + pub ty: MlirType, +} + +/// One operation in generic MLIR form. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirOperation { + pub name: String, + pub results: Vec, + pub operands: Vec, + pub successors: Vec, + /// Inherent operation properties rendered as `<{...}>`. + /// + /// Properties are not ordinary discardable attributes in modern MLIR. + /// Keeping them separate prevents a mapping from accidentally emitting a + /// structurally valid-looking operation whose verifier sees defaults. + pub properties: BTreeMap, + pub regions: Vec, + pub attributes: BTreeMap, + pub location: MlirLocation, +} + +impl MlirOperation { + /// Start an operation with no operands, results, regions, or attributes. + pub fn new(name: impl Into) -> Result { + let name = name.into(); + if name.is_empty() || has_forbidden_text(&name) || name.contains('"') { + return Err(format!("invalid MLIR operation name: {name:?}")); + } + Ok(Self { + name, + results: vec![], + operands: vec![], + successors: vec![], + properties: BTreeMap::new(), + regions: vec![], + attributes: BTreeMap::new(), + location: MlirLocation::Unknown, + }) + } +} + +/// A typed argument at the start of an MLIR block. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirBlockArgument { + pub id: MlirValueId, + pub ty: MlirType, + pub location: MlirLocation, +} + +/// A target MLIR block. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirBlock { + pub id: MlirBlockId, + pub arguments: Vec, + pub operations: Vec, +} + +/// A target MLIR region. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirRegion { + pub blocks: Vec, +} + +/// A translated module and the profile that defined its mappings. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MlirModule { + pub root: MlirOperation, + pub profile: String, +} + +fn has_forbidden_text(value: &str) -> bool { + value.chars().any(|ch| ch.is_control()) +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/error.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/error.rs new file mode 100644 index 0000000000..7ab8b384e5 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/error.rs @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::collections::BTreeSet; + +use thiserror::Error; + +/// Every source item that has no mapping in the selected profile. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MissingMappings { + pub operations: BTreeSet, + pub types: BTreeSet, + pub attributes: BTreeSet, +} + +impl MissingMappings { + pub fn is_empty(&self) -> bool { + self.operations.is_empty() && self.types.is_empty() && self.attributes.is_empty() + } +} + +impl std::fmt::Display for MissingMappings { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!( + f, + "the selected MLIR profile does not cover this Pliron module" + )?; + if !self.operations.is_empty() { + writeln!(f, " operations: {}", join(&self.operations))?; + } + if !self.types.is_empty() { + writeln!(f, " types: {}", join(&self.types))?; + } + if !self.attributes.is_empty() { + writeln!(f, " attributes: {}", join(&self.attributes))?; + } + Ok(()) + } +} + +fn join(values: &BTreeSet) -> String { + values.iter().cloned().collect::>().join(", ") +} + +/// A checked translation failure. +#[derive(Debug, Error)] +pub enum TranslationError { + #[error("duplicate {kind} mapping for {id}")] + DuplicateMapping { kind: &'static str, id: String }, + #[error("the translation registry is sealed")] + RegistrySealed, + #[error("{0}")] + MissingMappings(MissingMappings), + #[error("translation of {source_item} for profile {profile} failed: {message}")] + Mapping { + source_item: String, + profile: String, + message: String, + }, + #[error("translated operation {source_item} did not define source result {result:?}")] + MissingResult { + source_item: String, + result: crate::MlirValueId, + }, + #[error("translated operation {source_item} defined source result {result:?} more than once")] + DuplicateResult { + source_item: String, + result: crate::MlirValueId, + }, + #[error("internal translation state is missing {kind} {id}")] + MissingIndexedValue { kind: &'static str, id: String }, + #[error("invalid target MLIR: {0}")] + InvalidTarget(String), +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/mod.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/mod.rs new file mode 100644 index 0000000000..30066643cc --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/mod.rs @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit mapping plans and checked translation of the prepared source module. + +mod ast; +mod error; +mod registry; +mod translate; + +pub(crate) use translate::validate_target; + +pub use ast::{ + MlirAttribute, MlirBlock, MlirBlockArgument, MlirBlockId, MlirFloatType, MlirLocation, + MlirModule, MlirOperation, MlirRegion, MlirResult, MlirType, MlirValueId, MlirValueUse, +}; +pub use error::{MissingMappings, TranslationError}; +pub use registry::{ + AttributeTranslation, DropAttribute, FixedType, OneToOneOperation, OperationTranslation, + TranslationRegistry, TypeTranslation, +}; +pub use translate::{ + LocationMode, OperationInput, TranslationConfig, TranslationSession, translate_module, +}; diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/registry.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/registry.rs new file mode 100644 index 0000000000..02f2002d3f --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/registry.rs @@ -0,0 +1,253 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::collections::{HashMap, hash_map::Entry}; + +use pliron::{ + attribute::{AttrId, AttrObj, Attribute}, + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::{Type, TypeHandle, TypeId}, +}; + +use crate::TranslationSession; +use crate::{MlirAttribute, MlirOperation, MlirType, OperationInput, TranslationError}; + +/// Translate one source type. Implementations may recursively ask the registry +/// to translate contained types. +pub trait TypeTranslation: Send + Sync { + fn translate( + &self, + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + ) -> Result; +} + +/// Translate one source attribute. `None` is an explicit, registered drop. +pub trait AttributeTranslation: Send + Sync { + fn translate( + &self, + ctx: &Context, + source: &AttrObj, + registry: &TranslationRegistry, + ) -> Result, String>; +} + +/// Translate one source operation into zero, one, or several target ops. +pub trait OperationTranslation: Send + Sync { + fn translate( + &self, + ctx: &Context, + source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String>; +} + +/// Explicit mappings for one target MLIR profile. +#[derive(Default)] +pub struct TranslationRegistry { + operations: HashMap>, + types: HashMap>, + attributes: HashMap>, + sealed: bool, +} + +impl TranslationRegistry { + pub fn new() -> Self { + Self::default() + } + + /// Prevent every later registration. Profiles normally seal their + /// registry after all mapping packs have been composed. + pub fn seal(&mut self) { + self.sealed = true; + } + + pub fn is_sealed(&self) -> bool { + self.sealed + } + + /// Translate a nested source type from inside another mapping recipe. + /// + /// Container mappings use this for their element, field, and function + /// signature types. Missing mappings remain explicit errors; there is no + /// name-based fallback. + pub fn translate_type(&self, ctx: &Context, source: TypeHandle) -> Result { + let id = source.deref(ctx).get_type_id(); + self.type_translation(&id) + .ok_or_else(|| format!("missing nested type mapping for `{id}`"))? + .translate(ctx, source, self) + } + + pub fn register_operation_id( + &mut self, + id: OpId, + translation: impl OperationTranslation + 'static, + ) -> Result<(), TranslationError> { + if self.sealed { + return Err(TranslationError::RegistrySealed); + } + match self.operations.entry(id.clone()) { + Entry::Vacant(entry) => { + entry.insert(Box::new(translation)); + Ok(()) + } + Entry::Occupied(_) => Err(TranslationError::DuplicateMapping { + kind: "operation", + id: id.to_string(), + }), + } + } + + pub fn register_operation( + &mut self, + translation: impl OperationTranslation + 'static, + ) -> Result<(), TranslationError> { + self.register_operation_id(O::get_opid_static(), translation) + } + + pub fn register_type_id( + &mut self, + id: TypeId, + translation: impl TypeTranslation + 'static, + ) -> Result<(), TranslationError> { + if self.sealed { + return Err(TranslationError::RegistrySealed); + } + match self.types.entry(id.clone()) { + Entry::Vacant(entry) => { + entry.insert(Box::new(translation)); + Ok(()) + } + Entry::Occupied(_) => Err(TranslationError::DuplicateMapping { + kind: "type", + id: id.to_string(), + }), + } + } + + pub fn register_type( + &mut self, + translation: impl TypeTranslation + 'static, + ) -> Result<(), TranslationError> { + self.register_type_id(T::get_type_id_static(), translation) + } + + pub fn register_attribute_id( + &mut self, + id: AttrId, + translation: impl AttributeTranslation + 'static, + ) -> Result<(), TranslationError> { + if self.sealed { + return Err(TranslationError::RegistrySealed); + } + match self.attributes.entry(id.clone()) { + Entry::Vacant(entry) => { + entry.insert(Box::new(translation)); + Ok(()) + } + Entry::Occupied(_) => Err(TranslationError::DuplicateMapping { + kind: "attribute", + id: id.to_string(), + }), + } + } + + pub fn register_attribute( + &mut self, + translation: impl AttributeTranslation + 'static, + ) -> Result<(), TranslationError> { + self.register_attribute_id(A::get_attr_id_static(), translation) + } + + pub(crate) fn operation(&self, id: &OpId) -> Option<&dyn OperationTranslation> { + self.operations.get(id).map(Box::as_ref) + } + + pub(crate) fn type_translation(&self, id: &TypeId) -> Option<&dyn TypeTranslation> { + self.types.get(id).map(Box::as_ref) + } + + pub(crate) fn attribute(&self, id: &AttrId) -> Option<&dyn AttributeTranslation> { + self.attributes.get(id).map(Box::as_ref) + } + + pub(crate) fn has_operation(&self, id: &OpId) -> bool { + self.operations.contains_key(id) + } + + pub(crate) fn has_type(&self, id: &TypeId) -> bool { + self.types.contains_key(id) + } + + pub(crate) fn has_attribute(&self, id: &AttrId) -> bool { + self.attributes.contains_key(id) + } +} + +/// A fixed target type, useful for builtin scalar mappings. +#[derive(Clone)] +pub struct FixedType(pub MlirType); + +impl TypeTranslation for FixedType { + fn translate( + &self, + _ctx: &Context, + _source: TypeHandle, + _registry: &TranslationRegistry, + ) -> Result { + Ok(self.0.clone()) + } +} + +/// An explicitly dropped source attribute. +pub struct DropAttribute; + +impl AttributeTranslation for DropAttribute { + fn translate( + &self, + _ctx: &Context, + _source: &AttrObj, + _registry: &TranslationRegistry, + ) -> Result, String> { + Ok(None) + } +} + +/// Rename an operation while preserving its typed operands, results, regions, +/// successors, attributes, and location. +pub struct OneToOneOperation { + target_name: String, +} + +impl OneToOneOperation { + pub fn new(target_name: impl Into) -> Self { + Self { + target_name: target_name.into(), + } + } +} + +impl OperationTranslation for OneToOneOperation { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + _session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mut operation = MlirOperation::new(&self.target_name)?; + operation.results = input.results; + operation.operands = input.operands; + operation.successors = input.successors; + operation.regions = input.regions; + operation.attributes = input.attributes; + operation.location = input.location; + Ok(vec![operation]) + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/translate.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/translate.rs new file mode 100644 index 0000000000..fe8b8e41b0 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/translation/translate.rs @@ -0,0 +1,842 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use std::collections::{BTreeMap, HashMap, HashSet}; + +use pliron::{ + attribute::AttrObj, + basic_block::BasicBlock, + builtin::ops::ModuleOp, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + location::{Location, Source}, + op::Op, + operation::{Operation, verify_operation}, + region::Region, + r#type::{TypeHandle, Typed}, + uniqued_any, + value::Value, +}; + +use crate::{ + MissingMappings, MlirAttribute, MlirBlock, MlirBlockArgument, MlirBlockId, MlirLocation, + MlirModule, MlirOperation, MlirRegion, MlirResult, MlirType, MlirValueId, MlirValueUse, + TranslationError, TranslationRegistry, +}; + +/// How much source location data to keep. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum LocationMode { + Off, + FileLineCol, + #[default] + Full, +} + +/// Stable identity and rendering choices for one export. +#[derive(Clone, Debug)] +pub struct TranslationConfig { + pub profile: String, + pub locations: LocationMode, +} + +impl TranslationConfig { + pub fn new(profile: impl Into) -> Self { + Self { + profile: profile.into(), + locations: LocationMode::Full, + } + } +} + +/// Fully translated inputs supplied to one operation recipe. +#[derive(Clone, Debug)] +pub struct OperationInput { + pub results: Vec, + pub operands: Vec, + pub successors: Vec, + pub regions: Vec, + pub attributes: BTreeMap, + pub location: MlirLocation, +} + +/// Mutable state shared by operation recipes during one export. +pub struct TranslationSession<'a> { + ctx: &'a Context, + registry: &'a TranslationRegistry, + config: &'a TranslationConfig, + values: HashMap, + blocks: HashMap, MlirBlockId>, + next_value: u64, + next_block: u64, +} + +impl<'a> TranslationSession<'a> { + pub fn profile(&self) -> &str { + &self.config.profile + } + + pub fn fresh_value(&mut self) -> MlirValueId { + let value = MlirValueId(self.next_value); + self.next_value += 1; + value + } + + /// Allocate a deterministic block ID for a region created by a mapping. + /// + /// Source blocks are indexed before translation. Mapping-created nested + /// regions therefore receive IDs strictly after every source block, which + /// prevents collisions without coupling a dialect recipe to traversal + /// order or renderer internals. + pub fn fresh_block(&mut self) -> MlirBlockId { + let block = MlirBlockId(self.next_block); + self.next_block += 1; + block + } + + pub fn translate_type(&self, source: TypeHandle) -> Result { + translate_type(self.ctx, self.registry, self.config, source) + } + + pub fn translate_attribute( + &self, + source: &AttrObj, + ) -> Result, TranslationError> { + translate_attribute(self.ctx, self.registry, self.config, source) + } + + /// Return the target SSA use assigned to any indexed source value. + /// + /// Most recipes consume only their operation's direct operands, which are + /// already present in [`OperationInput`]. A semantic recipe may also need + /// a value from a module-validated producer story—for example the original + /// tensor base, length, or tile coordinate. This lookup keeps that need + /// typed and deterministic instead of making dialect packs reconstruct + /// target names or reach into exporter internals. + pub fn target_value_use(&self, source: Value) -> Result { + let id = self.values.get(&source).copied().ok_or_else(|| { + TranslationError::MissingIndexedValue { + kind: "value", + id: format!("{source:?}"), + } + })?; + Ok(MlirValueUse { + id, + ty: self.translate_type(source.get_type(self.ctx))?, + }) + } +} + +/// Verify, preflight, translate, and build one target module. +pub fn translate_module( + ctx: &Context, + module: &ModuleOp, + registry: &TranslationRegistry, + config: &TranslationConfig, +) -> Result { + if config.profile.trim().is_empty() { + return Err(TranslationError::InvalidTarget( + "translation profile name must not be empty".into(), + )); + } + verify_operation(module.get_operation(), ctx).map_err(|error| TranslationError::Mapping { + source_item: "builtin.module".into(), + profile: config.profile.clone(), + message: format!("source verification failed: {error}"), + })?; + + let missing = collect_missing(ctx, module.get_operation(), registry); + if !missing.is_empty() { + return Err(TranslationError::MissingMappings(missing)); + } + + let mut session = TranslationSession { + ctx, + registry, + config, + values: HashMap::new(), + blocks: HashMap::new(), + next_value: 0, + next_block: 0, + }; + let mut next_block = 0; + index_operation(module.get_operation(), ctx, &mut session, &mut next_block); + session.next_block = next_block; + let mut translated = translate_operation(module.get_operation(), &mut session)?; + if translated.len() != 1 { + return Err(TranslationError::InvalidTarget(format!( + "builtin.module mapping must emit exactly one root operation, got {}", + translated.len() + ))); + } + let module = MlirModule { + root: translated.remove(0), + profile: config.profile.clone(), + }; + validate_target(&module)?; + Ok(module) +} + +fn collect_missing( + ctx: &Context, + root: Ptr, + registry: &TranslationRegistry, +) -> MissingMappings { + let mut missing = MissingMappings::default(); + scan_operation(ctx, root, registry, &mut missing); + missing +} + +fn scan_operation( + ctx: &Context, + source: Ptr, + registry: &TranslationRegistry, + missing: &mut MissingMappings, +) { + let id = Operation::get_opid(source, ctx); + if !registry.has_operation(&id) { + missing.operations.insert(id.to_string()); + } + let operation = source.deref(ctx); + for value in operation.operands().chain(operation.results()) { + scan_type(ctx, value.get_type(ctx), registry, missing); + } + scan_attributes(&operation.attributes.0, registry, missing); + scan_location(operation.loc(), registry, missing); + for region in operation.regions() { + for block in region.deref(ctx).iter(ctx) { + let block_ref = block.deref(ctx); + for argument in block_ref.arguments() { + scan_type(ctx, argument.get_type(ctx), registry, missing); + } + scan_attributes(&block_ref.attributes.0, registry, missing); + scan_location(block_ref.loc(), registry, missing); + for child in block_ref.iter(ctx) { + scan_operation(ctx, child, registry, missing); + } + } + } +} + +fn scan_location( + location: &Location, + registry: &TranslationRegistry, + missing: &mut MissingMappings, +) { + match location { + Location::Fused { + metadata, + locations, + } => { + if let Some(metadata) = metadata { + let id = metadata.get_attr_id(); + if !registry.has_attribute(&id) { + missing.attributes.insert(id.to_string()); + } + } + for location in locations { + scan_location(location, registry, missing); + } + } + Location::Named { child_loc, .. } => scan_location(child_loc, registry, missing), + Location::CallSite { callee, caller } => { + scan_location(callee, registry, missing); + scan_location(caller, registry, missing); + } + Location::SrcPos { .. } | Location::Unknown => {} + } +} + +fn scan_type( + ctx: &Context, + source: TypeHandle, + registry: &TranslationRegistry, + missing: &mut MissingMappings, +) { + let id = source.deref(ctx).get_type_id(); + if !registry.has_type(&id) { + missing.types.insert(id.to_string()); + } +} + +fn scan_attributes( + attributes: &pliron::attribute::AttributeDictContainer, + registry: &TranslationRegistry, + missing: &mut MissingMappings, +) { + for attribute in attributes.values() { + let id = attribute.get_attr_id(); + if !registry.has_attribute(&id) { + missing.attributes.insert(id.to_string()); + } + } +} + +fn index_operation( + source: Ptr, + ctx: &Context, + session: &mut TranslationSession<'_>, + next_block: &mut u64, +) { + for result in source.deref(ctx).results() { + let id = session.fresh_value(); + session.values.insert(result, id); + } + for region in source.deref(ctx).regions() { + let blocks = region.deref(ctx).iter(ctx).collect::>(); + for block in &blocks { + session.blocks.insert(*block, MlirBlockId(*next_block)); + *next_block += 1; + for argument in block.deref(ctx).arguments() { + let id = session.fresh_value(); + session.values.insert(argument, id); + } + } + for block in blocks { + for child in block.deref(ctx).iter(ctx) { + index_operation(child, ctx, session, next_block); + } + } + } +} + +fn translate_operation( + source: Ptr, + session: &mut TranslationSession<'_>, +) -> Result, TranslationError> { + let ctx = session.ctx; + let source_id = Operation::get_opid(source, ctx); + let operation = source.deref(ctx); + let mut regions = Vec::with_capacity(operation.num_regions()); + for region in operation.regions() { + regions.push(translate_region(region, session)?); + } + + let mut operands = Vec::with_capacity(operation.get_num_operands()); + for value in operation.operands() { + let id = session.values.get(&value).copied().ok_or_else(|| { + TranslationError::MissingIndexedValue { + kind: "value", + id: format!("{value:?}"), + } + })?; + operands.push(MlirValueUse { + id, + ty: session.translate_type(value.get_type(ctx))?, + }); + } + + let mut results = Vec::with_capacity(operation.get_num_results()); + for value in operation.results() { + let id = session.values.get(&value).copied().ok_or_else(|| { + TranslationError::MissingIndexedValue { + kind: "value", + id: format!("{value:?}"), + } + })?; + results.push(MlirResult { + id, + ty: session.translate_type(value.get_type(ctx))?, + }); + } + + let mut successors = Vec::with_capacity(operation.get_num_successors()); + for successor in operation.successors() { + successors.push(session.blocks.get(&successor).copied().ok_or_else(|| { + TranslationError::MissingIndexedValue { + kind: "block", + id: format!("{successor:?}"), + } + })?); + } + + let attributes = translate_attributes( + ctx, + session.registry, + session.config, + &operation.attributes.0, + )?; + let location = translate_location(ctx, session, operation.loc())?; + drop(operation); + + let input = OperationInput { + results: results.clone(), + operands, + successors, + regions, + attributes, + location, + }; + let translation = session.registry.operation(&source_id).ok_or_else(|| { + let mut missing = MissingMappings::default(); + missing.operations.insert(source_id.to_string()); + TranslationError::MissingMappings(missing) + })?; + let translated = translation + .translate(ctx, source, input, session) + .map_err(|message| TranslationError::Mapping { + source_item: source_id.to_string(), + profile: session.config.profile.clone(), + message, + })?; + check_source_results(&source_id.to_string(), &results, &translated)?; + Ok(translated) +} + +fn translate_region( + source: Ptr, + session: &mut TranslationSession<'_>, +) -> Result { + let ctx = session.ctx; + let blocks = source.deref(ctx).iter(ctx).collect::>(); + let mut translated_blocks = Vec::with_capacity(blocks.len()); + for block in blocks { + let block_ref = block.deref(ctx); + if !block_ref.attributes.0.is_empty() { + let translated = translate_attributes( + ctx, + session.registry, + session.config, + &block_ref.attributes.0, + )?; + if !translated.is_empty() { + return Err(TranslationError::InvalidTarget(format!( + "MLIR blocks cannot carry the translated attributes on {block:?}" + ))); + } + } + let id = session.blocks.get(&block).copied().ok_or_else(|| { + TranslationError::MissingIndexedValue { + kind: "block", + id: format!("{block:?}"), + } + })?; + let mut arguments = Vec::with_capacity(block_ref.get_num_arguments()); + for argument in block_ref.arguments() { + arguments.push(MlirBlockArgument { + id: session.values[&argument], + ty: session.translate_type(argument.get_type(ctx))?, + location: translate_location(ctx, session, block_ref.loc())?, + }); + } + let operations = block_ref.iter(ctx).collect::>(); + drop(block_ref); + let mut translated_operations = vec![]; + for operation in operations { + translated_operations.extend(translate_operation(operation, session)?); + } + translated_blocks.push(MlirBlock { + id, + arguments, + operations: translated_operations, + }); + } + Ok(MlirRegion { + blocks: translated_blocks, + }) +} + +fn translate_type( + ctx: &Context, + registry: &TranslationRegistry, + config: &TranslationConfig, + source: TypeHandle, +) -> Result { + let id = source.deref(ctx).get_type_id(); + let translation = registry.type_translation(&id).ok_or_else(|| { + let mut missing = MissingMappings::default(); + missing.types.insert(id.to_string()); + TranslationError::MissingMappings(missing) + })?; + translation + .translate(ctx, source, registry) + .map_err(|message| TranslationError::Mapping { + source_item: id.to_string(), + profile: config.profile.clone(), + message, + }) +} + +fn translate_attribute( + ctx: &Context, + registry: &TranslationRegistry, + config: &TranslationConfig, + source: &AttrObj, +) -> Result, TranslationError> { + let id = source.get_attr_id(); + let translation = registry.attribute(&id).ok_or_else(|| { + let mut missing = MissingMappings::default(); + missing.attributes.insert(id.to_string()); + TranslationError::MissingMappings(missing) + })?; + translation + .translate(ctx, source, registry) + .map_err(|message| TranslationError::Mapping { + source_item: id.to_string(), + profile: config.profile.clone(), + message, + }) +} + +fn translate_attributes( + ctx: &Context, + registry: &TranslationRegistry, + config: &TranslationConfig, + source: &pliron::attribute::AttributeDictContainer, +) -> Result, TranslationError> { + let mut translated = BTreeMap::new(); + for (name, attribute) in source { + if let Some(attribute) = translate_attribute(ctx, registry, config, attribute)? { + translated.insert(name.to_string(), attribute); + } + } + Ok(translated) +} + +fn translate_location( + ctx: &Context, + session: &TranslationSession<'_>, + source: &Location, +) -> Result { + match session.config.locations { + LocationMode::Off => Ok(MlirLocation::Unknown), + LocationMode::FileLineCol => Ok(first_file_location(ctx, source)), + LocationMode::Full => full_location(ctx, session, source), + } +} + +fn first_file_location(ctx: &Context, source: &Location) -> MlirLocation { + match source { + Location::SrcPos { src, pos } => { + source_position(ctx, *src, pos.line.max(0) as u32, pos.column.max(0) as u32) + } + Location::Fused { locations, .. } => locations + .iter() + .map(|location| first_file_location(ctx, location)) + .find(|location| *location != MlirLocation::Unknown) + .unwrap_or(MlirLocation::Unknown), + Location::Named { child_loc, .. } => first_file_location(ctx, child_loc), + Location::CallSite { caller, callee } => { + let caller = first_file_location(ctx, caller); + if caller != MlirLocation::Unknown { + caller + } else { + first_file_location(ctx, callee) + } + } + Location::Unknown => MlirLocation::Unknown, + } +} + +fn full_location( + ctx: &Context, + session: &TranslationSession<'_>, + source: &Location, +) -> Result { + Ok(match source { + Location::SrcPos { src, pos } => { + source_position(ctx, *src, pos.line.max(0) as u32, pos.column.max(0) as u32) + } + Location::Fused { + metadata, + locations, + } => MlirLocation::Fused { + metadata: match metadata { + Some(metadata) => session.translate_attribute(metadata)?.map(Box::new), + None => None, + }, + locations: locations + .iter() + .map(|location| full_location(ctx, session, location)) + .collect::>()?, + }, + Location::Named { name, child_loc } => MlirLocation::Named { + name: name.clone(), + child: Box::new(full_location(ctx, session, child_loc)?), + }, + Location::CallSite { callee, caller } => MlirLocation::CallSite { + callee: Box::new(full_location(ctx, session, callee)?), + caller: Box::new(full_location(ctx, session, caller)?), + }, + Location::Unknown => MlirLocation::Unknown, + }) +} + +fn source_position(ctx: &Context, source: Source, line: u32, column: u32) -> MlirLocation { + match source { + Source::File(path) => MlirLocation::FileLineCol { + file: uniqued_any::get(ctx, path).display().to_string(), + line, + column, + }, + Source::InMemory => MlirLocation::Unknown, + } +} + +fn check_source_results( + source: &str, + expected: &[MlirResult], + translated: &[MlirOperation], +) -> Result<(), TranslationError> { + let mut definitions = HashMap::>::new(); + for operation in translated { + for result in &operation.results { + definitions + .entry(result.id) + .or_default() + .push(result.ty.clone()); + } + } + for result in expected { + match definitions.get(&result.id).map_or(0, Vec::len) { + 0 => { + return Err(TranslationError::MissingResult { + source_item: source.into(), + result: result.id, + }); + } + 1 => { + let actual = &definitions[&result.id][0]; + if actual != &result.ty { + return Err(TranslationError::InvalidTarget(format!( + "mapping for `{source}` defines {:?} as {actual:?}, expected {:?}", + result.id, result.ty + ))); + } + } + _ => { + return Err(TranslationError::DuplicateResult { + source_item: source.into(), + result: result.id, + }); + } + } + } + Ok(()) +} + +pub(crate) fn validate_target(module: &MlirModule) -> Result<(), TranslationError> { + let mut values = HashMap::::new(); + let mut blocks = HashSet::::new(); + collect_target_definitions(&module.root, &mut values, &mut blocks)?; + check_target_uses(&module.root, &values, &blocks) +} + +fn collect_target_definitions( + operation: &MlirOperation, + values: &mut HashMap, + blocks: &mut HashSet, +) -> Result<(), TranslationError> { + for result in &operation.results { + if values.insert(result.id, result.ty.clone()).is_some() { + return Err(TranslationError::InvalidTarget(format!( + "SSA value {:?} is defined more than once", + result.id + ))); + } + } + for region in &operation.regions { + for block in ®ion.blocks { + if !blocks.insert(block.id) { + return Err(TranslationError::InvalidTarget(format!( + "block {:?} is defined more than once", + block.id + ))); + } + for argument in &block.arguments { + if values.insert(argument.id, argument.ty.clone()).is_some() { + return Err(TranslationError::InvalidTarget(format!( + "SSA value {:?} is defined more than once", + argument.id + ))); + } + } + for child in &block.operations { + collect_target_definitions(child, values, blocks)?; + } + } + } + Ok(()) +} + +fn check_target_uses( + operation: &MlirOperation, + values: &HashMap, + blocks: &HashSet, +) -> Result<(), TranslationError> { + for operand in &operation.operands { + let Some(definition_type) = values.get(&operand.id) else { + return Err(TranslationError::InvalidTarget(format!( + "operand {:?} has no definition", + operand.id + ))); + }; + if definition_type != &operand.ty { + return Err(TranslationError::InvalidTarget(format!( + "operand {:?} uses type {:?}, but its definition has type {:?}", + operand.id, operand.ty, definition_type + ))); + } + } + for successor in &operation.successors { + if !blocks.contains(successor) { + return Err(TranslationError::InvalidTarget(format!( + "successor {:?} has no block definition", + successor + ))); + } + } + for region in &operation.regions { + for block in ®ion.blocks { + for child in &block.operations { + check_target_uses(child, values, blocks)?; + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use pliron::{ + builtin::{ + attributes::{IdentifierAttr, StringAttr}, + ops::ModuleOp, + types::IntegerType, + }, + context::Context, + identifier::Identifier, + op::Op, + }; + + use crate::{ + OperationTranslation, + translation::registry::{DropAttribute, FixedType, OneToOneOperation}, + }; + + use super::*; + + struct MappingCreatedRegion; + + impl OperationTranslation for MappingCreatedRegion { + fn translate( + &self, + _ctx: &Context, + _source: Ptr, + input: OperationInput, + session: &mut TranslationSession<'_>, + ) -> Result, String> { + let mut root = MlirOperation::new("builtin.module")?; + root.regions = input.regions; + root.attributes = input.attributes; + root.location = input.location.clone(); + let mut generated = MlirOperation::new("test.mapping_region")?; + generated.regions.push(MlirRegion { + blocks: vec![MlirBlock { + id: session.fresh_block(), + arguments: vec![], + operations: vec![], + }], + }); + generated.location = input.location; + root.regions[0].blocks[0].operations.push(generated); + Ok(vec![root]) + } + } + + #[test] + fn preflight_collects_every_missing_kind() { + let mut ctx = Context::default(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("missing").unwrap()); + let error = translate_module( + &ctx, + &module, + &TranslationRegistry::new(), + &TranslationConfig::new("test"), + ) + .unwrap_err(); + let TranslationError::MissingMappings(missing) = error else { + panic!("expected missing mapping error"); + }; + assert!(missing.operations.contains("builtin.module")); + assert!( + missing + .attributes + .iter() + .any(|id| id.contains("identifier")) + ); + } + + #[test] + fn module_round_trips_through_generic_ast() { + let mut ctx = Context::default(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("empty").unwrap()); + let mut registry = TranslationRegistry::new(); + registry + .register_operation::(OneToOneOperation::new("builtin.module")) + .unwrap(); + registry + .register_attribute::(DropAttribute) + .unwrap(); + // Some Pliron revisions store symbol names as StringAttr instead. + let _ = registry.register_attribute::(DropAttribute); + registry + .register_type::(FixedType(MlirType::Integer(32))) + .unwrap(); + let translated = + translate_module(&ctx, &module, ®istry, &TranslationConfig::new("test")).unwrap(); + let text = crate::render_module(&translated).unwrap(); + assert!(text.starts_with("\"builtin.module\"() ({")); + assert!(text.trim_end().ends_with(": () -> ()")); + } + + #[test] + fn mapping_created_block_ids_are_non_colliding_and_deterministic() { + let mut ctx = Context::default(); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("nested").unwrap()); + let mut registry = TranslationRegistry::new(); + registry + .register_operation::(MappingCreatedRegion) + .unwrap(); + registry + .register_attribute::(DropAttribute) + .unwrap(); + let _ = registry.register_attribute::(DropAttribute); + + let first = + translate_module(&ctx, &module, ®istry, &TranslationConfig::new("test")).unwrap(); + let second = + translate_module(&ctx, &module, ®istry, &TranslationConfig::new("test")).unwrap(); + let source_block = first.root.regions[0].blocks[0].id; + let generated_block = first.root.regions[0].blocks[0].operations[0].regions[0].blocks[0].id; + assert_ne!(source_block, generated_block); + assert_eq!(source_block, MlirBlockId(0)); + assert_eq!(generated_block, MlirBlockId(1)); + assert_eq!( + crate::render_module(&first).unwrap(), + crate::render_module(&second).unwrap() + ); + } + + #[test] + fn registry_rejects_duplicates_and_late_changes() { + let mut registry = TranslationRegistry::new(); + registry + .register_operation::(OneToOneOperation::new("builtin.module")) + .unwrap(); + assert!(matches!( + registry.register_operation::(OneToOneOperation::new("other.module")), + Err(TranslationError::DuplicateMapping { .. }) + )); + assert!(registry.has_operation(&ModuleOp::get_opid_static())); + + registry.seal(); + assert!(registry.is_sealed()); + assert!(matches!( + registry.register_attribute::(DropAttribute), + Err(TranslationError::RegistrySealed) + )); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/address_contract.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/address_contract.rs new file mode 100644 index 0000000000..7f1fd7bcc9 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/address_contract.rs @@ -0,0 +1,155 @@ +//! Named shared offsets versus Rust generic pointer-address semantics. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use cuda_oxide_mlir_export::{CutlassPrimitivesMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirOperation, MlirType}; +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{MirCastOp, MirFuncOp, MirReturnOp}, + types::MirPtrType, +}; +use dialect_nvvm::ops::{CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +/// Build pointer-address operations for a selected source space and export through the primitives profile. +fn translate( + space: u32, + width: u32, + kind: Option, + cluster: bool, +) -> Vec { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, width, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, word, true, space).into(); + let module = ModuleOp::new(&mut ctx, "addresses".try_into().unwrap()); + let fty = FunctionType::get(&ctx, vec![ptr], vec![word]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = MirFuncOp::new(&mut ctx, f, TypeAttr::new(fty.into())); + func.set_symbol_name(&mut ctx, "address".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let block = BasicBlock::new(&mut ctx, None, vec![ptr]); + block.insert_at_back(f.deref(&ctx).get_region(0), &ctx); + let arg = block.deref(&ctx).get_argument(0); + let info = if kind.is_some() { + MirCastOp::get_concrete_op_info() + } else if cluster { + CvtaGenericToClusterSharedOffsetOp::get_concrete_op_info() + } else { + CvtaGenericToSharedOffsetOp::get_concrete_op_info() + }; + let cast = Operation::new(&mut ctx, info, vec![word], vec![arg], vec![], 0); + if let Some(kind) = kind { + MirCastOp::new(cast).set_attr_cast_kind(&ctx, kind); + } + cast.insert_at_back(block, &ctx); + let value = cast.deref(&ctx).get_result(0); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![value], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + let target = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + target.root.regions[0].blocks[0].operations[0].regions[0].blocks[0].operations[0].regions[0] + .blocks[0] + .operations + .clone() +} + +/// Check explicit shared offsets retain named spaces while Rust address casts use generic addresses. +#[test] +fn explicit_shared_offsets_preserve_cluster_rank_while_rust_addresses_genericize() { + for space in [0, 3, 7] { + for width in [32, 64] { + for kind in [ + None, + Some(MirCastKindAttr::PointerExposeAddress), + Some(MirCastKindAttr::Transmute), + ] { + if kind == Some(MirCastKindAttr::Transmute) && width != 64 { + continue; + } + let expose = kind.is_some(); + let ops = translate(space, width, kind, false); + let casts: Vec<_> = ops + .iter() + .filter(|o| o.name == "llvm.addrspacecast") + .collect(); + let ptrtoint = ops.iter().find(|o| o.name == "llvm.ptrtoint").unwrap(); + let needs_cast = if expose { space != 0 } else { space == 0 }; + assert_eq!(casts.len(), usize::from(needs_cast)); + let result_space = if expose { + 0 + } else if space == 0 { + 3 + } else { + space + }; + let spelling = if result_space == 0 { + "!llvm.ptr".into() + } else { + format!("!llvm.ptr<{result_space}>") + }; + assert_eq!( + ptrtoint.operands[0].ty, + MlirType::dialect(spelling).unwrap() + ); + assert_eq!(ptrtoint.results[0].ty, MlirType::Integer(width)); + if needs_cast { + assert_eq!(ptrtoint.operands[0].id, casts[0].results[0].id); + } + } + } + } +} + +/// Check generic pointers enter cluster shared space before their integer offset is formed. +#[test] +fn explicit_cluster_offsets_target_as7_even_from_generic_pointers() { + for space in [0, 3, 7] { + for width in [32, 64] { + let ops = translate(space, width, None, true); + let casts: Vec<_> = ops + .iter() + .filter(|o| o.name == "llvm.addrspacecast") + .collect(); + let int = ops.iter().find(|o| o.name == "llvm.ptrtoint").unwrap(); + assert_eq!(casts.len(), usize::from(space != 7)); + assert_eq!( + int.operands[0].ty, + MlirType::dialect("!llvm.ptr<7>").unwrap() + ); + if space != 7 { + assert_eq!(int.operands[0].id, casts[0].results[0].id); + } + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs new file mode 100644 index 0000000000..b750f39779 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_cute.rs @@ -0,0 +1,178 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Standalone SM100 attention MMA and TMA source-plan integration. + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_cute::sm100_attention_ops::*; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +/// Check standalone attention MMA/TMA plans lower without requiring the paired-GEMM lifecycle. +#[test] +fn independent_attention_mma_and_rank_three_copy_lower_without_gemm_story() { + for group in [1, 2] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "attention".try_into().unwrap()); + let f16 = MirFP16Type::get(&ctx).into(); + let u32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let b: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get(&mut ctx, f16, true, 3).into(); + let desc: TypeHandle = MirPtrType::get(&mut ctx, u64, false, 0).into(); + let barrier: TypeHandle = MirPtrType::get(&mut ctx, u64, true, 3).into(); + let args = vec![u32, shared, shared, b, u32, desc, barrier, u32, u32]; + let ty = FunctionType::get(&ctx, args.clone(), vec![]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, f, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "attention_probe".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(f.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + let plan = CuteSm100MmaAttr { + m: 128, + n: 128, + k: 128, + cta_group: group, + a_mn: false, + b_mn: false, + a_layout: shared_layout(128 / group, 128, false).unwrap(), + b_layout: shared_layout(128 / group, 128, false).unwrap(), + }; + CuteSm100MmaSsOp::new( + &mut ctx, + vec![a[0], a[1], a[2], a[3], a[7], a[8]], + plan.clone(), + ) + .get_operation() + .insert_at_back(block, &ctx); + if group == 1 { + let mut pv = plan.clone(); + pv.b_mn = true; + pv.b_layout = shared_layout(128, 128, true).unwrap(); + CuteSm100MmaTsOp::new(&mut ctx, vec![a[0], a[4], a[2], a[3], a[8]], pv) + .get_operation() + .insert_at_back(block, &ctx); + } + let copy = CuteSm100TmaAttr { + layout: plan.a_layout, + cta_group: group, + }; + CuteSm100TmaLoadOp::new( + &mut ctx, + vec![a[5], a[1], a[4], a[4], a[4], a[6], a[4]], + copy.clone(), + ) + .get_operation() + .insert_at_back(block, &ctx); + CuteSm100TmaStoreOp::new(&mut ctx, vec![a[5], a[1], a[4], a[4], a[4]], copy) + .get_operation() + .insert_at_back(block, &ctx); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + let target = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let text = cuda_oxide_mlir_export::render_module(&target).unwrap(); + assert_eq!( + text.matches("\"nvvm.tcgen05.mma\"").count(), + if group == 1 { 16 } else { 8 } + ); + // The allocation bases construct descriptors before dynamic byte offsets + // are converted to the descriptor's 16-byte address units. + assert_eq!( + text.matches("\"cute_nvgpu.make_umma_smem_desc\"").count(), + if group == 1 { 3 } else { 2 } + ); + assert_eq!( + text.matches("\"arith.shrui\"").count(), + if group == 1 { 3 } else { 2 } + ); + assert!(!text.contains("llvm.getelementptr"), "{text}"); + for shift in text.lines().filter(|line| line.contains("\"arith.shrui\"")) { + let (_, operands) = shift.split_once("\"arith.shrui\"(").unwrap(); + let (input, rest) = operands.split_once(", ").unwrap(); + let amount = rest.split_once(')').unwrap().0; + assert!( + text.lines().any( + |line| line.trim_start().starts_with(&format!("{amount} = ")) + && line.contains("value = 4 : i64") + ), + "{text}" + ); + assert!( + text.lines() + .any(|line| line.trim_start().starts_with(&format!("{input} = ")) + && line.contains("\"arith.extui\"") + && line.contains("(i32) -> (i64)")), + "{text}" + ); + } + assert!(text.contains("#cute_nvgpu.major"), "{text}"); + assert!(text.contains("num_cta = "), "{text}"); + assert!( + text.contains("operandSegmentSizes = array"), + "{text}" + ); + assert!(!text.contains("\"arith.shli\""), "{text}"); + if group == 2 { + assert!(text.contains("num_cta = 2 : i32"), "{text}"); + let mapped = text + .lines() + .find(|line| line.contains("\"nvvm.mapa\"")) + .unwrap(); + let (_, operands) = mapped.split_once("\"nvvm.mapa\"(").unwrap(); + let rank = operands + .split_once(", ") + .unwrap() + .1 + .split_once(')') + .unwrap() + .0; + assert!( + text.lines() + .any(|line| line.trim_start().starts_with(&format!("{rank} = ")) + && line.contains("value = 0 : i32")), + "{text}" + ); + } + if group == 1 { + assert!(text.contains("#cute_nvgpu.major"), "{text}"); + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs new file mode 100644 index 0000000000..2106668abe --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_descriptor.rs @@ -0,0 +1,227 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Execute the actual exported descriptor arithmetic on the CPU. Native CuTe +//! descriptor construction is an input to this probe; its layout is unchanged. +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirAttribute, MlirBlock, MlirOperation, MlirType, MlirValueId}; +use dialect_cute::sm100_attention_ops::*; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; +use std::collections::HashMap; + +/// Exported descriptor-building SSA graph evaluated independently of GPU execution. +struct Probe { + /// Target operation defining each SSA result inspected by the probe. + definitions: HashMap, + /// SSA argument representing the dynamic shared-memory byte offset. + offset: MlirValueId, + /// Descriptor result IDs whose address fields are checked. + descriptors: Vec, +} + +/// Find the translated block containing the TCGen05 MMA instruction. +fn function_block(op: &MlirOperation) -> Option<&MlirBlock> { + for region in &op.regions { + for block in ®ion.blocks { + if block + .operations + .iter() + .any(|op| op.name == "nvvm.tcgen05.mma") + { + return Some(block); + } + for child in &block.operations { + if let Some(found) = function_block(child) { + return Some(found); + } + } + } + } + None +} + +impl Probe { + /// Export an attention MMA plan for the requested shape and shared-memory major mode. + fn new(rows: u32, k: u32, mn: bool) -> Self { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "descriptor_probe".try_into().unwrap()); + let u32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let b: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let f16 = MirFP16Type::get(&ctx).into(); + let shared: TypeHandle = MirPtrType::get(&mut ctx, f16, true, 3).into(); + let args = vec![u32, shared, shared, b, u32, u32]; + let ty = FunctionType::get(&ctx, args.clone(), vec![]); + let f = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, f, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "descriptor_probe".try_into().unwrap()); + module.append_operation(&mut ctx, f, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(f.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + let group = if rows <= 128 { 2 } else { 1 }; + CuteSm100MmaSsOp::new( + &mut ctx, + a, + CuteSm100MmaAttr { + m: 128, + n: rows * group, + k, + cta_group: group, + a_mn: false, + b_mn: mn, + a_layout: shared_layout(128 / group, k, false).unwrap(), + b_layout: shared_layout(rows, k, mn).unwrap(), + }, + ) + .get_operation() + .insert_at_back(block, &ctx); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + let target = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let block = function_block(&target.root).unwrap(); + Self { + offset: block.arguments[5].id, + descriptors: block + .operations + .iter() + .filter(|op| op.name == "nvvm.tcgen05.mma") + .map(|op| op.operands[2].id) + .collect(), + definitions: block + .operations + .iter() + .filter(|op| !op.results.is_empty()) + .map(|op| (op.results[0].id, op.clone())) + .collect(), + } + } + + /// Interpret descriptor arithmetic using a chosen base descriptor and dynamic byte offset. + fn evaluate(&self, id: MlirValueId, base: u64, byte_offset: u32) -> u64 { + if id == self.offset { + return byte_offset as u64; + } + let op = &self.definitions[&id]; + let operand = |i: usize| self.evaluate(op.operands[i].id, base, byte_offset); + let value = match op.name.as_str() { + "cute_nvgpu.make_umma_smem_desc" => base, + "builtin.unrealized_conversion_cast" | "arith.extui" | "arith.trunci" => operand(0), + "arith.constant" => match op.properties["value"] { + MlirAttribute::Integer { value, .. } => value as u64, + _ => panic!("non-integer constant"), + }, + "arith.addi" => operand(0).wrapping_add(operand(1)), + "arith.andi" => operand(0) & operand(1), + "arith.ori" => operand(0) | operand(1), + "arith.shrui" => operand(0) >> operand(1), + other => panic!("unexpected descriptor arithmetic: {other}"), + }; + match op.results[0].ty { + MlirType::Integer(32) => value as u32 as u64, + _ => value, + } + } +} + +/// Check address updates leave every descriptor metadata field unchanged across boundary cases. +#[test] +fn descriptor_update_preserves_metadata_across_the_entire_address_field() { + let probe = Probe::new(64, 64, false); + // Both supported swizzles, stride extremes, and non-address bit patterns. + // The last pattern deliberately fills reserved bits to catch accidental + // masking of metadata independently of native descriptor construction. + for metadata in [0x4000_4040_0001_0000, 0x8000_7fff_3fff_0000, !0x3fffu64] { + for address in 0..=0x3fffu64 { + let remaining = 0x3fff - address; + for delta in [0, remaining / 2, remaining] { + let base = metadata | address; + let got = probe.evaluate(probe.descriptors[0], base, (delta * 16) as u32); + assert_eq!(got, base + delta, "address={address:#x}, delta={delta:#x}"); + assert_eq!(got & !0x3fff, metadata); + } + } + } +} + +/// Check K-major and MN-major descriptor updates agree with independent byte-address arithmetic. +#[test] +fn tile_and_k_offsets_match_addition_for_both_layout_families() { + for rows in (8..=256).step_by(8) { + if rows > 128 && rows % 16 != 0 { + continue; + } + for k in [32, 64, 128, 256] { + for mn in [false, true] { + if shared_layout(rows, k, mn).is_err() { + continue; + } + let probe = Probe::new(rows, k, mn); + let tile_bytes = rows * k * 2; + let span = if mn { rows } else { k }.min(64); + let period = if span == 32 { 512 } else { 1024 }; + let last_base = (0x40000 - tile_bytes) / period * period; + for base_bytes in [0, period, last_base] { + let last_tile = (0x40000 - base_bytes) / tile_bytes - 1; + for tile in [0, last_tile / 2, last_tile] { + let byte_offset = tile * tile_bytes; + let base = 0x4000_4040_0001_0000 | (base_bytes / 16) as u64; + for (i, descriptor) in probe.descriptors.iter().enumerate() { + // Physical matrix strides, expressed directly in + // bytes rather than the exporter's address units. + let col = i as u32 * 16; + let within_tile = if mn { + col * span * 2 + } else { + (col / span * rows * span + col % span) * 2 + }; + let expected = base + (byte_offset + within_tile) as u64 / 16; + assert_eq!( + probe.evaluate(*descriptor, base, byte_offset), + expected, + "rows={rows}, k={k}, mn={mn}, base={base_bytes}, tile={tile}, col={col}" + ); + } + } + } + } + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs new file mode 100644 index 0000000000..2ab62cc118 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/attention_primitives.rs @@ -0,0 +1,222 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Attention arithmetic retains exact precision selectors and lane ordering. +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_mir::ops::{MirFuncOp, MirReturnOp}; +use dialect_nvvm::ops::*; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::{TypeHandle, Typed}, +}; + +/// Attach the canonical catalog identity and append the source intrinsic to the test block. +fn mark(ctx: &mut Context, op: Ptr, block: Ptr, marker: &str) { + op.deref_mut(ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + op.insert_at_back(block, ctx); + if op.deref(ctx).get_num_results() == 1 { + let value = op.deref(ctx).get_result(0); + let args = block.deref(ctx).arguments().collect::>(); + let ptr = args + .iter() + .rev() + .find(|arg| { + arg.get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|p| p.pointee == value.get_type(ctx)) + }) + .copied() + .expect("output pointer"); + let store = Operation::new( + ctx, + dialect_mir::ops::MirStoreOp::get_concrete_op_info(), + vec![], + vec![ptr, value], + vec![], + 0, + ); + dialect_mir::ops::MirStoreOp::new(store).set_volatile(ctx, true); + store.insert_at_back(block, ctx); + } +} + +/// Check attention arithmetic uses the reviewed precise native intrinsic spellings. +#[test] +fn full_cute_attention_numeric_pack_keeps_precise_intrinsics() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "numeric".try_into().unwrap()); + let f32: TypeHandle = FP32Type::get(&ctx).into(); + let i32: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let pred = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let shared = dialect_mir::types::MirPtrType::get(&mut ctx, i32, true, 3).into(); + let global = dialect_mir::types::MirPtrType::get(&mut ctx, i32, false, 1).into(); + let u64 = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let bar = dialect_mir::types::MirPtrType::get(&mut ctx, u64, true, 3).into(); + let f_out = dialect_mir::types::MirPtrType::get(&mut ctx, f32, true, 1).into(); + let u_out = dialect_mir::types::MirPtrType::get(&mut ctx, i32, true, 1).into(); + let boolean = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let b_out = dialect_mir::types::MirPtrType::get(&mut ctx, boolean, true, 1).into(); + let packed_out = dialect_mir::types::MirPtrType::get(&mut ctx, u64, true, 1).into(); + let args = vec![ + f32, f32, f32, i32, i32, pred, shared, global, bar, f_out, u_out, b_out, u64, u64, u64, + packed_out, + ]; + let fty = FunctionType::get(&ctx, args.clone(), vec![]); + let op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let f = MirFuncOp::new(&mut ctx, op, TypeAttr::new(fty.into())); + f.set_symbol_name(&mut ctx, "numeric_probe".try_into().unwrap()); + module.append_operation(&mut ctx, op, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(op.deref(&ctx).get_region(0), &ctx); + let a = block.deref(&ctx).arguments().collect::>(); + // Distinct dynamic packed words retain both FP32 lanes through native arithmetic. + let add = AddF32x2Op::build(&mut ctx, a[12], a[13]); + mark(&mut ctx, add, block, "v1:i0995"); + let mul = MulF32x2Op::build(&mut ctx, a[12], a[13]); + mark(&mut ctx, mul, block, "v1:i0999"); + let fma = FmaF32x2Op::build(&mut ctx, a[12], a[13], a[14]); + mark(&mut ctx, fma, block, "v1:i1001"); + let shuffle = ShflSyncDownF32Op::build(&mut ctx, a[3], a[0], a[4]); + mark(&mut ctx, shuffle, block, "v1:i0056"); + let shuffle = ShflSyncBflyI32Op::build(&mut ctx, a[3], a[4], a[4]); + mark(&mut ctx, shuffle, block, "v1:i0051"); + // Ordinary Rust selection and saturation survive the same consumer boundary. + use dialect_mir::{ + ops::{MirCallOp, MirStoreOp}, + rust_intrinsics::{CALLEE_SATURATING_SUB, CALLEE_SELECT_UNPREDICTABLE}, + }; + let select = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![i32], + vec![a[5], a[3], a[4]], + vec![], + 0, + ); + MirCallOp::new(select) + .set_attr_callee(&ctx, StringAttr::new(CALLEE_SELECT_UNPREDICTABLE.into())); + select.insert_at_back(block, &ctx); + let selected = select.deref(&ctx).get_result(0); + let sub = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + vec![i32], + vec![selected, a[4]], + vec![], + 0, + ); + MirCallOp::new(sub).set_attr_callee(&ctx, StringAttr::new(CALLEE_SATURATING_SUB.into())); + sub.insert_at_back(block, &ctx); + let value = sub.deref(&ctx).get_result(0); + let store = Operation::new( + &mut ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![a[10], value], + vec![], + 0, + ); + MirStoreOp::new(store).set_volatile(&mut ctx, true); + store.insert_at_back(block, &ctx); + + for (kind, ftz, marker) in [ + (ScalarMathOperationAttr::Ex2, false, "v1:i0819"), + (ScalarMathOperationAttr::Lg2, false, "v1:i0786"), + (ScalarMathOperationAttr::Rcp, true, "v1:i0788"), + ] { + let op = ScalarMathOp::build( + &mut ctx, + a[0], + ScalarMathFormatAttr::F32, + kind, + ScalarMathPrecisionAttr::Approx, + if ftz { + ScalarMathSubnormalAttr::Ftz + } else { + ScalarMathSubnormalAttr::Preserve + }, + ); + mark(&mut ctx, op, block, marker); + } + let op = ScalarArithmeticOp::build( + &mut ctx, + a[..3].to_vec(), + ScalarArithmeticFormatAttr::F32, + ScalarArithmeticOperationAttr::Fma, + ScalarArithmeticRoundingAttr::Rn, + ScalarArithmeticSubnormalAttr::Preserve, + ScalarArithmeticSaturationAttr::None, + ); + mark(&mut ctx, op, block, "v1:i0418"); + let op = CvtF16x2F32Op::build(&mut ctx, a[0], a[1]); + mark(&mut ctx, op, block, "v1:i0081"); + let op = ReduxSyncFmaxOp::build(&mut ctx, a[3], a[0]); + mark(&mut ctx, op, block, "v1:i1007"); + let op = SetMaxNRegIncOp::build(&mut ctx, 192); + mark(&mut ctx, op, block, "v1:i0915"); + let op = CpAsyncCa4Op::build(&mut ctx, a[6], a[7]); + mark(&mut ctx, op, block, "v1:i0086"); + let op = CpAsyncMbarrierArriveNoIncOp::build(&mut ctx, a[8]); + mark(&mut ctx, op, block, "v1:i0103"); + let op = VoteSyncAllOp::build(&mut ctx, a[3], a[5]); + mark(&mut ctx, op, block, "v1:i0040"); + let op = BarWarpSyncOp::build(&mut ctx, a[3]); + mark(&mut ctx, op, block, "v1:i0049"); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + let result = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap(); + let text = cuda_oxide_mlir_export::render_module(&result).unwrap(); + for name in [ + "llvm.nvvm.ex2.approx.f32", + "llvm.nvvm.lg2.approx.f", + "llvm.nvvm.rcp.approx.ftz.f", + "llvm.nvvm.fma.rn.f", + "nvvm.redux.sync", + "nvvm.cp.async.shared.global", + "nvvm.cp.async.mbarrier.arrive", + "nvvm.vote.sync", + "nvvm.bar.warp.sync", + "llvm.nvvm.ff2f16x2.rn", + "nvvm.add.packed.f32x2", + "nvvm.mul.packed.f32x2", + "nvvm.fma.packed.f32x2", + ] { + assert!(text.contains(name), "{text}"); + } + assert!(text.contains("192 : i32"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs new file mode 100644 index 0000000000..c27fd82489 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/checked_storage.rs @@ -0,0 +1,453 @@ +//! Checked arithmetic, enum byte storage, shared globals, and Rust call contracts. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassPrimitivesMlir22, MlirConsumerProfile}; +use dialect_mir::{ + attributes::{FieldIndexAttr, VariantIndexAttr}, + ops::*, + types::*, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::TypeHandle, + utils::apint::APInt, + value::Value, +}; +/// Append a typed source operation to the fixture block and return its operation handle. +fn append( + ctx: &mut Context, + b: Ptr, + types: Vec, + values: Vec, +) -> Ptr { + let o = Operation::new(ctx, O::get_concrete_op_info(), types, values, vec![], 0); + o.insert_at_back(b, ctx); + o +} +/// Create a named MIR function and its argument-bearing entry block for a storage test. +fn function( + ctx: &mut Context, + args: Vec, + returns: Vec, +) -> (ModuleOp, Ptr) { + let m = ModuleOp::new(ctx, "storage".try_into().unwrap()); + let ft = FunctionType::get(ctx, args.clone(), returns); + let f = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(ctx, f, TypeAttr::new(ft.into())) + .set_symbol_name(ctx, "test".try_into().unwrap()); + m.append_operation(ctx, f, 0); + let b = BasicBlock::new(ctx, None, args); + b.insert_at_back(f.deref(ctx).get_region(0), ctx); + (m, b) +} +/// Translate a prepared source module with the SM100 primitives profile and render its MLIR. +fn render(ctx: &Context, m: &ModuleOp) -> Result { + CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(ctx, m) + .and_then(|m| render_module(&m).map_err(Into::into)) + .map_err(|e| e.to_string()) +} +/// Check overflow intrinsics use source signedness and return the exact padded Rust tuple. +#[test] +fn checked_arithmetic_preserves_signedness_and_padded_overflow_tuple() { + for signedness in [Signedness::Signed, Signedness::Unsigned] { + for mul in [false, true] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, 64, signedness).into(); + let flag: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let tuple: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![word, flag], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let (m, b) = function(&mut ctx, vec![word, word], vec![tuple]); + let args = b.deref(&ctx).arguments().collect(); + let op = if mul { + append::(&mut ctx, b, vec![tuple], args) + } else { + append::(&mut ctx, b, vec![tuple], args) + }; + let v = op.deref(&ctx).get_result(0); + append::(&mut ctx, b, vec![], vec![v]); + let text = render(&ctx, &m).unwrap(); + let sign = if signedness == Signedness::Signed { + "s" + } else { + "u" + }; + let name = if mul { "mul" } else { "add" }; + assert!( + text.contains(&format!("llvm.{sign}{name}.with.overflow.i64")), + "{text}" + ); + assert!( + text.contains("!llvm.struct<(i64, i1, !llvm.array<7 x i8>)>"), + "{text}" + ); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), 2); + assert_eq!(text.matches("\"llvm.insertvalue\"").count(), 2); + } + } +} +/// Check enum tag and payload accesses use their recorded byte offsets and storage size. +#[test] +fn direct_enum_uses_physical_offsets_for_tag_and_active_payload() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let tag: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let word: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let e: TypeHandle = MirEnumType::get_with_layout( + &mut ctx, + "Result".into(), + tag, + vec![7, 19], + vec![ + EnumVariant::new_with_layout("Ok".into(), vec![word], vec![8], vec![8]), + EnumVariant::new_with_layout("Err".into(), vec![tag], vec![1], vec![1]), + ], + 0, + 16, + 8, + ) + .into(); + let (m, b) = function(&mut ctx, vec![word], vec![word]); + let arg = b.deref(&ctx).get_argument(0); + let c = append::(&mut ctx, b, vec![e], vec![arg]); + MirConstructEnumOp::new(c).set_attr_construct_enum_variant_index(&ctx, VariantIndexAttr(0)); + let v = c.deref(&ctx).get_result(0); + append::(&mut ctx, b, vec![tag], vec![v]); + let p = append::(&mut ctx, b, vec![word], vec![v]); + let p = MirEnumPayloadOp::new(p); + p.set_attr_payload_variant_index(&ctx, VariantIndexAttr(0)); + p.set_attr_payload_field_index(&ctx, FieldIndexAttr(0)); + let v = p.get_operation().deref(&ctx).get_result(0); + append::(&mut ctx, b, vec![], vec![v]); + let text = render(&ctx, &m).unwrap(); + assert!(text.contains("value = 7 : i8")); + assert_eq!( + text.matches("rawConstantIndices = array").count(), + 2 + ); + assert_eq!( + text.matches("rawConstantIndices = array").count(), + 2 + ); + assert!(text.contains("!llvm.struct<(array<16 x i8>, array<0 x i64>)>")); +} +/// Build repeated shared declarations, optionally introducing a named storage-contract conflict. +fn shared_case(conflict: &str) -> Result { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let unsigned = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let signed = IntegerType::get(&ctx, 32, Signedness::Signed); + let size_ty = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let (m, b) = function(&mut ctx, vec![], vec![]); + for i in 0..2 { + let element: TypeHandle = if conflict == "type" && i == 1 { + signed.into() + } else { + unsigned.into() + }; + let ptr: TypeHandle = MirPtrType::get(&mut ctx, element, true, 3).into(); + let a = append::(&mut ctx, b, vec![ptr], vec![]); + let a = MirSharedAllocOp::new(a); + a.set_attr_elem_type(&ctx, TypeAttr::new(element)); + a.set_attr_size( + &ctx, + IntegerAttr::new( + size_ty, + APInt::from_u64( + if conflict == "size" && i == 1 { 8 } else { 4 }, + 64.try_into().unwrap(), + ), + ), + ); + a.set_attr_alloc_key( + &ctx, + StringAttr::new(if conflict == "distinct" { + format!("key{i}") + } else { + "key".into() + }), + ); + // An absent alignment means the element's natural alignment; the dialect + // rejects an explicit zero. + if conflict == "alignment" && i == 1 { + a.set_alignment_value(&mut ctx, 16); + } + } + append::(&mut ctx, b, vec![], vec![]); + render(&ctx, &m) +} +/// Check matching shared declarations merge, conflicting ones fail, and an absent alignment uses the natural value. +#[test] +fn shared_declarations_deduplicate_only_identical_storage_and_absent_alignment_means_natural() { + let text = shared_case("same").unwrap(); + assert_eq!(text.matches("\"llvm.mlir.global\"").count(), 1); + assert!(text.contains("alignment = 4 : i64")); + assert!(!text.contains("__source_element_identity")); + assert_eq!( + shared_case("distinct") + .unwrap() + .matches("\"llvm.mlir.global\"") + .count(), + 2 + ); + for conflict in ["type", "size", "alignment"] { + assert!( + shared_case(conflict) + .unwrap_err() + .contains("incompatible declarations"), + "{conflict}" + ); + } +} + +/// Check dynamic shared storage uses one external AS3 global and byte-correct aligned offsets. +#[test] +fn primitive_dynamic_shared_uses_external_as3_storage_with_alignment_and_offsets() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let byte: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, byte, true, 3).into(); + let (module, block) = function(&mut ctx, vec![], vec![ptr]); + for (offset, alignment) in [(0, 128), (4096, 1024)] { + let op = append::(&mut ctx, block, vec![ptr], vec![]); + let shared = MirExternSharedOp::new(op); + shared.set_alignment_value(&mut ctx, alignment); + shared.set_byte_offset_value(&mut ctx, offset); + if offset != 0 { + let value = op.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![value]); + } + } + let text = render(&ctx, &module).unwrap(); + assert_eq!(text.matches("\"llvm.mlir.global\"").count(), 1, "{text}"); + assert!( + text.contains("!llvm.array<0 x i8>") && text.contains("alignment = 1024 : i64"), + "{text}" + ); + assert!( + text.contains("#llvm.linkage") && text.contains("addr_space = 3 : i32"), + "{text}" + ); + assert!( + text.contains("4096") && text.contains("\"llvm.getelementptr\""), + "{text}" + ); + assert!( + !text.contains("\"cute") && !text.contains("!cute"), + "{text}" + ); +} + +/// Check an unreachable device contract emits a trap before the LLVM terminator. +#[test] +fn primitive_contract_failure_emits_trap_before_unreachable() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let (module, block) = function(&mut ctx, vec![], vec![]); + let trap = append::(&mut ctx, block, vec![], vec![]); + trap.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0295".into()), + ); + append::(&mut ctx, block, vec![], vec![]); + let text = render(&ctx, &module).unwrap(); + assert!( + text.find("\"llvm.intr.trap\"") + .expect("trap must use the registered LLVM intrinsic") + < text.find("\"llvm.unreachable\"").unwrap(), + "{text}" + ); +} + +/// Check unit-returning calls use a void callee signature while retaining needed empty-value carriers. +#[test] +fn logical_unit_call_uses_void_signature_and_preserves_value_uses() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let unit: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); + let pointer: TypeHandle = MirPtrType::get(&mut ctx, unit, true, 0).into(); + let (module, block) = function(&mut ctx, vec![pointer], vec![]); + let destination = block.deref(&ctx).get_argument(0); + let call = append::(&mut ctx, block, vec![unit], vec![destination]); + MirCallOp::new(call).set_attr_callee(&ctx, StringAttr::new("test".into())); + let value = call.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![destination, value]); + append::(&mut ctx, block, vec![], vec![]); + let text = render(&ctx, &module).unwrap(); + let call = text + .lines() + .find(|line| line.contains("\"func.call\"")) + .unwrap(); + assert!(call.ends_with("-> ()"), "{text}"); + assert!(call.trim().starts_with("\"func.call\""), "{text}"); + let carrier = text + .lines() + .find(|line| line.contains("\"ub.poison\"")) + .unwrap(); + let value = carrier.trim().split(" = ").next().unwrap(); + assert!(text.contains(&format!("\"llvm.store\"({value},")), "{text}"); +} + +/// Keep explicit empty aggregate return values distinct from logical unit calls. +#[test] +fn explicitly_returned_zero_sized_aggregates_keep_their_call_result() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let tuple: TypeHandle = MirTupleType::get(&mut ctx, vec![]).into(); + let structure: TypeHandle = + MirStructType::get(&mut ctx, "Marker".into(), vec![], vec![]).into(); + for result in [tuple, structure] { + let (module, block) = function(&mut ctx, vec![], vec![result]); + let call = append::(&mut ctx, block, vec![result], vec![]); + MirCallOp::new(call).set_attr_callee(&ctx, StringAttr::new("test".into())); + let value = call.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![value]); + let text = render(&ctx, &module).unwrap(); + let call = text + .lines() + .find(|line| line.contains("\"func.call\"")) + .unwrap(); + assert!( + call.contains(" = ") && call.ends_with("-> (!llvm.struct<()>)"), + "{text}" + ); + assert!(!text.contains("\"ub.poison\""), "{text}"); + } +} + +/// Check reviewed scalar intrinsics become target operations without unresolved helper symbols. +#[test] +fn rust_scalar_intrinsics_are_operations_not_external_calls() { + use dialect_mir::rust_intrinsics::{CALLEE_SATURATING_SUB, CALLEE_SELECT_UNPREDICTABLE}; + for sign in [Signedness::Signed, Signedness::Unsigned] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, 32, sign).into(); + let flag: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let (m, b) = function(&mut ctx, vec![flag, word, word], vec![word]); + let args: Vec<_> = b.deref(&ctx).arguments().collect(); + let select = append::(&mut ctx, b, vec![word], args.clone()); + MirCallOp::new(select) + .set_attr_callee(&ctx, StringAttr::new(CALLEE_SELECT_UNPREDICTABLE.into())); + let selected = select.deref(&ctx).get_result(0); + let sub = append::(&mut ctx, b, vec![word], vec![selected, args[2]]); + MirCallOp::new(sub).set_attr_callee(&ctx, StringAttr::new(CALLEE_SATURATING_SUB.into())); + let result = sub.deref(&ctx).get_result(0); + append::(&mut ctx, b, vec![], vec![result]); + let text = render(&ctx, &m).unwrap(); + assert!(text.contains("\"arith.select\""), "{text}"); + let sign = if sign == Signedness::Signed { "s" } else { "u" }; + assert!(text.contains(&format!("llvm.{sign}sub.sat.i32")), "{text}"); + assert!(!text.contains("__cuda_oxide_rust_intrinsic_"), "{text}"); + } +} + +/// Check constants introduced by unrolling preserve their complete source bit patterns. +#[test] +fn unrolling_builtin_integer_constants_retain_width_and_bits() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let word = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let (m, b) = function(&mut ctx, vec![], vec![word.into()]); + let constant = pliron::builtin::ops::ConstantOp::new( + &mut ctx, + Box::new(IntegerAttr::new( + word, + APInt::from_u64(u32::MAX as u64, 32.try_into().unwrap()), + )), + ); + constant.get_operation().insert_at_back(b, &ctx); + let value = constant.get_operation().deref(&ctx).get_result(0); + append::(&mut ctx, b, vec![], vec![value]); + let text = render(&ctx, &m).unwrap(); + assert!( + text.contains("\"arith.constant\"() <{value = -1 : i32}>"), + "{text}" + ); + assert!(!text.contains("builtin.constant"), "{text}"); +} + +/// Check the Rust maximum maps to native NaN-ignoring FP32 behavior without FTZ modifiers. +#[test] +fn rust_fp32_max_preserves_nan_ignoring_and_subnormal_contract() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let scalar: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); + let (module, block) = function(&mut ctx, vec![scalar, scalar], vec![scalar]); + let args = block.deref(&ctx).arguments().collect(); + let call = append::(&mut ctx, block, vec![scalar], args); + MirCallOp::new(call).set_attr_callee( + &ctx, + StringAttr::new(dialect_mir::rust_intrinsics::CALLEE_MAXNUM_NSZ_F32.into()), + ); + let value = call.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![value]); + let text = render(&ctx, &module).unwrap(); + let native = text + .lines() + .find(|line| line.contains("\"nvvm.fmax\"")) + .unwrap(); + assert!( + !native.contains("nan") && !native.contains("ftz") && !native.contains("abs"), + "{text}" + ); + assert!( + !text.contains("llvm.maxnum") && !text.contains("__cuda_oxide_rust_intrinsic_"), + "{text}" + ); +} + +/// Reject attempts to use the FP32 maximum mapping with incompatible source types. +#[test] +fn rust_fp32_max_rejects_other_signatures() { + for (operand_width, result_width, operand_count) in [(64, 32, 2), (32, 64, 2), (32, 32, 1)] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let fp32: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); + let fp64: TypeHandle = pliron::builtin::types::FP64Type::get(&ctx).into(); + let operand = if operand_width == 32 { fp32 } else { fp64 }; + let result = if result_width == 32 { fp32 } else { fp64 }; + let (module, block) = function(&mut ctx, vec![operand; operand_count], vec![result]); + let args = block.deref(&ctx).arguments().collect(); + let call = append::(&mut ctx, block, vec![result], args); + MirCallOp::new(call).set_attr_callee( + &ctx, + StringAttr::new(dialect_mir::rust_intrinsics::CALLEE_MAXNUM_NSZ_F32.into()), + ); + let value = call.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![value]); + let error = render(&ctx, &module).unwrap_err(); + assert!( + error.contains("requires exactly two FP32 operands and one FP32 result"), + "{error}" + ); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs new file mode 100644 index 0000000000..443ad81b24 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_elementwise.rs @@ -0,0 +1,485 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Elementwise view/copy mappings and production kernel launch-ABI integration. + +use std::num::NonZero; + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_cute::{ + attributes::CuteTensorAccessAttr, + tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, + }, +}; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirFP16Type, MirPtrType, MirSliceType, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::TypeHandle, + utils::apint::APInt, +}; + +/// Append a typed source operation to the fixture block and return its operation handle. +fn append(block: Ptr, operation: O, ctx: &Context) -> Ptr { + let operation = operation.get_operation(); + operation.insert_at_back(block, ctx); + operation +} + +/// Check full-tile CuTe operations, scalar tail accesses, and pointer-space bridges together. +#[test] +fn all_eight_elementwise_ops_map_to_cutlass_cute_and_core_mlir() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, Identifier::try_from("cute_elementwise").unwrap()); + let f32_type: TypeHandle = FP32Type::get(&ctx).into(); + let f16_type: TypeHandle = MirFP16Type::get(&ctx).into(); + let u64_type: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let i1_type: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + // Rust slice fields arrive through the generic pointer ABI even though + // the resulting CuTe view is global memory. Exercise that real bridge + // here; the exporter must spell the AS0 -> AS1 transition explicitly. + let global_ro: TypeHandle = + MirPtrType::get(&mut ctx, f32_type, false, address_space::GENERIC).into(); + let global_rw: TypeHandle = + MirPtrType::get(&mut ctx, f32_type, true, address_space::GENERIC).into(); + let local_rw: TypeHandle = + MirPtrType::get(&mut ctx, f32_type, true, address_space::LOCAL).into(); + let global_ro_f16: TypeHandle = + MirPtrType::get(&mut ctx, f16_type, false, address_space::GENERIC).into(); + let local_rw_f16: TypeHandle = + MirPtrType::get(&mut ctx, f16_type, true, address_space::LOCAL).into(); + let arguments = vec![ + global_ro, + global_rw, + local_rw, + u64_type, + u64_type, + f32_type, + global_ro_f16, + local_rw_f16, + ]; + let results = vec![i1_type, u64_type]; + let function_type = FunctionType::get(&ctx, arguments.clone(), results); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("elementwise").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let region = function_op.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, arguments); + entry.insert_at_back(region, &ctx); + let arguments = entry.deref(&ctx).arguments().collect::>(); + let read_ptr = arguments[0]; + let write_ptr = arguments[1]; + let carrier = arguments[2]; + let len = arguments[3]; + let tile_index = arguments[4]; + let scalar = arguments[5]; + let read_ptr_f16 = arguments[6]; + let carrier_f16 = arguments[7]; + + let read_make = append( + entry, + CuteTensorMakeOp::new( + &mut ctx, + read_ptr, + len, + f32_type, + f32_type, + CuteTensorAccessAttr::ReadOnly, + 4, + ), + &ctx, + ); + let read = read_make.deref(&ctx).get_result(0); + let read_zipped = append( + entry, + CuteTensorZippedDivideOp::new(&mut ctx, read, 4), + &ctx, + ); + let read_zipped = read_zipped.deref(&ctx).get_result(0); + let read_tile = append( + entry, + CuteTensorSliceOp::new(&mut ctx, read_zipped, tile_index), + &ctx, + ); + let read_tile = read_tile.deref(&ctx).get_result(0); + let is_full = append(entry, CuteTensorIsFullOp::new(&mut ctx, read_tile), &ctx); + let base = append(entry, CuteTensorBaseOp::new(&mut ctx, read_tile), &ctx); + append( + entry, + CuteTensorLoadIntoOp::new(&mut ctx, read_tile, carrier, 16), + &ctx, + ); + + let write_make = append( + entry, + CuteTensorMakeOp::new( + &mut ctx, + write_ptr, + len, + f32_type, + f32_type, + CuteTensorAccessAttr::ReadWrite, + 4, + ), + &ctx, + ); + let write = write_make.deref(&ctx).get_result(0); + let write_zipped = append( + entry, + CuteTensorZippedDivideOp::new(&mut ctx, write, 4), + &ctx, + ); + let write_zipped = write_zipped.deref(&ctx).get_result(0); + let write_tile = append( + entry, + CuteTensorSliceOp::new(&mut ctx, write_zipped, tile_index), + &ctx, + ); + let write_tile = write_tile.deref(&ctx).get_result(0); + append( + entry, + CuteTensorStoreFromOp::new(&mut ctx, carrier, write_tile, 16), + &ctx, + ); + append( + entry, + CuteTensorStoreElementAbsOp::new(&mut ctx, write_tile, tile_index, scalar), + &ctx, + ); + + // The production example has a second instantiation: eight f16 values + // still form one 16-byte transaction. Keep that real type/shape pair in + // the parser fixture rather than assuming the f32 spelling generalizes. + let read_make_f16 = append( + entry, + CuteTensorMakeOp::new( + &mut ctx, + read_ptr_f16, + len, + f16_type, + f16_type, + CuteTensorAccessAttr::ReadOnly, + 2, + ), + &ctx, + ); + let read_f16 = read_make_f16.deref(&ctx).get_result(0); + let read_zipped_f16 = append( + entry, + CuteTensorZippedDivideOp::new(&mut ctx, read_f16, 8), + &ctx, + ); + let read_zipped_f16 = read_zipped_f16.deref(&ctx).get_result(0); + let read_tile_f16 = append( + entry, + CuteTensorSliceOp::new(&mut ctx, read_zipped_f16, tile_index), + &ctx, + ); + let read_tile_f16 = read_tile_f16.deref(&ctx).get_result(0); + append( + entry, + CuteTensorLoadIntoOp::new(&mut ctx, read_tile_f16, carrier_f16, 16), + &ctx, + ); + let is_full = is_full.deref(&ctx).get_result(0); + let base = base.deref(&ctx).get_result(0); + let return_op = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![is_full, base], + vec![], + 0, + ); + return_op.insert_at_back(entry, &ctx); + + // The profile runs the shared whole-module CuTe verifier, then + // composes builtin, MIR, NVVM, and CuTe mappings. + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_module(&target).unwrap(); + + assert!(text.contains("\"cute.make_view\""), "{text}"); + assert!(text.contains("\"cute.zipped_divide\""), "{text}"); + assert!(text.contains("\"cute.slice\""), "{text}"); + assert!( + text.contains("!cute_nvgpu.atom.universal_copy"), + "{text}" + ); + assert!( + text.contains("!cute_nvgpu.atom.universal_copy"), + "{text}" + ); + assert_eq!(text.matches("\"cute.copy\"").count(), 3, "{text}"); + assert_eq!( + text.matches("operandSegmentSizes = array") + .count(), + 3, + "{text}" + ); + assert_eq!(text.matches("\"arith.ori\"").count(), 2, "{text}"); + assert!(text.contains("\"arith.andi\""), "{text}"); + assert!(text.contains("value = 4611686018427387903 : i64"), "{text}"); + assert!(text.contains("predicate = 9 : i64"), "{text}"); + assert!(text.contains("\"llvm.getelementptr\""), "{text}"); + assert!(text.contains("\"llvm.store\""), "{text}"); + assert_eq!(text.matches("\"llvm.addrspacecast\"").count(), 6, "{text}"); + assert_eq!( + text.matches("(!llvm.ptr) -> (!llvm.ptr<1>)").count(), + 3, + "{text}" + ); + assert_eq!( + text.matches("(!llvm.ptr<5>) -> (!llvm.ptr)").count(), + 3, + "{text}" + ); + assert!(!text.contains("cute.tensor_"), "{text}"); +} + +/// Check one executable GPU envelope and the pointer/length expansion of three slice arguments. +#[test] +fn production_kernel_gets_one_device_module_and_the_six_argument_launch_abi() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let module = ModuleOp::new(&mut ctx, Identifier::try_from("elementwise_abi").unwrap()); + let f32_type: TypeHandle = FP32Type::get(&ctx).into(); + let slice_type: TypeHandle = MirSliceType::get(&mut ctx, f32_type).into(); + let function_type = FunctionType::get(&ctx, vec![slice_type, slice_type, slice_type], vec![]); + let kernel_operation = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let kernel = MirFuncOp::new( + &mut ctx, + kernel_operation, + TypeAttr::new(function_type.into()), + ); + kernel.set_symbol_name(&mut ctx, Identifier::try_from("add_f32").unwrap()); + { + let u32_type = IntegerType::get(&ctx, 32, Signedness::Unsigned); + let width = NonZero::new(32).unwrap(); + let mut operation = kernel_operation.deref_mut(&ctx); + operation.attributes.set( + Identifier::try_from("gpu_kernel").unwrap(), + StringAttr::new("true".into()), + ); + for (name, value) in [ + ("maxntid", 256), + ("minctasm", 1), + ("reqntid_x", 256), + ("reqntid_y", 1), + ("reqntid_z", 1), + ("cluster_dim_x", 2), + ("cluster_dim_y", 1), + ("cluster_dim_z", 1), + ] { + operation.attributes.set( + Identifier::try_from(name).unwrap(), + IntegerAttr::new(u32_type, APInt::from_u32(value, width)), + ); + } + } + module.append_operation(&mut ctx, kernel_operation, 0); + let kernel_region = kernel_operation.deref(&ctx).get_region(0); + let kernel_entry = BasicBlock::new(&mut ctx, None, vec![slice_type, slice_type, slice_type]); + kernel_entry.insert_at_back(kernel_region, &ctx); + append( + kernel_entry, + MirReturnOp::new(Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + )), + &ctx, + ); + + let helper_type = FunctionType::get(&ctx, vec![], vec![]); + let helper_operation = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let helper = MirFuncOp::new( + &mut ctx, + helper_operation, + TypeAttr::new(helper_type.into()), + ); + helper.set_symbol_name(&mut ctx, Identifier::try_from("device_helper").unwrap()); + module.append_operation(&mut ctx, helper_operation, 0); + let helper_region = helper_operation.deref(&ctx).get_region(0); + let helper_entry = BasicBlock::new(&mut ctx, None, vec![]); + helper_entry.insert_at_back(helper_region, &ctx); + append( + helper_entry, + MirReturnOp::new(Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + )), + &ctx, + ); + + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let root = &target.root; + assert_eq!( + root.attributes.get("gpu.container_module"), + Some(&cuda_oxide_mlir_export::MlirAttribute::Unit) + ); + assert_eq!(root.regions[0].blocks[0].operations.len(), 1); + let device_module = &root.regions[0].blocks[0].operations[0]; + assert_eq!(device_module.name, "gpu.module"); + assert_eq!( + device_module.properties.get("sym_name"), + Some(&cuda_oxide_mlir_export::MlirAttribute::String( + "kernels".into() + )) + ); + assert_eq!(device_module.regions[0].blocks[0].operations.len(), 2); + + let kernel = &device_module.regions[0].blocks[0].operations[0]; + assert_eq!( + kernel.attributes.get("cute.kernel"), + Some(&cuda_oxide_mlir_export::MlirAttribute::Unit) + ); + assert_eq!( + kernel.attributes.get("gpu.kernel"), + Some(&cuda_oxide_mlir_export::MlirAttribute::Unit) + ); + assert_eq!( + kernel.attributes.get("nvvm.reqntid"), + Some(&cuda_oxide_mlir_export::MlirAttribute::DenseI32Array(vec![ + 256, 1, 1 + ])) + ); + assert_eq!( + kernel.attributes.get("nvvm.minctasm"), + Some(&cuda_oxide_mlir_export::MlirAttribute::Integer { + value: 1, + ty: cuda_oxide_mlir_export::MlirType::Integer(32) + }) + ); + assert!(!kernel.attributes.contains_key("minctasm")); + assert!(!kernel.attributes.contains_key("maxntid")); + assert!(!kernel.attributes.contains_key("nvvm.maxntid")); + let cuda_oxide_mlir_export::MlirAttribute::Type(cuda_oxide_mlir_export::MlirType::Function { + inputs, + .. + }) = kernel.properties.get("function_type").unwrap() + else { + panic!("kernel function_type is not a function type") + }; + let pointer = cuda_oxide_mlir_export::MlirType::dialect("!llvm.ptr").unwrap(); + let length = cuda_oxide_mlir_export::MlirType::Integer(64); + let expected_inputs = vec![ + pointer.clone(), + length.clone(), + pointer.clone(), + length.clone(), + pointer, + length, + ]; + assert_eq!(inputs, &expected_inputs); + let entry = &kernel.regions[0].blocks[0]; + assert_eq!( + entry + .arguments + .iter() + .map(|argument| argument.ty.clone()) + .collect::>(), + expected_inputs + ); + assert_eq!( + entry + .operations + .iter() + .filter(|operation| operation.name == "llvm.mlir.undef") + .count(), + 3 + ); + assert_eq!( + entry + .operations + .iter() + .filter(|operation| operation.name == "llvm.insertvalue") + .count(), + 6 + ); + + let helper = &device_module.regions[0].blocks[0].operations[1]; + assert_eq!(helper.name, "func.func"); + assert!(!helper.attributes.contains_key("gpu.kernel")); + + let text = render_module(&target).unwrap(); + assert_eq!(text.matches("\"gpu.module\"").count(), 1, "{text}"); + assert!(text.contains("gpu.container_module = unit"), "{text}"); + assert!( + text.contains("nvvm.reqntid = array"), + "{text}" + ); + assert!( + text.contains("nvvm.cluster_dim = array"), + "{text}" + ); + assert!(!text.contains("cluster_dim_x"), "{text}"); + for key in ["reqntid_x", "reqntid_y", "reqntid_z"] { + kernel_operation + .deref_mut(&ctx) + .attributes + .0 + .remove(&Identifier::try_from(key).unwrap()); + } + let bounded = profile.translate_module(&ctx, &module).unwrap(); + let text = render_module(&bounded).unwrap(); + assert!( + text.contains("nvvm.maxntid = array"), + "{text}" + ); + assert!(text.contains("nvvm.minctasm = 1 : i32"), "{text}"); + assert!(!text.contains("nvvm.reqntid"), "{text}"); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs new file mode 100644 index 0000000000..9b5313bbc3 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/cute_gemv.rs @@ -0,0 +1,267 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Production-shaped NVFP4 GEMV view, narrow-load, and ordered-dot integration. + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirModule, render_module}; +use dialect_cute::{ + attributes::CuteTensorRoleAttr, + gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, + }, +}; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirPtrType, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +/// Append a typed source operation to the fixture block and return its operation handle. +fn append(block: Ptr, operation: O, ctx: &Context) -> Ptr { + let operation = operation.get_operation(); + operation.insert_at_back(block, ctx); + operation +} + +/// Construct the production-shaped scaled-view K64 load and ordered-dot source graph. +fn build_gemv_module(ctx: &mut Context) -> ModuleOp { + dialect_mir::register(ctx); + dialect_cute::register(ctx); + + let module = ModuleOp::new(ctx, Identifier::try_from("nvfp4_gemv").unwrap()); + let u8_type: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let u64_type: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let f32_type: TypeHandle = FP32Type::get(ctx).into(); + let pointer_type: TypeHandle = + MirPtrType::get(ctx, u8_type, false, address_space::GENERIC).into(); + let argument_types = vec![ + pointer_type, + u64_type, + u64_type, + u64_type, + pointer_type, + u64_type, + pointer_type, + u64_type, + u64_type, + u64_type, + pointer_type, + u64_type, + u64_type, + u64_type, + u64_type, + f32_type, + ]; + let function_type = FunctionType::get(ctx, argument_types.clone(), vec![f32_type]); + let function_operation = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(ctx, function_operation, TypeAttr::new(function_type.into())); + function.set_symbol_name(ctx, Identifier::try_from("nvfp4_gemv_step").unwrap()); + module.append_operation(ctx, function_operation, 0); + let region = function_operation.deref(ctx).get_region(0); + let entry = BasicBlock::new(ctx, None, argument_types); + entry.insert_at_back(region, ctx); + let arguments = entry.deref(ctx).arguments().collect::>(); + + let a_values = append( + entry, + CuteTensorMake2DOp::new_e2m1( + ctx, + arguments[0], + arguments[1], + arguments[2], + arguments[3], + CuteTensorRoleAttr::Mkl, + 1, + ), + ctx, + ) + .deref(ctx) + .get_result(0); + let a_scales = append( + entry, + CuteTensorMake2DOp::new_ue8m0( + ctx, + arguments[4], + arguments[5], + arguments[2], + arguments[3], + CuteTensorRoleAttr::Mkl, + 1, + 16, + ), + ctx, + ) + .deref(ctx) + .get_result(0); + let a = append( + entry, + CuteScaledViewMakeOp::new(ctx, a_values, a_scales), + ctx, + ) + .deref(ctx) + .get_result(0); + let a_row = append( + entry, + CuteScaledViewRowOp::new(ctx, a, arguments[14], arguments[2]), + ctx, + ) + .deref(ctx) + .get_result(0); + let a_tile = append( + entry, + CuteScaledViewKTileOp::new(ctx, a_row, arguments[14]), + ctx, + ) + .deref(ctx) + .get_result(0); + let a_fragment = append(entry, CuteScaledViewLoadOp::new(ctx, a_tile, 16, 4), ctx) + .deref(ctx) + .get_result(0); + + let b_values = append( + entry, + CuteTensorMake2DOp::new_e2m1( + ctx, + arguments[6], + arguments[7], + arguments[8], + arguments[9], + CuteTensorRoleAttr::Nkl, + 1, + ), + ctx, + ) + .deref(ctx) + .get_result(0); + let b_scales = append( + entry, + CuteTensorMake2DOp::new_ue8m0( + ctx, + arguments[10], + arguments[11], + arguments[8], + arguments[9], + CuteTensorRoleAttr::Nkl, + 1, + 16, + ), + ctx, + ) + .deref(ctx) + .get_result(0); + let b = append( + entry, + CuteScaledViewMakeOp::new(ctx, b_values, b_scales), + ctx, + ) + .deref(ctx) + .get_result(0); + let b_row = append( + entry, + CuteScaledViewRowOp::new(ctx, b, arguments[14], arguments[8]), + ctx, + ) + .deref(ctx) + .get_result(0); + let b_tile = append( + entry, + CuteScaledViewKTileOp::new(ctx, b_row, arguments[14]), + ctx, + ) + .deref(ctx) + .get_result(0); + let b_fragment = append(entry, CuteScaledViewLoadOp::new(ctx, b_tile, 16, 4), ctx) + .deref(ctx) + .get_result(0); + + let dot = append( + entry, + CuteDotOp::new(ctx, a_fragment, b_fragment, arguments[15]), + ctx, + ) + .deref(ctx) + .get_result(0); + append( + entry, + MirReturnOp::new(Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![dot], + vec![], + 0, + )), + ctx, + ); + module +} + +/// Export the source GEMV graph using the full CuTe SM120 profile. +fn translate_gemv_target(ctx: &Context, module: &ModuleOp) -> MlirModule { + CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(ctx, module) + .unwrap() +} + +/// Render the exported GEMV target for textual contract assertions. +fn translate_gemv(ctx: &Context, module: &ModuleOp) -> String { + render_module(&translate_gemv_target(ctx, module)).unwrap() +} + +/// Check the exporter retains narrow FP4/scale loads and the required FP32 accumulation order. +#[test] +fn production_profile_maps_gemv_to_exact_narrow_types_and_ordered_dot() { + let mut ctx = Context::new(); + let module = build_gemv_module(&mut ctx); + let text = translate_gemv(&ctx, &module); + + assert!( + text.contains("!cute.ptr>"), + "{text}" + ); + assert!( + text.contains("!cute.ptr>"), + "{text}" + ); + assert_eq!( + text.matches("\"cute.memref.load_vec\"").count(), + 6, + "{text}" + ); + assert_eq!(text.matches("\"nvgpu.cvt_fpext\"").count(), 6, "{text}"); + assert_eq!(text.matches("vector<32xf4E2M1FN>").count(), 8, "{text}"); + assert_eq!(text.matches("vector<4xf8E8M0FNU>").count(), 4, "{text}"); + assert_eq!(text.matches("\"arith.mulf\"").count(), 192, "{text}"); + assert_eq!(text.matches("\"arith.addf\"").count(), 64, "{text}"); + assert_eq!(text.matches("\"llvm.addrspacecast\"").count(), 6, "{text}"); + assert_eq!(text.matches("!cute.i64").count(), 8, "{text}"); + assert_eq!(text.matches("!cute.i64").count(), 4, "{text}"); + assert!(!text.contains("cute.scaled_view_"), "{text}"); + assert!(!text.contains("cute.tensor_make_2d"), "{text}"); + assert_eq!(text.matches("\"gpu.module\"").count(), 1, "{text}"); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/grid_constant.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/grid_constant.rs new file mode 100644 index 0000000000..ca68c0587a --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/grid_constant.rs @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! By-value kernel parameter layout and grid-constant metadata validation. + +use std::num::NonZero; + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirAttribute, render_module}; +use dialect_mir::{ + attributes::ReferenceParamValidityAttr, + ops::{MirFuncOp, MirReturnOp}, + types::{MirArrayType, MirPointerKind, MirPtrType, MirSliceType, MirStructType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, + utils::apint::APInt, +}; +use reserved_oxide_symbols::{ + MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX, MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX, +}; + +/// Build the chosen grid-constant metadata case and return rendered MLIR plus launch argument attributes. +#[allow(clippy::disallowed_methods)] // Synthetic fixture models rustc-proven reference metadata. +fn fixture(case: &str) -> Result<(String, Vec), String> { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let byte: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let word_type = IntegerType::get(&ctx, 64, Signedness::Unsigned); + let u64: TypeHandle = word_type.into(); + let bytes: TypeHandle = + MirArrayType::get(&mut ctx, byte, if case == "zero" { 0 } else { 128 }).into(); + let descriptor: TypeHandle = if case == "zero" { + bytes + } else { + MirStructType::get_with_full_layout( + &mut ctx, + "Descriptor".into(), + vec!["bytes".into()], + vec![bytes], + vec![0], + vec![0], + 128, + 64, + ) + .into() + }; + let pointer: TypeHandle = MirPtrType::get( + &mut ctx, + descriptor, + case == "mutable", + if case == "space" { 1 } else { 0 }, + ) + .into(); + let slice: TypeHandle = + MirSliceType::get_with_kind(&mut ctx, byte, MirPointerKind::SharedRef).into(); + let args = vec![ + slice, + if case == "scalar" { u64 } else { pointer }, + u64, + pointer, + ]; + let module = ModuleOp::new(&mut ctx, "grid_constants".try_into().unwrap()); + let function_type = FunctionType::get(&ctx, args.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, "copy_maps".try_into().unwrap()); + if case != "nonkernel" { + function_op.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + function.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr::new(1)); + } + for index in [1, 3] { + let suffix = if case == "noncanonical" && index == 1 { + "01".into() + } else if case == "range" && index == 1 { + "4".into() + } else { + index.to_string() + }; + let mut op = function_op.deref_mut(&ctx); + op.attributes.set( + format!("{MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX}{suffix}") + .as_str() + .try_into() + .unwrap(), + TypeAttr::new(if case == "pointee" { bytes } else { descriptor }), + ); + if case != "incomplete" || index != 1 { + op.attributes.set( + format!("{MIR_GRID_CONSTANT_ALIGN_ATTR_PREFIX}{suffix}") + .as_str() + .try_into() + .unwrap(), + IntegerAttr::new( + word_type, + APInt::from_u64( + if case == "alignment" { 3 } else { 64 }, + NonZero::new(64).unwrap(), + ), + ), + ); + } + } + module.append_operation(&mut ctx, function_op, 0); + let entry = BasicBlock::new(&mut ctx, None, args); + entry.insert_at_back(function_op.deref(&ctx).get_region(0), &ctx); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(entry, &ctx); + let profile = CutlassFullCuteMlir22::new("sm_100a").unwrap(); + let target = profile + .translate_module(&ctx, &module) + .map_err(|e| e.to_string())?; + let device = &target.root.regions[0].blocks[0].operations[0]; + let function = &device.regions[0].blocks[0].operations[0]; + let Some(MlirAttribute::Array(attributes)) = function.properties.get("arg_attrs") else { + return Err("missing argument attributes".into()); + }; + Ok(( + render_module(&target).map_err(|error| error.to_string())?, + attributes.clone(), + )) +} + +/// Check descriptor byte size, alignment, and grid-constant attributes follow the expanded argument index. +#[test] +fn descriptors_remain_by_value_and_grid_constant_after_slice_flattening() { + let (text, attrs) = fixture("valid").unwrap(); + assert_eq!(attrs.len(), 5); + for index in [0, 1, 3] { + assert_eq!(attrs[index], MlirAttribute::Dictionary(Default::default())); + } + assert_eq!(attrs[2], attrs[4]); + assert!( + text.contains("llvm.byval = !llvm.array<128 x i8>"), + "{text}" + ); + assert!(text.contains("llvm.align = 64 : i64"), "{text}"); + assert!(text.contains("nvvm.grid_constant"), "{text}"); + assert!( + !text.contains(MIR_GRID_CONSTANT_POINTEE_ATTR_PREFIX), + "{text}" + ); + assert!(!text.contains("reference_param_validity"), "{text}"); + assert!( + text.contains("llvm.insertvalue"), + "slice reconstruction must survive: {text}" + ); +} + +/// Reject incomplete, inconsistent, or illegal by-value metadata before a pointer launch ABI can escape. +#[test] +fn invalid_grid_constant_contracts_fail_before_emitting_an_ordinary_pointer_abi() { + for (case, expected) in [ + ("nonkernel", "require a kernel entry"), + ("noncanonical", "malformed or out-of-range"), + ("range", "malformed or out-of-range"), + ("incomplete", "incomplete grid-constant"), + ("alignment", "power of two"), + ("scalar", "not a pointer"), + ("mutable", "immutable generic pointer"), + ("space", "immutable generic pointer"), + ("zero", "sized, nonzero pointee"), + ("pointee", "pointee metadata differs"), + ] { + let error = fixture(case).unwrap_err(); + assert!(error.contains(expected), "{case}: {error}"); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs new file mode 100644 index 0000000000..3b394e3a79 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/nvvm_attention.rs @@ -0,0 +1,609 @@ +//! Attention math and asynchronous-copy semantics across both consumer profiles. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, CutlassPrimitivesMlir22, MlirConsumerProfile}; +use dialect_mir::ops::{MirConstantOp, MirFuncOp, MirReturnOp}; +use dialect_nvvm::ops::*; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FP64Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + utils::apint::APInt, + value::Value, +}; + +/// Typed scalar and asynchronous-copy source function for comparing both consumer profiles. +struct Fixture { + /// Source IR context owning the fixture's registered types and operations. + ctx: Context, + /// Source module containing the operations under test. + module: ModuleOp, + /// Function entry block where the fixture appends source operations. + block: Ptr, +} + +impl Fixture { + /// Build a MIR/NVVM context and function with the operands required by attention intrinsics. + fn new() -> Self { + Self::with_float_width(false) + } + + /// Select FP64 only when testing rejection of an otherwise valid scalar-math operation. + fn with_float_width(double: bool) -> Self { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let scalar = if double { + FP64Type::get(&ctx).into() + } else { + FP32Type::get(&ctx).into() + }; + let u32 = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let args = vec![u32, scalar, u32]; + let module = ModuleOp::new(&mut ctx, "attention_primitives".try_into().unwrap()); + let ty = FunctionType::get(&ctx, args.clone(), vec![]); + let raw = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, raw, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "attention_primitives".try_into().unwrap()); + module.append_operation(&mut ctx, raw, 0); + let block = BasicBlock::new(&mut ctx, None, args); + block.insert_at_back(raw.deref(&ctx).get_region(0), &ctx); + Self { ctx, module, block } + } + + /// Return the indexed source function argument for constructing a test operation. + fn arg(&self, index: usize) -> Value { + self.block.deref(&self.ctx).get_argument(index) + } + + /// Append a typed source operation to the fixture block and return its operation handle. + fn append(&mut self, op: Ptr, marker: &str) -> Ptr { + if !marker.is_empty() { + op.deref_mut(&self.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + } + op.insert_at_back(self.block, &self.ctx); + op + } + + /// Append an unsigned compile-time wait count to the source block. + fn constant(&mut self, value: u64) -> Value { + let ty = IntegerType::get(&self.ctx, 32, Signedness::Unsigned); + let op = Operation::new( + &mut self.ctx, + MirConstantOp::get_concrete_op_info(), + vec![ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(op).set_attr_value( + &self.ctx, + IntegerAttr::new(ty, APInt::from_u64(value, 32usize.try_into().unwrap())), + ); + self.append(op, ""); + op.deref(&self.ctx).get_result(0) + } + + /// Append scalar math with explicit precision, format, subnormal mode, and catalog identity. + fn math( + &mut self, + operation: ScalarMathOperationAttr, + subnormal: ScalarMathSubnormalAttr, + marker: &str, + ) -> Ptr { + let input = self.arg(1); + let op = ScalarMathOp::build( + &mut self.ctx, + input, + ScalarMathFormatAttr::F32, + operation, + ScalarMathPrecisionAttr::Approx, + subnormal, + ); + self.append(op, marker) + } + + /// Append the source function return after the tested instructions. + fn finish(&mut self) { + let op = Operation::new( + &mut self.ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + self.append(op, ""); + } + + /// Export the fixture using the default full-CuTe attention profile. + fn translate(&self) -> Result { + self.translate_with_profile(&CutlassPrimitivesMlir22::new("sm_120a").unwrap()) + } + + /// Export and render the same source fixture with the supplied consumer profile. + fn translate_with_profile(&self, profile: &dyn MlirConsumerProfile) -> Result { + profile + .translate_module(&self.ctx, &self.module) + .and_then(|m| render_module(&m).map_err(Into::into)) + .map_err(|e| e.to_string()) + } +} + +/// RMSNorm's approximate reciprocal square root preserves subnormals in each consumer profile. +#[test] +fn rsqrt_preserves_subnormals_and_native_intrinsic_identity() { + let mut f = Fixture::new(); + f.math( + ScalarMathOperationAttr::Rsqrt, + ScalarMathSubnormalAttr::Preserve, + "v1:i0802", + ); + f.finish(); + for arch in ["sm_100a", "sm_120a"] { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new(arch).unwrap(), + &CutlassFullCuteMlir22::new(arch).unwrap(), + ]; + for profile in profiles { + let text = f.translate_with_profile(profile).unwrap(); + assert!(text.contains("\"llvm.call_intrinsic\"(%v1)"), "{text}"); + assert!( + text.contains("intrin = \"llvm.nvvm.rsqrt.approx.f\""), + "{text}" + ); + assert!(text.contains(": (f32) -> (f32)"), "{text}"); + assert!(!text.contains("rsqrt.approx.ftz"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } + } +} + +/// Only the reviewed FP32, approximate, preserve-subnormal rsqrt contract is admitted. +#[test] +fn rsqrt_rejects_unsupported_variants_and_forged_identity() { + for case in 0..6 { + let mut f = Fixture::with_float_width(case == 2); + let input = f.arg(if case == 3 { 0 } else { 1 }); + let op = ScalarMathOp::build( + &mut f.ctx, + input, + if case == 2 { + ScalarMathFormatAttr::F64 + } else { + ScalarMathFormatAttr::F32 + }, + ScalarMathOperationAttr::Rsqrt, + if case == 1 { + ScalarMathPrecisionAttr::Rn + } else { + ScalarMathPrecisionAttr::Approx + }, + if case == 0 { + ScalarMathSubnormalAttr::Ftz + } else { + ScalarMathSubnormalAttr::Preserve + }, + ); + let marker = match case { + 0 | 4 => "v1:i0803", + 2 => "v1:i0804", + 5 => "", + _ => "v1:i0802", + }; + f.append(op, marker); + f.finish(); + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new("sm_120a").unwrap(), + &CutlassFullCuteMlir22::new("sm_120a").unwrap(), + ]; + for profile in profiles { + let error = f.translate_with_profile(profile).unwrap_err(); + if case >= 4 { + assert!( + error.contains("expected intrinsic marker v1:i0802"), + "{error}" + ); + } + } + } +} + +/// SnakeBeta's approximate sine preserves subnormals in each consumer profile. +#[test] +fn sin_preserves_subnormals_and_native_intrinsic_identity() { + let mut f = Fixture::new(); + f.math( + ScalarMathOperationAttr::Sin, + ScalarMathSubnormalAttr::Preserve, + "v1:i0782", + ); + f.finish(); + for arch in ["sm_100a", "sm_120a"] { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new(arch).unwrap(), + &CutlassFullCuteMlir22::new(arch).unwrap(), + ]; + for profile in profiles { + let text = f.translate_with_profile(profile).unwrap(); + assert!(text.contains("\"llvm.call_intrinsic\"(%v1)"), "{text}"); + assert!( + text.contains("intrin = \"llvm.nvvm.sin.approx.f\""), + "{text}" + ); + assert!(text.contains(": (f32) -> (f32)"), "{text}"); + assert!(!text.contains("sin.approx.ftz"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } + } +} + +/// Only the reviewed FP32, approximate, preserve-subnormal sin contract is admitted. +#[test] +fn sin_rejects_unsupported_variants_and_forged_identity() { + for case in 0..6 { + let mut f = Fixture::with_float_width(case == 2); + let input = f.arg(if case == 3 { 0 } else { 1 }); + let op = ScalarMathOp::build( + &mut f.ctx, + input, + if case == 2 { + ScalarMathFormatAttr::F64 + } else { + ScalarMathFormatAttr::F32 + }, + ScalarMathOperationAttr::Sin, + if case == 1 { + ScalarMathPrecisionAttr::Rn + } else { + ScalarMathPrecisionAttr::Approx + }, + if case == 0 { + ScalarMathSubnormalAttr::Ftz + } else { + ScalarMathSubnormalAttr::Preserve + }, + ); + let marker = match case { + 0 | 4 => "v1:i0783", + 2 => "v1:i0784", + 5 => "", + _ => "v1:i0782", + }; + f.append(op, marker); + f.finish(); + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new("sm_120a").unwrap(), + &CutlassFullCuteMlir22::new("sm_120a").unwrap(), + ]; + for profile in profiles { + let error = f.translate_with_profile(profile).unwrap_err(); + if case >= 4 { + assert!( + error.contains("expected intrinsic marker v1:i0782"), + "{error}" + ); + } + } + } +} + +/// Check both profiles retain the exact FTZ exp2 intrinsic and mode. +#[test] +fn ftz_exp2_keeps_its_exact_subnormal_mode_in_both_profiles() { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new("sm_120a").unwrap(), + &CutlassFullCuteMlir22::new("sm_120a").unwrap(), + ]; + let mut f = Fixture::new(); + f.math( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Ftz, + "v1:i0820", + ); + f.finish(); + for profile in profiles { + let text = f.translate_with_profile(profile).unwrap(); + assert!(text.contains("\"llvm.inline_asm\"(%v1)"), "{text}"); + assert!( + text.contains("asm_string = \"ex2.approx.ftz.f32 $0, $1;\""), + "{text}" + ); + assert!(text.contains("constraints = \"=f,f\""), "{text}"); + assert!(text.contains(": (f32) -> (f32)"), "{text}"); + assert!(!text.contains("ex2.approx.f32"), "{text}"); + assert!(!text.contains("has_side_effects"), "{text}"); + assert!(!text.contains("llvm.call"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } +} + +/// Reject an FTZ exp2 operation mislabeled with another canonical identity. +#[test] +fn ftz_exp2_rejects_other_canonical_identities_in_both_profiles() { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new("sm_120a").unwrap(), + &CutlassFullCuteMlir22::new("sm_120a").unwrap(), + ]; + for marker in ["v1:i0819", "v1:i0788", ""] { + let mut f = Fixture::new(); + f.math( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Ftz, + marker, + ); + f.finish(); + for profile in profiles { + let error = f.translate_with_profile(profile).unwrap_err(); + assert!( + error.contains("expected intrinsic marker v1:i0820"), + "{error}" + ); + } + } +} + +/// Reject FTZ exp2 operands or precision selectors outside its reviewed contract. +#[test] +fn ftz_exp2_rejects_wrong_type_and_selectors() { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new("sm_120a").unwrap(), + &CutlassFullCuteMlir22::new("sm_120a").unwrap(), + ]; + for case in 0..5 { + let mut f = Fixture::new(); + let input = f.arg(if case == 0 { 0 } else { 1 }); + let op = ScalarMathOp::build( + &mut f.ctx, + input, + if case == 1 { + ScalarMathFormatAttr::F64 + } else { + ScalarMathFormatAttr::F32 + }, + if case == 2 { + ScalarMathOperationAttr::Lg2 + } else { + ScalarMathOperationAttr::Ex2 + }, + if case == 3 { + ScalarMathPrecisionAttr::Rn + } else { + ScalarMathPrecisionAttr::Approx + }, + if case == 4 { + ScalarMathSubnormalAttr::Preserve + } else { + ScalarMathSubnormalAttr::Ftz + }, + ); + f.append(op, "v1:i0820"); + f.finish(); + for profile in profiles { + assert!(f.translate_with_profile(profile).is_err(), "case {case}"); + } + } +} + +/// Check softmax exp/reciprocal operations retain their native approximation and subnormal semantics. +#[test] +fn softmax_math_keeps_approximation_and_subnormal_modes() { + let mut f = Fixture::new(); + f.math( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Preserve, + "v1:i0819", + ); + f.math( + ScalarMathOperationAttr::Rcp, + ScalarMathSubnormalAttr::Ftz, + "v1:i0788", + ); + f.finish(); + let text = f.translate().unwrap(); + assert!(text.contains("\"llvm.inline_asm\"(%v1)"), "{text}"); + assert!( + text.contains("asm_string = \"ex2.approx.f32 $0, $1;\""), + "{text}" + ); + assert!(text.contains("constraints = \"=f,f\""), "{text}"); + assert!(!text.contains("ex2.approx.ftz"), "{text}"); + assert!(!text.contains("llvm.call"), "{text}"); + assert!( + text.contains("\"nvvm.rcp.approx.ftz.f\"(%v1) : (f32) -> (f32)"), + "{text}" + ); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); +} + +/// Reject unreviewed math modes and fabricated intrinsic identities. +#[test] +fn unsupported_math_and_forged_identity_are_rejected() { + let mut f = Fixture::new(); + let op = f.math( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Preserve, + "v1:i0788", + ); + f.finish(); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i0819") + ); + ScalarMathOp::new(op).set_attr_nvvm_scalar_math_operation(&f.ctx, ScalarMathOperationAttr::Cos); + assert!( + f.translate() + .unwrap_err() + .contains("supports F32 approximate exp2 and FTZ reciprocal") + ); +} + +/// Check the shuffle keeps its membership mask, lane XOR offset, and full-warp clamp. +#[test] +fn butterfly_shuffle_preserves_mask_offset_and_warp_clamp() { + let mut f = Fixture::new(); + let (mask, value, offset) = (f.arg(0), f.arg(1), f.arg(2)); + let op = ShflSyncBflyF32Op::build(&mut f.ctx, mask, value, offset); + f.append(op, "v1:i0055"); + f.finish(); + let text = f.translate().unwrap(); + assert!(text.contains("value = 31 : i32"), "{text}"); + assert!(text.contains("\"nvvm.shfl.sync\"(%v0, %v1, %v2,"), "{text}"); + assert!(text.contains("kind = #nvvm"), "{text}"); + assert!(!text.contains("return_value_and_is_valid")); + op.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0054".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i0055") + ); +} + +/// Check asynchronous-copy grouping preserves explicit commits and constant wait counts. +#[test] +fn cp_async_commit_and_wait_keep_the_group_count() { + let mut f = Fixture::new(); + let commit = CpAsyncCommitGroupOp::build(&mut f.ctx); + f.append(commit, "v1:i0094"); + let count = f.constant(3); + let wait = CpAsyncWaitGroupOp::build(&mut f.ctx, count); + f.append(wait, "v1:i0096"); + f.finish(); + let text = f.translate().unwrap(); + let commit_pos = text + .find("\"nvvm.cp.async.commit.group\"() : () -> ()") + .unwrap(); + let wait_pos = text + .find("\"nvvm.cp.async.wait.group\"() <{n = 3 : i32}> : () -> ()") + .unwrap(); + assert!(commit_pos < wait_pos, "{text}"); + wait.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0095".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i0096") + ); +} + +/// Reject wait selectors that cannot be represented by the hardware immediate contract. +#[test] +fn cp_async_wait_rejects_runtime_or_out_of_range_counts() { + for count in [None, Some(8)] { + let mut f = Fixture::new(); + let count = match count { + Some(n) => f.constant(n), + None => f.arg(0), + }; + let wait = CpAsyncWaitGroupOp::build(&mut f.ctx, count); + f.append(wait, "v1:i0096"); + f.finish(); + assert!(f.translate().is_err()); + } +} + +/// Check scalar numeric and asynchronous-copy registrations compose without duplicate ownership. +#[test] +fn numeric_and_async_packs_coexist_on_both_architectures() { + let mut f = Fixture::new(); + for (operation, subnormal, marker) in [ + ( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Preserve, + "v1:i0819", + ), + ( + ScalarMathOperationAttr::Ex2, + ScalarMathSubnormalAttr::Ftz, + "v1:i0820", + ), + ( + ScalarMathOperationAttr::Lg2, + ScalarMathSubnormalAttr::Preserve, + "v1:i0786", + ), + ( + ScalarMathOperationAttr::Lg2, + ScalarMathSubnormalAttr::Ftz, + "v1:i0787", + ), + ] { + f.math(operation, subnormal, marker); + } + let (mask, value, offset) = (f.arg(0), f.arg(1), f.arg(2)); + let shuffle = ShflSyncBflyF32Op::build(&mut f.ctx, mask, value, offset); + f.append(shuffle, "v1:i0055"); + let commit = CpAsyncCommitGroupOp::build(&mut f.ctx); + f.append(commit, "v1:i0094"); + let count = f.constant(3); + let wait = CpAsyncWaitGroupOp::build(&mut f.ctx, count); + f.append(wait, "v1:i0096"); + f.finish(); + + for arch in ["sm_100a", "sm_120a"] { + let profiles: [&dyn MlirConsumerProfile; 2] = [ + &CutlassPrimitivesMlir22::new(arch).unwrap(), + &CutlassFullCuteMlir22::new(arch).unwrap(), + ]; + for profile in profiles { + let text = f.translate_with_profile(profile).unwrap(); + for name in [ + "llvm.nvvm.lg2.approx.f", + "llvm.nvvm.lg2.approx.ftz.f", + "nvvm.cp.async.commit.group", + "nvvm.cp.async.wait.group", + ] { + assert!(text.contains(name), "{arch}: {text}"); + } + assert!(text.contains("n = 3 : i32"), "{arch}: {text}"); + assert!(text.contains("value = 31 : i32"), "{arch}: {text}"); + if arch == "sm_100a" { + for name in [ + "llvm.nvvm.ex2.approx.f32", + "llvm.nvvm.ex2.approx.ftz.f32", + "llvm.nvvm.shfl.sync.bfly.f32", + ] { + assert!(text.contains(name), "{text}"); + } + assert!(!text.contains("llvm.inline_asm"), "{text}"); + } else { + for name in [ + "ex2.approx.f32 $0, $1;", + "ex2.approx.ftz.f32 $0, $1;", + "nvvm.shfl.sync", + ] { + assert!(text.contains(name), "{text}"); + } + assert!(!text.contains("llvm.nvvm.ex2"), "{text}"); + assert!(!text.contains("llvm.nvvm.shfl.sync.bfly"), "{text}"); + } + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/primitives_profile.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/primitives_profile.rs new file mode 100644 index 0000000000..94ef86dba1 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/primitives_profile.rs @@ -0,0 +1,178 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Primitives-only profile isolation, source verification, and reproducibility metadata. + +use cuda_oxide_mlir_export::{CutlassPrimitivesMlir22, MlirConsumerProfile}; +use pliron::builtin::op_interfaces::SingleBlockRegionInterface; +use pliron::{ + builtin::ops::ModuleOp, context::Context, identifier::Identifier, op::Op, printable::Printable, +}; + +/// Reject CuTe pointees even when the runtime carrier would become an opaque pointer. +#[test] +fn primitives_rejects_cute_hidden_in_pointer_type_metadata() { + use dialect_cute::{attributes::CuteTileGridAttr, types::CuteWorkTileType}; + use dialect_mir::types::{MirPtrType, MirTupleType}; + use pliron::builtin::attributes::{TypeAttr, VecAttr}; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "nested".try_into().unwrap()); + let cute = CuteWorkTileType::get(&mut ctx, CuteTileGridAttr::new(1, 1, 1)); + let tuple = MirTupleType::get(&mut ctx, vec![cute.into()]); + let pointer = MirPtrType::get_generic(&mut ctx, tuple.into(), false); + module.get_operation().deref_mut(&ctx).attributes.set( + "metadata".try_into().unwrap(), + VecAttr(vec![Box::new(TypeAttr::new(pointer.into()))]), + ); + let error = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .verify_source_module(&ctx, &module) + .unwrap_err() + .to_string(); + assert!(error.contains("cute.work_tile"), "{error}"); +} + +/// Check ordinary MIR exports with a primitives-only registry. +#[test] +fn primitives_profile_translates_without_cute_registration() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, Identifier::try_from("primitives").unwrap()); + let before = module.disp(&ctx).to_string(); + let profile = CutlassPrimitivesMlir22::new("sm_100a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + assert_eq!( + target.root.regions[0].blocks[0].operations[0].name, + "gpu.module" + ); + assert_eq!(module.disp(&ctx).to_string(), before); + assert!(profile.build_registry().unwrap().is_sealed()); +} + +/// Check profile identity, compiler revision, and architecture are recorded in the manifest. +#[test] +fn primitives_profile_pins_consumer_and_target() { + let profile = CutlassPrimitivesMlir22::new("sm_100a").unwrap(); + let manifest = profile.manifest(); + assert_eq!(manifest.profile, "cutlass-primitives-mlir22"); + assert_eq!(manifest.consumer_revision, "cutlass-compiler-v4.7.0"); + assert_eq!(manifest.llvm_major, 22); + assert_eq!( + manifest.pliron_revision, + "435536c8611af34c48c8f73d78e77894a05d103d" + ); + assert_eq!( + manifest.interop_revision, + cuda_oxide_mlir_export::INTEROP_REVISION + ); + assert_eq!( + CutlassPrimitivesMlir22::new("sm_120a") + .unwrap() + .manifest() + .gpu_arch, + "sm_120a" + ); + for arch in ["sm_90a", "sm_120", "sm_100", "garbage"] { + assert!(CutlassPrimitivesMlir22::new(arch).is_err()); + } +} + +/// Reject a CuTe operation before its own verifier can obscure the profile violation. +#[test] +fn primitives_rejects_cute_operations_before_local_verification() { + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "cute_op".try_into().unwrap()); + let op = pliron::operation::Operation::new( + &mut ctx, + dialect_cute::ops::CuteCopyG2SOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + module.append_operation(&mut ctx, op, 0); + let error = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .verify_source_module(&ctx, &module) + .unwrap_err() + .to_string(); + assert!(error.contains("CuTe operation"), "{error}"); +} + +/// Find forbidden CuTe metadata recursively inside nested attribute dictionaries. +#[test] +fn primitives_rejects_cute_attributes_in_dictionaries() { + use pliron::{attribute::AttributeDict, builtin::attributes::DictAttr}; + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "cute_attr".try_into().unwrap()); + let mut dictionary = AttributeDict::default(); + dictionary.set( + "grid".try_into().unwrap(), + dialect_cute::attributes::CuteTileGridAttr::new(1, 1, 1), + ); + module + .get_operation() + .deref_mut(&ctx) + .attributes + .set("metadata".try_into().unwrap(), DictAttr(dictionary)); + let error = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .verify_source_module(&ctx, &module) + .unwrap_err() + .to_string(); + assert!(error.contains("cute.tile_grid"), "{error}"); +} + +/// Reject unreviewed NVVM instructions instead of translating them by name alone. +#[test] +fn primitives_rejects_unmapped_nvvm_instructions() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "unmapped".try_into().unwrap()); + let op = pliron::operation::Operation::new( + &mut ctx, + dialect_nvvm::ops::Barrier0Op::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + module.append_operation(&mut ctx, op, 0); + let error = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(&ctx, &module) + .unwrap_err() + .to_string(); + assert!(error.contains("nvvm.barrier0"), "{error}"); +} + +/// Check the primitives profile also enforces ordinary source operation validity. +#[test] +fn primitives_still_checks_ordinary_ir_verification() { + let mut ctx = Context::new(); + dialect_nvvm::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "invalid".try_into().unwrap()); + // A special-register read must produce one result. + let op = pliron::operation::Operation::new( + &mut ctx, + dialect_nvvm::ops::ReadPtxSregTidXOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + module.append_operation(&mut ctx, op, 0); + assert!( + CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .verify_source_module(&ctx, &module) + .is_err() + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/register_storage.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/register_storage.rs new file mode 100644 index 0000000000..660cb55a81 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/register_storage.rs @@ -0,0 +1,558 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Rust array construction needs byte-preserving unions, slices and aggregate casts. + +// These fixtures intentionally forge pointer contracts to test verifier rejection. +#![allow(clippy::disallowed_methods)] + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirAttribute, MlirOperation, MlirType, render_module}; +use dialect_mir::{ + attributes::{FieldIndexAttr, MirCastKindAttr, MirPointerKindAuthorityAttr}, + ops::*, + types::*, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness, UnitType}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::TypeHandle, + value::Value, +}; + +/// Create a source context with the dialects used by the full CuTe profile. +fn context() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + ctx +} + +/// Append an operation and return its handle for setting semantic attributes. +fn append( + ctx: &mut Context, + block: Ptr, + types: Vec, + values: Vec, +) -> Ptr { + let operation = Operation::new(ctx, O::get_concrete_op_info(), types, values, vec![], 0); + operation.insert_at_back(block, ctx); + operation +} + +/// Create one ordinary MIR function in a module. +fn function( + ctx: &mut Context, + args: Vec, + results: Vec, +) -> (ModuleOp, Ptr) { + let module = ModuleOp::new(ctx, "register_storage".try_into().unwrap()); + let signature = FunctionType::get(ctx, args.clone(), results); + let function = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(ctx, function, TypeAttr::new(signature.into())) + .set_symbol_name(ctx, "test".try_into().unwrap()); + module.append_operation(ctx, function, 0); + let block = BasicBlock::new(ctx, None, args); + block.insert_at_back(function.deref(ctx).get_region(0), ctx); + (module, block) +} + +/// Translate through the production SM120 CuTe profile. +fn render(ctx: &Context, module: &ModuleOp) -> Result { + CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(ctx, module) + .and_then(|target| render_module(&target).map_err(Into::into)) + .map_err(|error| error.to_string()) +} + +/// Model the standard library's uninitialized-or-value storage layout. +fn maybe_uninit(ctx: &mut Context, value: TypeHandle, size: u64, align: u64) -> TypeHandle { + let unit = UnitType::get(ctx).into(); + MirUnionType::get( + ctx, + "MaybeUninit".into(), + vec!["uninit".into(), "value".into()], + vec![unit, value], + size, + align, + ) + .into() +} + +/// Field one of a union overlaps field zero; it is not LLVM aggregate slot one. +#[test] +fn union_value_insert_and_extract_share_byte_zero() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let union = maybe_uninit(&mut ctx, scalar, 4, 4); + let (module, block) = function(&mut ctx, vec![scalar], vec![scalar]); + let argument = block.deref(&ctx).get_argument(0); + let undef = append::(&mut ctx, block, vec![union], vec![]); + let empty = undef.deref(&ctx).get_result(0); + let insert = append::(&mut ctx, block, vec![union], vec![empty, argument]); + MirInsertFieldOp::new(insert).set_attr_insert_index(&ctx, FieldIndexAttr(1)); + let initialized = insert.deref(&ctx).get_result(0); + let extract = append::(&mut ctx, block, vec![scalar], vec![initialized]); + MirExtractFieldOp::new(extract).set_attr_index(&ctx, FieldIndexAttr(1)); + let result = extract.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let text = render(&ctx, &module).unwrap(); + assert!(text.contains("array<4 x i8>"), "{text}"); + assert!(text.contains("array<0 x i32>"), "{text}"); + assert!( + text.contains("\"llvm.store\"") && text.contains("\"llvm.load\""), + "{text}" + ); + assert!(!text.contains("position = array"), "{text}"); +} + +/// `array::from_fn` reinterprets initialized storage without numeric conversion. +#[test] +fn initialized_array_transmute_preserves_storage_and_alignment() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let union = maybe_uninit(&mut ctx, scalar, 4, 4); + let source = MirArrayType::get(&mut ctx, union, 4).into(); + let target = MirArrayType::get(&mut ctx, scalar, 4).into(); + let (module, block) = function(&mut ctx, vec![source], vec![target]); + let argument = block.deref(&ctx).get_argument(0); + let cast = append::(&mut ctx, block, vec![target], vec![argument]); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let result = cast.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let text = render(&ctx, &module).unwrap(); + assert!(text.contains("!llvm.array<4 x f32>"), "{text}"); + assert!( + text.contains("\"llvm.store\"") && text.contains("\"llvm.load\""), + "{text}" + ); + assert!( + !text.contains("arith.sitofp") && !text.contains("arith.fptosi"), + "{text}" + ); +} + +/// The mutable slice used by array initialization preserves both carrier fields. +#[test] +fn constructed_slice_keeps_pointer_and_length() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let union = maybe_uninit(&mut ctx, scalar, 4, 4); + let pointer = + MirPtrType::get_with_kind(&mut ctx, union, true, 0, MirPointerKind::RawMut).into(); + let slice = MirSliceType::get_with_kind(&mut ctx, union, MirPointerKind::RawMut).into(); + let length = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let (module, block) = function(&mut ctx, vec![pointer, length], vec![slice]); + let args = block.deref(&ctx).arguments().collect(); + let construct = append::(&mut ctx, block, vec![slice], args); + let result = construct.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let text = render(&ctx, &module).unwrap(); + assert_eq!(text.matches("\"llvm.insertvalue\"").count(), 2, "{text}"); + assert!(text.contains("position = array"), "{text}"); + assert!(text.contains("position = array"), "{text}"); + assert!(text.contains("!llvm.struct<(!llvm.ptr, i64)>"), "{text}"); +} + +/// Array-to-slice coercion retains the original pointer and supplies the static extent. +#[test] +fn array_unsizing_preserves_extent_and_pointer_kind() { + for (length, kind) in [ + (0, MirPointerKind::SharedRef), + (7, MirPointerKind::UniqueRef), + ] { + let mut ctx = context(); + let element = FP32Type::get(&ctx).into(); + let array = MirArrayType::get(&mut ctx, element, length).into(); + let mutable = kind == MirPointerKind::UniqueRef; + let pointer = MirPtrType::get_with_kind(&mut ctx, array, mutable, 0, kind).into(); + let slice = MirSliceType::get_with_kind(&mut ctx, element, kind).into(); + let (module, block) = function(&mut ctx, vec![pointer], vec![slice]); + let argument = block.deref(&ctx).get_argument(0); + let cast = append::(&mut ctx, block, vec![slice], vec![argument]); + let semantic = MirCastOp::new(cast); + semantic.set_attr_cast_kind(&ctx, MirCastKindAttr::PointerCoercionUnsize); + semantic.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + let result = cast.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let text = render(&ctx, &module).unwrap(); + assert_eq!(text.matches("\"llvm.insertvalue\"").count(), 2, "{text}"); + assert!(text.contains(&format!("value = {length} : i64")), "{text}"); + assert!( + !text.contains("llvm.alloca") && !text.contains("llvm.load"), + "{text}" + ); + } +} + +/// Coercion cannot erase a named address space or strengthen a borrow. +#[test] +fn array_unsizing_rejects_incompatible_pointer_contracts() { + for (space, target_kind, wrong_element, not_array) in [ + (3, MirPointerKind::SharedRef, false, false), + (0, MirPointerKind::UniqueRef, false, false), + (0, MirPointerKind::SharedRef, true, false), + (0, MirPointerKind::SharedRef, false, true), + ] { + let mut ctx = context(); + let element = FP32Type::get(&ctx).into(); + let pointee = if not_array { + element + } else { + MirArrayType::get(&mut ctx, element, 7).into() + }; + let pointer = + MirPtrType::get_with_kind(&mut ctx, pointee, false, space, MirPointerKind::SharedRef) + .into(); + let target_element = if wrong_element { + MirFP16Type::get(&ctx).into() + } else { + element + }; + let slice = MirSliceType::get_with_kind(&mut ctx, target_element, target_kind).into(); + let (module, block) = function(&mut ctx, vec![pointer], vec![slice]); + let argument = block.deref(&ctx).get_argument(0); + let cast = append::(&mut ctx, block, vec![slice], vec![argument]); + let semantic = MirCastOp::new(cast); + semantic.set_attr_cast_kind(&ctx, MirCastKindAttr::PointerCoercionUnsize); + semantic.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RustCast); + let result = cast.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + assert!(render(&ctx, &module).is_err()); + } +} + +/// Infallible array construction retains an uninhabited residual in unreachable IR. +#[test] +fn empty_residual_has_zero_storage_but_no_discriminant() { + let mut ctx = context(); + let tag = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let empty = MirEnumType::get_with_encoding( + &mut ctx, + "NeverShortCircuitResidual".into(), + tag, + vec![], + vec![], + EnumEncoding { + abi_align: 1, + layout_kind: EnumLayoutKind::Empty, + ..Default::default() + }, + ) + .into(); + let unit = UnitType::get(&ctx).into(); + let empty_flow = MirEnumType::get_with_encoding( + &mut ctx, + "ControlFlow".into(), + tag, + vec![0, 1], + vec![ + EnumVariant::new_with_layout("Continue".into(), vec![unit], vec![0], vec![0]), + EnumVariant::new_with_layout("Break".into(), vec![empty], vec![0], vec![0]), + ], + EnumEncoding { + abi_align: 1, + layout_kind: EnumLayoutKind::Empty, + variant_inhabited: vec![0, 0], + ..Default::default() + }, + ) + .into(); + let (flow_module, flow_block) = function(&mut ctx, vec![empty_flow], vec![]); + append::(&mut ctx, flow_block, vec![], vec![]); + assert!(render(&ctx, &flow_module).is_ok()); + let (module, block) = function(&mut ctx, vec![empty], vec![]); + append::(&mut ctx, block, vec![], vec![]); + let text = render(&ctx, &module).unwrap(); + assert!(text.contains("array<0 x i8>"), "{text}"); + let (module, block) = function(&mut ctx, vec![empty], vec![tag]); + let value = block.deref(&ctx).get_argument(0); + let discriminant = append::(&mut ctx, block, vec![tag], vec![value]); + let result = discriminant.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + assert!(render(&ctx, &module).is_err()); +} + +/// Reject unreviewed byte representations and malformed ABI metadata. +#[test] +fn union_storage_rejects_bool_pointer_and_invalid_layouts() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let flag = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let pointer = MirPtrType::get_generic(&mut ctx, scalar, false).into(); + for (value, size, align) in [ + (flag, 1, 1), + (pointer, 8, 8), + (scalar, 2, 2), + (scalar, 4, 3), + ] { + let union = maybe_uninit(&mut ctx, value, size, align); + let (module, block) = function(&mut ctx, vec![union], vec![union]); + let argument = block.deref(&ctx).get_argument(0); + append::(&mut ctx, block, vec![], vec![argument]); + assert!( + render(&ctx, &module).is_err(), + "accepted invalid union size={size} align={align}" + ); + } +} + +/// Aggregate reinterpretation must not read past storage or bypass byte adapters. +#[test] +fn aggregate_transmute_rejects_size_mismatch_and_bool_storage() { + let mut ctx = context(); + let word = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let flag = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let byte = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let words = MirArrayType::get(&mut ctx, word, 4).into(); + let short = MirArrayType::get(&mut ctx, word, 3).into(); + let flags = MirArrayType::get(&mut ctx, flag, 4).into(); + let bytes = MirArrayType::get(&mut ctx, byte, 4).into(); + for (source, target) in [(words, short), (flags, bytes), (bytes, flags)] { + let (module, block) = function(&mut ctx, vec![source], vec![target]); + let argument = block.deref(&ctx).get_argument(0); + let cast = append::(&mut ctx, block, vec![target], vec![argument]); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let result = cast.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + assert!(render(&ctx, &module).is_err()); + } +} + +/// Return the exported function so storage tests can check operand identity and order. +fn exported_function(ctx: &Context, module: &ModuleOp) -> MlirOperation { + fn find(operation: &MlirOperation) -> Option<&MlirOperation> { + if operation.name == "func.func" { + return Some(operation); + } + operation + .regions + .iter() + .flat_map(|region| ®ion.blocks) + .flat_map(|block| &block.operations) + .find_map(find) + } + let target = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(ctx, module) + .unwrap(); + find(&target.root) + .expect("fixture contains one function") + .clone() +} + +/// Taking either union field's address must retain its space and overlap at byte zero. +#[test] +fn union_field_addresses_overlap_without_changing_address_space() { + for address_space in [0, 3] { + let mut ctx = context(); + let wide = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let narrow = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let union = MirUnionType::get( + &mut ctx, + "OverlappingWords".into(), + vec!["wide".into(), "narrow".into()], + vec![wide, narrow], + 8, + 8, + ) + .into(); + let pointer = MirPtrType::get(&mut ctx, union, true, address_space).into(); + let wide_pointer = MirPtrType::get(&mut ctx, wide, true, address_space).into(); + let narrow_pointer = MirPtrType::get(&mut ctx, narrow, true, address_space).into(); + let (module, block) = function(&mut ctx, vec![pointer], vec![wide_pointer, narrow_pointer]); + let base = block.deref(&ctx).get_argument(0); + let mut results = Vec::new(); + for (index, result_type) in [wide_pointer, narrow_pointer].into_iter().enumerate() { + let address = MirFieldAddrOp::build(&mut ctx, base, result_type, index as u32).unwrap(); + address.insert_at_back(block, &ctx); + results.push(address.deref(&ctx).get_result(0)); + } + append::(&mut ctx, block, vec![], results); + + let target = exported_function(&ctx, &module); + let entry = &target.regions[0].blocks[0]; + let addresses: Vec<_> = entry + .operations + .iter() + .filter(|operation| operation.name == "llvm.getelementptr") + .collect(); + assert_eq!(addresses.len(), 2); + for address in &addresses { + assert_eq!(address.operands[0].id, entry.arguments[0].id); + assert_ne!(address.results[0].id, entry.arguments[0].id); + assert_eq!(address.results[0].ty, entry.arguments[0].ty); + assert_eq!( + address.properties.get("elem_type"), + Some(&MlirAttribute::Type(MlirType::Integer(8))), + ); + assert_eq!( + address.properties.get("rawConstantIndices"), + Some(&MlirAttribute::DenseI32Array(vec![0])), + ); + } + assert_ne!(addresses[0].results[0].id, addresses[1].results[0].id); + assert!( + !entry + .operations + .iter() + .any(|operation| operation.name == "llvm.addrspacecast") + ); + } +} + +/// A short field write retains the high bytes of the existing wider union value. +#[test] +fn partial_union_overwrite_preserves_untouched_bytes() { + let mut ctx = context(); + let wide = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let narrow = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let union = MirUnionType::get( + &mut ctx, + "PartialOverwrite".into(), + vec!["wide".into(), "narrow".into()], + vec![wide, narrow], + 8, + 8, + ) + .into(); + let (module, block) = function(&mut ctx, vec![union, narrow], vec![wide]); + let args: Vec<_> = block.deref(&ctx).arguments().collect(); + let insert = append::(&mut ctx, block, vec![union], args); + MirInsertFieldOp::new(insert).set_attr_insert_index(&ctx, FieldIndexAttr(1)); + let updated = insert.deref(&ctx).get_result(0); + let extract = append::(&mut ctx, block, vec![wide], vec![updated]); + MirExtractFieldOp::new(extract).set_attr_index(&ctx, FieldIndexAttr(0)); + let value = extract.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![value]); + + let target = exported_function(&ctx, &module); + let entry = &target.regions[0].blocks[0]; + let memory: Vec<_> = entry + .operations + .iter() + .filter(|operation| matches!(operation.name.as_str(), "llvm.store" | "llvm.load")) + .collect(); + assert_eq!( + memory + .iter() + .map(|operation| operation.name.as_str()) + .collect::>(), + [ + "llvm.store", + "llvm.store", + "llvm.load", + "llvm.store", + "llvm.load" + ], + ); + // Preserve the complete original carrier before overwriting its low word. + assert_eq!(memory[0].operands[0].id, entry.arguments[0].id); + assert_eq!(memory[1].operands[0].id, entry.arguments[1].id); + assert_eq!(memory[1].operands[0].ty, MlirType::Integer(32)); + assert_eq!(memory[0].operands[1], memory[1].operands[1]); + assert_eq!(memory[2].operands[0], memory[1].operands[1]); + // The wider read must observe that updated carrier, including untouched bytes. + assert_eq!(memory[3].operands[0].id, memory[2].results[0].id); + assert_eq!(memory[4].operands[0], memory[3].operands[1]); + assert_eq!(memory[4].results[0].ty, MlirType::Integer(64)); +} + +/// Nesting unsupported storage inside an array, tuple or struct cannot hide it. +#[test] +fn byte_storage_rejects_nested_pointer_and_bool_payloads() { + let mut ctx = context(); + let word = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let flag = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let byte = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let pointer = MirPtrType::get_generic(&mut ctx, word, false).into(); + let pointers = MirArrayType::get(&mut ctx, pointer, 1).into(); + let tuple = MirTupleType::get(&mut ctx, vec![pointers]).into(); + let structure = MirStructType::get( + &mut ctx, + "PointerWrapper".into(), + vec!["pointer".into()], + vec![pointer], + ) + .into(); + let flags = MirArrayType::get(&mut ctx, flag, 1).into(); + let nested_flags = MirArrayType::get(&mut ctx, flags, 1).into(); + let bytes = MirArrayType::get(&mut ctx, byte, 1).into(); + let words = MirArrayType::get(&mut ctx, word, 1).into(); + for (payload, target, size, align, reason) in [ + (pointers, words, 8, 8, "pointer"), + (tuple, words, 8, 8, "pointer"), + (structure, words, 8, 8, "pointer"), + (nested_flags, bytes, 1, 1, "bool"), + ] { + let union = maybe_uninit(&mut ctx, payload, size, align); + let (module, block) = function(&mut ctx, vec![union], vec![union]); + let argument = block.deref(&ctx).get_argument(0); + append::(&mut ctx, block, vec![], vec![argument]); + let error = render(&ctx, &module).unwrap_err(); + assert!(error.contains(reason), "{error}"); + + for (source, destination) in [(payload, target), (target, payload)] { + let (module, block) = function(&mut ctx, vec![source], vec![destination]); + let argument = block.deref(&ctx).get_argument(0); + let cast = append::(&mut ctx, block, vec![destination], vec![argument]); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::Transmute); + let result = cast.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let error = render(&ctx, &module).unwrap_err(); + assert!(error.contains(reason), "{error}"); + } + } +} + +/// Slice construction cannot erase address spaces or invent a compatible element/kind/length. +#[test] +fn constructed_slice_rejects_incompatible_pointer_and_length() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let other = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let length64 = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let length32 = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let pointer = + MirPtrType::get_with_kind(&mut ctx, scalar, true, 0, MirPointerKind::RawMut).into(); + let shared = + MirPtrType::get_with_kind(&mut ctx, scalar, true, 3, MirPointerKind::RawMut).into(); + let wrong_element = + MirPtrType::get_with_kind(&mut ctx, other, true, 0, MirPointerKind::RawMut).into(); + let immutable = + MirPtrType::get_with_kind(&mut ctx, scalar, false, 0, MirPointerKind::RawConst).into(); + let slice = MirSliceType::get_with_kind(&mut ctx, scalar, MirPointerKind::RawMut).into(); + for (pointer, length, reason) in [ + (shared, length64, "address space"), + (wrong_element, length64, "pointee mismatch"), + (immutable, length64, "pointer kind"), + (pointer, length32, "64-bit length"), + (pointer, scalar, "integer type"), + ] { + let (module, block) = function(&mut ctx, vec![pointer, length], vec![slice]); + let arguments = block.deref(&ctx).arguments().collect(); + let construct = append::(&mut ctx, block, vec![slice], arguments); + let result = construct.deref(&ctx).get_result(0); + append::(&mut ctx, block, vec![], vec![result]); + let error = render(&ctx, &module).unwrap_err(); + assert!(error.contains(reason), "{error}"); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/rust_float_max.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/rust_float_max.rs new file mode 100644 index 0000000000..071c864e04 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/rust_float_max.rs @@ -0,0 +1,180 @@ +//! Rust FP32 maximum semantics and rejection of incompatible intrinsic signatures. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, CutlassPrimitivesMlir22, MlirConsumerProfile}; +use dialect_mir::{ + ops::{MirCallOp, MirFuncOp, MirReturnOp}, + rust_intrinsics::CALLEE_MAXNUM_NSZ_F32, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FP64Type, FunctionType, IntegerType, Signedness}, + }, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +/// Build and export a scalar call with independently controlled operand and result signatures. +fn translate_call( + profile: &dyn MlirConsumerProfile, + callee: &str, + operand_types: &[&str], + result_types: &[&str], +) -> Result { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let ty = |name| -> TypeHandle { + match name { + "f32" => FP32Type::get(&ctx).into(), + "f64" => FP64Type::get(&ctx).into(), + "i32" => IntegerType::get(&ctx, 32, Signedness::Signed).into(), + _ => unreachable!(), + } + }; + let operands = operand_types + .iter() + .map(|name| ty(*name)) + .collect::>(); + let results = result_types + .iter() + .map(|name| ty(*name)) + .collect::>(); + let function_type = FunctionType::get(&ctx, operands.clone(), results.clone()); + let module = ModuleOp::new(&mut ctx, "rust_float_max".try_into().unwrap()); + let function = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())) + .set_symbol_name(&mut ctx, "user_max".try_into().unwrap()); + module.append_operation(&mut ctx, function, 0); + let block = BasicBlock::new(&mut ctx, None, operands); + block.insert_at_back(function.deref(&ctx).get_region(0), &ctx); + let arguments = block.deref(&ctx).arguments().collect(); + let call = Operation::new( + &mut ctx, + MirCallOp::get_concrete_op_info(), + results, + arguments, + vec![], + 0, + ); + MirCallOp::new(call).set_attr_callee(&ctx, StringAttr::new(callee.into())); + call.insert_at_back(block, &ctx); + let values = call.deref(&ctx).results().collect(); + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + values, + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + profile + .translate_module(&ctx, &module) + .and_then(|module| render_module(&module).map_err(Into::into)) + .map_err(|error| error.to_string()) +} + +/// Create both reviewed profiles for the same SM100 target. +fn profiles() -> [Box; 2] { + [ + Box::new(CutlassFullCuteMlir22::new("sm_120a").unwrap()), + Box::new(CutlassPrimitivesMlir22::new("sm_120a").unwrap()), + ] +} + +/// Check both profiles preserve Rust max's NaN-ignoring and non-FTZ behavior. +#[test] +fn rust_f32_max_exports_nan_ignoring_subnormal_preserving_native_maximum() { + for profile in profiles() { + let text = + translate_call(&*profile, CALLEE_MAXNUM_NSZ_F32, &["f32", "f32"], &["f32"]).unwrap(); + // CUTLASS 4.7 cute/arch/nvvm_wrappers.py::fmax uses the modifier-free + // nvvm.fmax for maximumNumber: ignore qNaN/sNaN, preserve subnormals. + // PTX's preference for +0 over -0 is allowed by Rust's nsz contract. + assert!( + text.contains("\"nvvm.fmax\"(%v0, %v1) : (f32, f32) -> (f32)"), + "{}: {text}", + profile.name() + ); + assert!( + !text.contains("\"func.call\"") && !text.contains(CALLEE_MAXNUM_NSZ_F32), + "{text}" + ); + assert!( + text.contains("\"func.return\"(%v2) : (f32) -> ()"), + "{text}" + ); + } +} + +/// Reject max calls whose argument/result arity differs from the reviewed intrinsic. +#[test] +fn rust_f32_max_rejects_malformed_operand_and_result_counts() { + for profile in profiles() { + for (operands, results) in [ + (vec![], vec!["f32"]), + (vec!["f32"], vec!["f32"]), + (vec!["f32", "f32", "f32"], vec!["f32"]), + (vec!["f32", "f32"], vec![]), + (vec!["f32", "f32"], vec!["f32", "f32"]), + ] { + let result = translate_call(&*profile, CALLEE_MAXNUM_NSZ_F32, &operands, &results); + assert!( + result.is_err(), + "{} accepted {operands:?} -> {results:?}: {result:?}", + profile.name() + ); + } + } +} + +/// Reject max calls that try to reuse the FP32 mapping for another scalar format. +#[test] +fn rust_f32_max_rejects_non_f32_operands_and_results() { + for profile in profiles() { + for (operands, results) in [ + (["f64", "f32"], ["f32"]), + (["f32", "f64"], ["f32"]), + (["i32", "f32"], ["f32"]), + (["f32", "i32"], ["f32"]), + (["f32", "f32"], ["f64"]), + (["f32", "f32"], ["i32"]), + ] { + let result = translate_call(&*profile, CALLEE_MAXNUM_NSZ_F32, &operands, &results); + assert!( + result.is_err(), + "{} accepted {operands:?} -> {results:?}: {result:?}", + profile.name() + ); + } + } +} + +/// Ensure an unrelated scalar function is not mistaken for the Rust max intrinsic. +#[test] +fn ordinary_f32_function_calls_remain_calls() { + for profile in profiles() { + let text = translate_call(&*profile, "user_max", &["f32", "f32"], &["f32"]).unwrap(); + assert!( + text.contains("\"func.call\"(%v0, %v1) <{callee = @user_max}> : (f32, f32) -> (f32)"), + "{text}" + ); + assert!(!text.contains("\"nvvm.fmax\""), "{text}"); + } +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sm120_primitives.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sm120_primitives.rs new file mode 100644 index 0000000000..22d9073898 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sm120_primitives.rs @@ -0,0 +1,659 @@ +//! SM120 matrix instructions and TMA mappings with selector and identity rejection cases. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, CutlassPrimitivesMlir22, MlirConsumerProfile}; +use dialect_mir::{ + ops::{MirConstantOp, MirFuncOp, MirReturnOp}, + types::MirPtrType, +}; +use dialect_nvvm::ops::*; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{IntegerAttr, StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::TypeHandle, + utils::apint::APInt, + value::Value, +}; + +/// MIR/NVVM function with shared pointers and register carriers for SM120 primitive export tests. +struct Fixture { + /// Source IR context owning the fixture's registered types and operations. + ctx: Context, + /// Source module containing the operations under test. + module: ModuleOp, + /// Function entry block where the fixture appends source operations. + block: Ptr, + /// Source function argument types used to build well-typed primitive operations. + types: Vec, +} +impl Fixture { + /// Build the reusable SM120 source context, module, and typed argument block. + fn new() -> Self { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + let i8 = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let i32 = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let i64 = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let f32 = FP32Type::get(&ctx).into(); + let types = vec![ + MirPtrType::get(&mut ctx, i32, true, 0).into(), + MirPtrType::get(&mut ctx, i8, true, 3).into(), + MirPtrType::get(&mut ctx, i64, true, 3).into(), + MirPtrType::get(&mut ctx, i8, false, 0).into(), + i32, + i32, + f32, + f32, + f32, + f32, + i32, + i32, + i32, + i32, + i32, + i32, + i32, + i32, + ]; + let module = ModuleOp::new(&mut ctx, "sm120".try_into().unwrap()); + let ft = FunctionType::get(&ctx, types.clone(), vec![]); + let raw = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, raw, TypeAttr::new(ft.into())) + .set_symbol_name(&mut ctx, "sm120_primitives".try_into().unwrap()); + raw.deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".into()), + ); + module.append_operation(&mut ctx, raw, 0); + let block = BasicBlock::new(&mut ctx, None, types.clone()); + block.insert_at_back(raw.deref(&ctx).get_region(0), &ctx); + Self { + ctx, + module, + block, + types, + } + } + /// Return the indexed source function argument for constructing a test operation. + fn arg(&self, i: usize) -> Value { + self.block.deref(&self.ctx).get_argument(i) + } + /// Append a typed source operation to the fixture block and return its operation handle. + fn append( + &mut self, + marker: &str, + operands: Vec, + types: Vec, + ) -> Ptr { + let op = Operation::new( + &mut self.ctx, + O::get_concrete_op_info(), + types, + operands, + vec![], + 0, + ); + if !marker.is_empty() { + op.deref_mut(&self.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + } + op.insert_at_back(self.block, &self.ctx); + op + } + /// Append an unsigned constant with a chosen bit width for intrinsic selectors. + fn constant(&mut self, width: u32, value: u64) -> Value { + let ty = IntegerType::get(&self.ctx, width, Signedness::Unsigned); + let op = self.append::("", vec![], vec![ty.into()]); + MirConstantOp::new(op).set_attr_value( + &self.ctx, + IntegerAttr::new( + ty, + APInt::from_u64(value, (width as usize).try_into().unwrap()), + ), + ); + op.deref(&self.ctx).get_result(0) + } + /// Terminate the source function after the test's primitive operations. + fn finish(&mut self) { + self.append::("", vec![], vec![]); + } + /// Export the fixture with the SM120 primitives profile and render its target MLIR. + fn translate(&self) -> Result { + CutlassPrimitivesMlir22::new("sm_120a") + .unwrap() + .translate_module(&self.ctx, &self.module) + .and_then(|m| render_module(&m).map_err(Into::into)) + .map_err(|e| e.to_string()) + } + /// Append the canonical reviewed shared matrix-load operation for mutation tests. + fn load(&mut self) -> Ptr { + let address = self.arg(0); + let load = LdmatrixOp::build( + &mut self.ctx, + address, + LdmatrixShapeAttr::M8n8, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B16, + LdmatrixStateSpaceAttr::Shared, + ); + load.deref_mut(&self.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0013".into()), + ); + load.insert_at_back(self.block, &self.ctx); + load + } + + /// Append the canonical reviewed MXF4 register-MMA operation for mutation tests. + fn mma(&mut self) -> Ptr { + let mut operands = (6..16).map(|i| self.arg(i)).collect::>(); + // Distinct definitions prove each selector and scale retains its slot. + operands.push(self.arg(16)); + operands.push(self.constant(16, 0)); + operands.push(self.constant(16, 0)); + operands.push(self.arg(17)); + operands.push(self.constant(16, 0)); + operands.push(self.constant(16, 0)); + let op = self.append::("v1:i1039", operands, vec![self.types[6]; 4]); + let mma = RegisterMmaOp::new(op); + mma.set_attr_nvvm_register_mma_shape(&self.ctx, RegisterMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_register_mma_operation(&self.ctx, RegisterMmaOperationAttr::Multiply); + mma.set_attr_nvvm_register_mma_kind(&self.ctx, RegisterMmaKindAttr::Mxf4); + mma.set_attr_nvvm_register_mma_accumulator(&self.ctx, RegisterMmaAccumulatorAttr::F32); + mma.set_attr_nvvm_register_mma_a_element(&self.ctx, RegisterMmaElementAttr::E2m1); + mma.set_attr_nvvm_register_mma_b_element(&self.ctx, RegisterMmaElementAttr::E2m1); + mma.set_attr_nvvm_register_mma_a_layout(&self.ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(&self.ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(&self.ctx, RegisterMmaOverflowAttr::NotApplicable); + op + } +} + +/// Check MXF4 operands enter NVVM in A/B/C/scale order and every accumulator result is recovered. +#[test] +fn sm120_mxf4_preserves_a_b_c_scale_order_and_extracts_all_results() { + let mut f = Fixture::new(); + let mma = f.mma(); + let results = mma.deref(&f.ctx).results().collect::>(); + let top = f.append::("v1:i0081", results[..2].to_vec(), vec![f.types[4]]); + let bottom = f.append::("v1:i0081", results[2..].to_vec(), vec![f.types[4]]); + let top = top.deref(&f.ctx).get_result(0); + let bottom = bottom.deref(&f.ctx).get_result(0); + f.append::("v1:i0301", vec![f.arg(1), top, bottom], vec![]); + f.finish(); + let text = f.translate().unwrap(); + assert!(text.contains("\"nvvm.mma.block_scale\"(%v10, %v11, %v12, %v13, %v14, %v15, %v6, %v7, %v8, %v9, %v16, %v18, %v19, %v17, %v20, %v21)"), "{text}"); + for expected in [ + "#nvvm.shape", + "#nvvm.scale_vec_size", + "#nvvm.block_scale_kind", + "#nvvm.block_scale_format", + "operandSegmentSizes = array", + ] { + assert!(text.contains(expected), "{text}"); + } + assert_eq!(text.matches("#nvvm.mma_type").count(), 2); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), 4); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); +} + +/// Reject unsupported MXF4 selectors and catalog identities before emission. +#[test] +fn sm120_mxf4_rejects_other_variants_and_forged_identity() { + let mut f = Fixture::new(); + let op = f.mma(); + f.finish(); + let mma = RegisterMmaOp::new(op); + mma.set_attr_nvvm_register_mma_shape(&f.ctx, RegisterMmaShapeAttr::M16n8k32); + mma.set_attr_nvvm_register_mma_kind(&f.ctx, RegisterMmaKindAttr::Mxf8f6f4); + let error = f.translate().unwrap_err(); + assert!(error.contains("supports m16n8k64 MXF4"), "{error}"); + mma.set_attr_nvvm_register_mma_shape(&f.ctx, RegisterMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_register_mma_kind(&f.ctx, RegisterMmaKindAttr::Mxf4); + op.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0858".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i1039") + ); +} + +/// Reject the SM120-only MXF4 instruction when targeting SM100. +#[test] +fn sm100_rejects_sm120_mxf4() { + let mut f = Fixture::new(); + f.mma(); + f.finish(); + let error = CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(&f.ctx, &f.module) + .unwrap_err() + .to_string(); + assert!( + error.contains("MXF4 register MMA requires sm_120a"), + "{error}" + ); + let error = CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&f.ctx, &f.module) + .unwrap_err() + .to_string(); + assert!( + error.contains("MXF4 register MMA requires sm_120a"), + "{error}" + ); +} + +/// Check packed FP16 operand lanes and scalar FP32 accumulator positions survive export. +#[test] +fn fp16_mma_preserves_packed_a_b_and_f32_accumulator_order() { + let mut f = Fixture::new(); + let operands = (6..16).map(|i| f.arg(i)).collect(); + let op = f.append::("v1:i0106", operands, vec![f.types[6]; 4]); + let mma = RegisterMmaOp::new(op); + mma.set_attr_nvvm_register_mma_shape(&f.ctx, RegisterMmaShapeAttr::M16n8k16); + mma.set_attr_nvvm_register_mma_operation(&f.ctx, RegisterMmaOperationAttr::Multiply); + mma.set_attr_nvvm_register_mma_kind(&f.ctx, RegisterMmaKindAttr::Standard); + mma.set_attr_nvvm_register_mma_accumulator(&f.ctx, RegisterMmaAccumulatorAttr::F32); + mma.set_attr_nvvm_register_mma_a_element(&f.ctx, RegisterMmaElementAttr::F16); + mma.set_attr_nvvm_register_mma_b_element(&f.ctx, RegisterMmaElementAttr::F16); + mma.set_attr_nvvm_register_mma_a_layout(&f.ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(&f.ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(&f.ctx, RegisterMmaOverflowAttr::NotApplicable); + f.finish(); + let text = f.translate().unwrap(); + for register in 10..16 { + assert!( + text.contains(&format!( + "\"llvm.bitcast\"(%v{register}) : (i32) -> (vector<2xf16>)" + )), + "{text}" + ); + } + let issue = text + .lines() + .find(|line| line.contains("\"nvvm.mma.sync\"")) + .unwrap(); + let packed_values: Vec<_> = (10..16) + .map(|register| { + text.lines() + .find(|line| line.contains(&format!("\"llvm.bitcast\"(%v{register})"))) + .unwrap() + .trim() + .split(" = ") + .next() + .unwrap() + }) + .collect(); + assert!( + issue.contains(&format!( + "\"nvvm.mma.sync\"({}, %v6, %v7, %v8, %v9)", + packed_values.join(", ") + )), + "{issue}" + ); + for arch in ["sm_100a", "sm_120a"] { + let text = CutlassFullCuteMlir22::new(arch) + .unwrap() + .translate_module(&f.ctx, &f.module) + .and_then(|m| render_module(&m).map_err(Into::into)) + .unwrap(); + assert!(text.contains("\"nvvm.mma.sync\""), "{arch}: {text}"); + } + assert!( + issue.contains("operandSegmentSizes = array"), + "{issue}" + ); + assert!( + issue.contains("#nvvm.shape"), + "{issue}" + ); + assert!(issue.contains("layoutA = #nvvm.mma_layout"), "{issue}"); + assert!(issue.contains("layoutB = #nvvm.mma_layout"), "{issue}"); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), 4); + + op.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i1039".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i0106") + ); +} + +/// Build the BF16 instruction with independent dtype and identity selectors. +fn bf16_mma_fixture() -> (Fixture, RegisterMmaOp) { + let mut f = Fixture::new(); + let operands = (6..16).map(|i| f.arg(i)).collect(); + let op = f.append::("v1:i0105", operands, vec![f.types[6]; 4]); + let mma = RegisterMmaOp::new(op); + mma.set_attr_nvvm_register_mma_shape(&f.ctx, RegisterMmaShapeAttr::M16n8k16); + mma.set_attr_nvvm_register_mma_operation(&f.ctx, RegisterMmaOperationAttr::Multiply); + mma.set_attr_nvvm_register_mma_kind(&f.ctx, RegisterMmaKindAttr::Standard); + mma.set_attr_nvvm_register_mma_accumulator(&f.ctx, RegisterMmaAccumulatorAttr::F32); + mma.set_attr_nvvm_register_mma_a_element(&f.ctx, RegisterMmaElementAttr::Bf16); + mma.set_attr_nvvm_register_mma_b_element(&f.ctx, RegisterMmaElementAttr::Bf16); + mma.set_attr_nvvm_register_mma_a_layout(&f.ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(&f.ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(&f.ctx, RegisterMmaOverflowAttr::NotApplicable); + f.finish(); + (f, mma) +} + +/// NVVM consumes BF16 register pairs as i32, with the original low/high bits unchanged. +#[test] +fn bf16_mma_preserves_integer_packing_and_fp32_accumulator_order() { + let (f, _) = bf16_mma_fixture(); + let text = f.translate().unwrap(); + let issue = text + .lines() + .find(|line| line.contains("\"nvvm.mma.sync\"")) + .unwrap(); + assert!( + issue.contains("\"nvvm.mma.sync\"(%v10, %v11, %v12, %v13, %v14, %v15, %v6, %v7, %v8, %v9)"), + "{issue}" + ); + assert!( + issue.contains("(i32, i32, i32, i32, i32, i32, f32, f32, f32, f32)"), + "{issue}" + ); + assert!( + issue.contains("!llvm.struct<(f32, f32, f32, f32)>"), + "{issue}" + ); + assert!( + issue.contains("operandSegmentSizes = array"), + "{issue}" + ); + assert!( + issue.contains("#nvvm.shape"), + "{issue}" + ); + assert!(issue.contains("layoutA = #nvvm.mma_layout"), "{issue}"); + assert!(issue.contains("layoutB = #nvvm.mma_layout"), "{issue}"); + assert!( + issue.contains("multiplicandAPtxType = #nvvm.mma_type"), + "{issue}" + ); + assert!( + issue.contains("multiplicandBPtxType = #nvvm.mma_type"), + "{issue}" + ); + assert!(!text.contains("\"llvm.bitcast\""), "{text}"); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), 4); + for arch in ["sm_100a", "sm_120a"] { + let text = CutlassFullCuteMlir22::new(arch) + .unwrap() + .translate_module(&f.ctx, &f.module) + .and_then(|m| render_module(&m).map_err(Into::into)) + .unwrap(); + assert!( + text.contains("multiplicandAPtxType = #nvvm.mma_type"), + "{arch}: {text}" + ); + } +} + +/// A BF16 shape cannot impersonate the FP16 intrinsic with an otherwise identical register ABI. +#[test] +fn bf16_mma_rejects_fp16_intrinsic_identity() { + let (f, mma) = bf16_mma_fixture(); + mma.get_operation().deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0106".into()), + ); + let error = f.translate().unwrap_err(); + assert!( + error.contains("expected intrinsic marker v1:i0105"), + "{error}" + ); +} + +/// Same-size FP16 and BF16 registers do not make mixed-dtype instructions valid. +#[test] +fn bf16_mma_rejects_mixed_operand_dtypes() { + for change_a in [false, true] { + let (f, mma) = bf16_mma_fixture(); + if change_a { + mma.set_attr_nvvm_register_mma_a_element(&f.ctx, RegisterMmaElementAttr::F16); + } else { + mma.set_attr_nvvm_register_mma_b_element(&f.ctx, RegisterMmaElementAttr::F16); + } + assert!(f.translate().is_err()); + } +} + +/// Check each reviewed matrix-load form retains its transpose flag and exact register count. +#[test] +fn matrix_loads_preserve_transposition_and_register_count() { + for (multiplicity, layout, marker, spelling, count) in [ + ( + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Normal, + "v1:i0027", + "row", + 2, + ), + ( + LdmatrixMultiplicityAttr::X2, + LdmatrixLayoutAttr::Transposed, + "v1:i0028", + "col", + 2, + ), + ( + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Transposed, + "v1:i0029", + "col", + 4, + ), + ] { + let mut f = Fixture::new(); + let address = f.arg(0); + let load = LdmatrixOp::build( + &mut f.ctx, + address, + LdmatrixShapeAttr::M8n8, + multiplicity, + layout, + LdmatrixElementAttr::B16, + LdmatrixStateSpaceAttr::Shared, + ); + load.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + load.insert_at_back(f.block, &f.ctx); + f.finish(); + let text = f.translate().unwrap(); + assert!(text.contains(&format!("num = {count} : i32")), "{text}"); + assert!( + text.contains(&format!("#nvvm.mma_layout<{spelling}>")), + "{text}" + ); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), count); + load.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0013".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains(&format!("expected intrinsic marker {marker}")) + ); + } +} + +/// Check a load/FP16-pack/store chain preserves scalar register and packed-lane ordering. +#[test] +fn sm120_load_pack_and_store_preserve_register_and_lane_order() { + let mut f = Fixture::new(); + let load = f.load(); + let pack = f.append::("v1:i0081", vec![f.arg(6), f.arg(7)], vec![f.types[4]]); + let registers = load.deref(&f.ctx).results().collect::>(); + let packed = pack.deref(&f.ctx).get_result(0); + f.append::("v1:i0301", vec![f.arg(1), packed, registers[3]], vec![]); + f.finish(); + let text = f.translate().unwrap(); + assert!(text.contains("\"nvvm.ldmatrix\""), "{text}"); + assert!(text.contains("num = 4 : i32"), "{text}"); + assert_eq!(text.matches("\"llvm.extractvalue\"").count(), 4); + assert!(text.contains("\"llvm.call_intrinsic\"(%v7, %v6)"), "{text}"); + assert!(text.contains("llvm.nvvm.ff2f16x2.rn"), "{text}"); + assert!(text.contains("\"nvvm.stmatrix\""), "{text}"); + assert!(text.contains("#nvvm.mma_layout"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker")); +} + +/// Reject matrix-load shapes, packing choices, or identities outside the reviewed subset. +#[test] +fn sm120_rejects_other_matrix_load_selectors_and_forged_identity() { + let mut f = Fixture::new(); + // An integer shared address is valid for both tested matrix element formats. + let address = f.arg(4); + let load = LdmatrixOp::build( + &mut f.ctx, + address, + LdmatrixShapeAttr::M8n8, + LdmatrixMultiplicityAttr::X4, + LdmatrixLayoutAttr::Normal, + LdmatrixElementAttr::B16, + LdmatrixStateSpaceAttr::Shared, + ); + load.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0013".into()), + ); + load.insert_at_back(f.block, &f.ctx); + f.finish(); + LdmatrixOp::new(load).set_attr_nvvm_ldmatrix_shape(&f.ctx, LdmatrixShapeAttr::M8n16); + LdmatrixOp::new(load) + .set_attr_nvvm_ldmatrix_element(&f.ctx, LdmatrixElementAttr::B8x16B4x16P64); + assert!(f.translate().unwrap_err().contains("m8n8.x2/x4.b16")); + LdmatrixOp::new(load).set_attr_nvvm_ldmatrix_shape(&f.ctx, LdmatrixShapeAttr::M8n8); + LdmatrixOp::new(load).set_attr_nvvm_ldmatrix_element(&f.ctx, LdmatrixElementAttr::B16); + load.deref_mut(&f.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new("v1:i0001".into()), + ); + assert!( + f.translate() + .unwrap_err() + .contains("expected intrinsic marker v1:i0013") + ); +} + +/// Check TMA uses the local shared completion barrier without introducing multicast. +#[test] +fn sm120_tma_uses_local_shared_barrier_and_no_multicast() { + let mut f = Fixture::new(); + let mask = f.constant(16, 0); + let hint = f.constant(64, 0); + f.append::( + "v1:i0329", + vec![f.arg(1), f.arg(2), f.arg(3), f.arg(4), f.arg(5), mask, hint], + vec![], + ); + f.finish(); + let text = f.translate().unwrap(); + assert!( + text.contains("\"nvvm.cp.async.bulk.tensor.shared.cluster.global\""), + "{text}" + ); + assert!( + text.contains("operandSegmentSizes = array"), + "{text}" + ); + assert!(text.contains("useIntrinsic = false"), "{text}"); + assert!(!text.contains("cta_2")); +} + +/// Reject unsupported TMA selectors instead of silently dropping their semantics. +#[test] +fn sm120_tma_rejects_nonzero_multicast_mask_or_cache_hint() { + for (mask, hint) in [(1, 0), (0, 1)] { + let mut f = Fixture::new(); + let mask = f.constant(16, mask); + let hint = f.constant(64, hint); + f.append::( + "v1:i0329", + vec![f.arg(1), f.arg(2), f.arg(3), f.arg(4), f.arg(5), mask, hint], + vec![], + ); + f.finish(); + assert!( + f.translate() + .unwrap_err() + .contains("zero multicast mask and cache hint") + ); + } +} + +/// Reject SM100-only tensor-memory instructions during source architecture validation. +#[test] +fn sm120_rejects_tcgen05_before_translation() { + let mut f = Fixture::new(); + f.append::("v1:i0357", vec![], vec![]); + f.finish(); + let error = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(&f.ctx, &f.module) + .unwrap_err() + .to_string(); + assert!( + error.contains("TCGen05 is unsupported on sm_120a"), + "{error}" + ); + assert!( + f.translate() + .unwrap_err() + .contains("TCGen05 is unsupported on sm_120a") + ); +} + +/// Reject a raw integer barrier carrier where the reviewed TMA contract requires a pointer. +#[test] +fn sm120_tma_requires_pointer_barrier_carrier() { + let mut f = Fixture::new(); + let mask = f.constant(16, 0); + let hint = f.constant(64, 0); + f.append::( + "v1:i0329", + vec![f.arg(1), f.arg(4), f.arg(3), f.arg(4), f.arg(5), mask, hint], + vec![], + ); + f.finish(); + assert!( + f.translate() + .unwrap_err() + .contains("TMA requires pointer carriers") + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs new file mode 100644 index 0000000000..9c2fbcf4fb --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/sync_restrict.rs @@ -0,0 +1,195 @@ +//! Restricted proxy-fence direction, ordering, identity, and optional native-parser checks. + +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, CutlassPrimitivesMlir22, MlirConsumerProfile}; +use dialect_mir::ops::{MirFuncOp, MirReturnOp}; +use dialect_nvvm::ops::{ + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceProxyAsyncGenericReleaseSharedCtaClusterOp, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::{StringAttr, TypeAttr}, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::FunctionType, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, +}; + +/// Source module containing release/acquire proxy fences whose contracts can be mutated independently. +struct Fixture { + /// Source IR context owning the fixture's registered types and operations. + ctx: Context, + /// Source module containing the operations under test. + module: ModuleOp, + /// Source release-fence operation used for identity and attribute rejection cases. + release: Ptr, + /// Source acquire-fence operation used for identity and attribute rejection cases. + acquire: Ptr, +} + +impl Fixture { + /// Build canonical restricted release/acquire fences with their native catalog identities. + fn new() -> Self { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_nvvm::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new(&mut ctx, "sync_restrict".try_into().unwrap()); + let ty = FunctionType::get(&ctx, vec![], vec![]); + let raw = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, raw, TypeAttr::new(ty.into())) + .set_symbol_name(&mut ctx, "proxy_fences".try_into().unwrap()); + module.append_operation(&mut ctx, raw, 0); + let block = BasicBlock::new(&mut ctx, None, vec![]); + block.insert_at_back(raw.deref(&ctx).get_region(0), &ctx); + let release = FenceProxyAsyncGenericReleaseSharedCtaClusterOp::build(&mut ctx); + let acquire = FenceProxyAsyncGenericAcquireSharedClusterClusterOp::build(&mut ctx); + for (op, marker) in [(release, "v1:i0314"), (acquire, "v1:i0315")] { + op.deref_mut(&ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(marker.into()), + ); + op.insert_at_back(block, &ctx); + } + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(block, &ctx); + Self { + ctx, + module, + release, + acquire, + } + } + + /// Export the fixture through the selected full-CuTe or primitives profile. + fn translate(&self, primitives: bool) -> Result { + let module = if primitives { + CutlassPrimitivesMlir22::new("sm_100a") + .unwrap() + .translate_module(&self.ctx, &self.module) + } else { + CutlassFullCuteMlir22::new("sm_100a") + .unwrap() + .translate_module(&self.ctx, &self.module) + }; + module + .and_then(|m| render_module(&m).map_err(Into::into)) + .map_err(|e| e.to_string()) + } +} + +/// Check both profiles keep acquire/release order and the corresponding shared-memory restriction. +#[test] +fn selected_cutlass_profiles_preserve_proxy_direction_and_each_memory_order() { + let fixture = Fixture::new(); + for primitives in [false, true] { + let text = fixture.translate(primitives).unwrap(); + let fences = text + .lines() + .filter(|line| line.contains("\"nvvm.fence.proxy.sync_restrict\"")) + .collect::>(); + assert_eq!(fences.len(), 2, "{text}"); + for (line, order) in fences.iter().zip(["release", "acquire"]) { + assert!( + line.contains(&format!("order = #nvvm.mem_order<{order}>")), + "{text}" + ); + assert!( + line.contains("fromProxy = #nvvm.proxy_kind"), + "{text}" + ); + assert!(line.contains("toProxy = #nvvm.proxy_kind"), "{text}"); + assert!(line.contains(": () -> ()"), "{text}"); + } + assert!(!text.contains("inline_asm"), "{text}"); + assert!(!text.contains("cuda_oxide_intrinsic_marker"), "{text}"); + } +} + +/// Reject fences carrying the other direction's catalog identity. +#[test] +fn proxy_fences_reject_swapped_intrinsic_markers() { + for acquire in [false, true] { + let fixture = Fixture::new(); + let (op, wrong, expected) = if acquire { + (fixture.acquire, "v1:i0314", "v1:i0315") + } else { + (fixture.release, "v1:i0315", "v1:i0314") + }; + op.deref_mut(&fixture.ctx).attributes.set( + "cuda_oxide_intrinsic_marker".try_into().unwrap(), + StringAttr::new(wrong.into()), + ); + let error = fixture.translate(false).unwrap_err(); + assert!( + error.contains(&format!("expected intrinsic marker {expected}")), + "{error}" + ); + } +} + +/// Reject source attributes the reviewed fence recipe has not explicitly consumed. +#[test] +fn proxy_fences_reject_unexpected_semantic_attributes() { + let fixture = Fixture::new(); + fixture.release.deref_mut(&fixture.ctx).attributes.set( + "order".try_into().unwrap(), + StringAttr::new("acquire".into()), + ); + let error = fixture.translate(false).unwrap_err(); + assert!(error.contains("residual attributes"), "{error}"); +} + +/// When the pinned parser is available, verify it accepts both exported fence forms. +#[test] +#[ignore = "requires CUDA_OXIDE_CUTLASS_PYTHON pointing to the selected CUTLASS Python environment"] +fn selected_cutlass_native_parser_verifies_both_proxy_fences() { + use std::{ + io::Write, + process::{Command, Stdio}, + }; + let python = std::env::var_os("CUDA_OXIDE_CUTLASS_PYTHON") + .expect("set CUDA_OXIDE_CUTLASS_PYTHON to the selected CUTLASS Python executable"); + let text = Fixture::new().translate(false).unwrap(); + let mut child = Command::new(python) + .args(["-c", "import sys\nfrom cutlass._mlir import ir\nfrom cutlass._mlir.dialects import nvvm\nwith ir.Context():\n module = ir.Module.parse(sys.stdin.read())\n assert module.operation.verify()\n"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + child + .stdin + .take() + .unwrap() + .write_all(text.as_bytes()) + .unwrap(); + let output = child.wait_with_output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/transparent_tensor.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/transparent_tensor.rs new file mode 100644 index 0000000000..3e6fa52677 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/transparent_tensor.rs @@ -0,0 +1,157 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Transparent numeric wrappers retain Rust identity across CuTe vector copies. + +use cuda_oxide_mlir_export::render_module; +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use dialect_cute::{ + attributes::CuteTensorAccessAttr, + tensor_ops::{ + CuteTensorLoadIntoOp, CuteTensorMakeOp, CuteTensorSliceOp, CuteTensorStoreElementAbsOp, + CuteTensorStoreFromOp, CuteTensorZippedDivideOp, + }, +}; +use dialect_mir::{ + ops::{MirFuncOp, MirReturnOp}, + types::{MirPtrType, MirStructType, StructAbiKind, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + identifier::Identifier, + op::Op, + operation::Operation, + r#type::TypeHandle, +}; + +/// Append one detached operation to the test function. +fn append(block: Ptr, operation: O, ctx: &Context) -> Ptr { + let operation = operation.get_operation(); + operation.insert_at_back(block, ctx); + operation +} + +/// Copies use the scalar representation; scalar tails retain the aggregate's exact stride. +#[test] +fn transparent_u16_tensor_uses_bit_preserving_vector_copies_and_tails() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let module = ModuleOp::new( + &mut ctx, + Identifier::try_from("transparent_tensor").unwrap(), + ); + let u16_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let wrapper: TypeHandle = MirStructType::get_with_full_layout_and_abi( + &mut ctx, + "NumericBits".into(), + vec!["0".into()], + vec![u16_ty], + vec![0], + vec![0], + 2, + 2, + StructAbiKind::TransparentScalar, + ) + .into(); + let read_ptr: TypeHandle = + MirPtrType::get(&mut ctx, wrapper, false, address_space::GENERIC).into(); + let write_ptr: TypeHandle = + MirPtrType::get(&mut ctx, wrapper, true, address_space::GENERIC).into(); + let carrier_ptr: TypeHandle = + MirPtrType::get(&mut ctx, wrapper, true, address_space::LOCAL).into(); + let arguments = vec![read_ptr, write_ptr, carrier_ptr, u64_ty, u64_ty, wrapper]; + let function_type = FunctionType::get(&ctx, arguments.clone(), vec![]); + let function_op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, function_op, TypeAttr::new(function_type.into())); + function.set_symbol_name(&mut ctx, Identifier::try_from("copy_numeric_bits").unwrap()); + module.append_operation(&mut ctx, function_op, 0); + let region = function_op.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, arguments); + entry.insert_at_back(region, &ctx); + let args = entry.deref(&ctx).arguments().collect::>(); + + for (pointer, access) in [ + (args[0], CuteTensorAccessAttr::ReadOnly), + (args[1], CuteTensorAccessAttr::ReadWrite), + ] { + let make = append( + entry, + CuteTensorMakeOp::new(&mut ctx, pointer, args[3], wrapper, wrapper, access, 2), + &ctx, + ); + let view = make.deref(&ctx).get_result(0); + let zipped = append( + entry, + CuteTensorZippedDivideOp::new(&mut ctx, view, 8), + &ctx, + ); + let divided = zipped.deref(&ctx).get_result(0); + let slice = append( + entry, + CuteTensorSliceOp::new(&mut ctx, divided, args[4]), + &ctx, + ); + let tile = slice.deref(&ctx).get_result(0); + if access == CuteTensorAccessAttr::ReadOnly { + append( + entry, + CuteTensorLoadIntoOp::new(&mut ctx, tile, args[2], 16), + &ctx, + ); + } else { + append( + entry, + CuteTensorStoreFromOp::new(&mut ctx, args[2], tile, 16), + &ctx, + ); + append( + entry, + CuteTensorStoreElementAbsOp::new(&mut ctx, tile, args[4], args[5]), + &ctx, + ); + } + } + let ret = Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + ret.insert_at_back(entry, &ctx); + + let profile = CutlassFullCuteMlir22::new("sm_120a").unwrap(); + let target = profile.translate_module(&ctx, &module).unwrap(); + let text = render_module(&target).unwrap(); + assert!(text.contains("!cute.memref"), + "{text}" + ); + assert_eq!(text.matches("\"cute.copy\"").count(), 2, "{text}"); + assert!(text.contains("elem_type = !llvm.struct<(i16)>"), "{text}"); + assert!(text.contains("\"llvm.store\""), "{text}"); + assert!( + !text.contains("\"arith.extf\"") && !text.contains("\"arith.truncf\""), + "{text}" + ); +} diff --git a/cuda-oxide/crates/cuda-oxide-mlir-export/tests/value_reference.rs b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/value_reference.rs new file mode 100644 index 0000000000..545c864345 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/tests/value_reference.rs @@ -0,0 +1,407 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! References to SSA values need fresh storage at their original control-flow position. + +// These fixtures intentionally forge pointer contracts to test verifier rejection. +#![allow(clippy::disallowed_methods)] + +use cuda_oxide_mlir_export::{CutlassFullCuteMlir22, MlirConsumerProfile}; +use cuda_oxide_mlir_export::{MlirAttribute, MlirOperation, MlirType}; +use dialect_mir::{ + attributes::MirPointerKindAuthorityAttr, + ops::{MirConstructStructOp, MirFuncOp, MirGotoOp, MirRefOp, MirReturnOp}, + types::{MirArrayType, MirPointerKind, MirPtrType, MirStructType, MirTupleType}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::TypeHandle, + value::Value, +}; + +/// Build the same full-profile context used by CuTe device compilation. +fn context() -> Context { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + ctx +} + +/// Append one source operation with no successors or nested regions. +fn append( + ctx: &mut Context, + block: Ptr, + results: Vec, + operands: Vec, +) -> Ptr { + let operation = Operation::new(ctx, O::get_concrete_op_info(), results, operands, vec![], 0); + operation.insert_at_back(block, ctx); + operation +} + +/// Put values in a second block so accidental allocation hoisting is observable. +fn function( + ctx: &mut Context, + args: Vec, +) -> (ModuleOp, Ptr, Ptr) { + let module = ModuleOp::new(ctx, "value_reference".try_into().unwrap()); + let signature = FunctionType::get(ctx, args.clone(), vec![]); + let function = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(ctx, function, TypeAttr::new(signature.into())) + .set_symbol_name(ctx, "test".try_into().unwrap()); + module.append_operation(ctx, function, 0); + let region = function.deref(ctx).get_region(0); + let entry = BasicBlock::new(ctx, None, args); + let body = BasicBlock::new(ctx, None, vec![]); + entry.insert_at_back(region, ctx); + body.insert_at_back(region, ctx); + let branch = Operation::new( + ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![body], + 0, + ); + branch.insert_at_back(entry, ctx); + (module, entry, body) +} + +/// Materialize a reference with the explicitly supplied Rust boundary evidence. +fn reference( + ctx: &mut Context, + block: Ptr, + value: Value, + pointer: TypeHandle, + mutable: bool, + authority: MirPointerKindAuthorityAttr, +) { + let reference = append::(ctx, block, vec![pointer], vec![value]); + let reference = MirRefOp::new(reference); + reference.set_mutable(ctx, mutable); + reference.set_pointer_kind_authority(ctx, authority); +} + +/// Return the translated test function, retaining individual operations for assertions. +fn translate(ctx: &Context, module: &ModuleOp) -> Result { + fn find(function: &MlirOperation) -> Option<&MlirOperation> { + if function.name == "func.func" { + return Some(function); + } + function + .regions + .iter() + .flat_map(|region| ®ion.blocks) + .flat_map(|block| &block.operations) + .find_map(find) + } + let translated = CutlassFullCuteMlir22::new("sm_120a") + .unwrap() + .translate_module(ctx, module) + .map_err(|error| error.to_string())?; + Ok(find(&translated.root) + .expect("translated test function") + .clone()) +} + +/// FnMut conversion closures have no captures, but still require a mutable reference. +#[test] +fn zero_sized_closure_reference_materializes_in_its_original_block() { + let mut ctx = context(); + let closure = MirStructType::get_with_full_layout( + &mut ctx, + "ConversionClosure".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 1, + ) + .into(); + let pointer = + MirPtrType::get_with_kind(&mut ctx, closure, true, 0, MirPointerKind::UniqueRef).into(); + let (module, _, body) = function(&mut ctx, vec![]); + let value = append::(&mut ctx, body, vec![closure], vec![]) + .deref(&ctx) + .get_result(0); + reference( + &mut ctx, + body, + value, + pointer, + true, + MirPointerKindAuthorityAttr::Reborrow, + ); + append::(&mut ctx, body, vec![], vec![]); + let function = translate(&ctx, &module).unwrap(); + let blocks = &function.regions[0].blocks; + assert!( + blocks[0] + .operations + .iter() + .all(|operation| operation.name != "llvm.alloca") + ); + let operations = &blocks[1].operations; + let alloca = operations + .iter() + .find(|operation| operation.name == "llvm.alloca") + .unwrap(); + let store = operations + .iter() + .find(|operation| operation.name == "llvm.store") + .unwrap(); + assert_eq!(store.operands[1].id, alloca.results[0].id); + assert_eq!( + alloca.properties["alignment"], + MlirAttribute::Integer { + value: 1, + ty: MlirType::Integer(64) + } + ); +} + +/// Distinct references must retain distinct storage, even when borrowing the same SSA value. +#[test] +fn scalar_references_initialize_distinct_slots_with_the_operand() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let pointer = + MirPtrType::get_with_kind(&mut ctx, scalar, false, 0, MirPointerKind::SharedRef).into(); + let (module, entry, body) = function(&mut ctx, vec![scalar]); + let value = entry.deref(&ctx).get_argument(0); + for _ in 0..2 { + reference( + &mut ctx, + body, + value, + pointer, + false, + MirPointerKindAuthorityAttr::Reborrow, + ); + } + append::(&mut ctx, body, vec![], vec![]); + let function = translate(&ctx, &module).unwrap(); + let blocks = &function.regions[0].blocks; + let operations = &blocks[1].operations; + let allocas: Vec<_> = operations + .iter() + .filter(|operation| operation.name == "llvm.alloca") + .collect(); + let stores: Vec<_> = operations + .iter() + .filter(|operation| operation.name == "llvm.store") + .collect(); + assert_eq!(allocas.len(), 2); + assert_eq!(stores.len(), 2); + assert_ne!(allocas[0].results[0].id, allocas[1].results[0].id); + for (alloca, store) in allocas.iter().zip(stores) { + assert_eq!(store.operands[0].id, blocks[0].arguments[0].id); + assert_eq!(store.operands[1].id, alloca.results[0].id); + for operation in [*alloca, store] { + assert_eq!( + operation.properties["alignment"], + MlirAttribute::Integer { + value: 4, + ty: MlirType::Integer(64) + } + ); + } + } +} + +/// Explicit Rust over-alignment belongs on the allocation and initialization store. +#[test] +fn value_reference_preserves_explicit_abi_alignment() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let tuple = + MirTupleType::get_with_layout(&mut ctx, vec![scalar], vec![0], vec![0], 32, 32).into(); + let pointer = + MirPtrType::get_with_kind(&mut ctx, tuple, false, 0, MirPointerKind::SharedRef).into(); + let (module, entry, body) = function(&mut ctx, vec![tuple]); + let value = entry.deref(&ctx).get_argument(0); + reference( + &mut ctx, + body, + value, + pointer, + false, + MirPointerKindAuthorityAttr::Reborrow, + ); + append::(&mut ctx, body, vec![], vec![]); + let function = translate(&ctx, &module).unwrap(); + for operation in function.regions[0].blocks[1] + .operations + .iter() + .filter(|operation| matches!(operation.name.as_str(), "llvm.alloca" | "llvm.store")) + { + assert_eq!( + operation.properties["alignment"], + MlirAttribute::Integer { + value: 32, + ty: MlirType::Integer(64) + } + ); + } +} + +/// An array inherits its element's explicit ABI alignment, beyond natural field alignment. +#[test] +fn array_reference_preserves_overaligned_element_abi() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let element = MirStructType::get_with_full_layout( + &mut ctx, + "AlignedValue".into(), + vec!["value".into()], + vec![scalar], + vec![0], + vec![0], + 32, + 32, + ) + .into(); + let array = MirArrayType::get(&mut ctx, element, 2).into(); + let pointer = + MirPtrType::get_with_kind(&mut ctx, array, false, 0, MirPointerKind::SharedRef).into(); + let (module, entry, body) = function(&mut ctx, vec![array]); + let value = entry.deref(&ctx).get_argument(0); + reference( + &mut ctx, + body, + value, + pointer, + false, + MirPointerKindAuthorityAttr::Reborrow, + ); + append::(&mut ctx, body, vec![], vec![]); + let function = translate(&ctx, &module).unwrap(); + let memory: Vec<_> = function.regions[0].blocks[1] + .operations + .iter() + .filter(|operation| matches!(operation.name.as_str(), "llvm.alloca" | "llvm.store")) + .collect(); + assert_eq!(memory.len(), 2); + for operation in memory { + assert_eq!( + operation.properties["alignment"], + MlirAttribute::Integer { + value: 32, + ty: MlirType::Integer(64) + } + ); + } +} + +/// Dropping source attributes must not permit forged kinds, mutability, or address spaces. +#[test] +fn references_reject_invalid_source_pointer_contracts() { + for case in 0..5 { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let wrong = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let (pointee, mutable, address_space, kind, authority) = match case { + 0 => ( + scalar, + false, + 0, + MirPointerKind::RawConst, + MirPointerKindAuthorityAttr::Reborrow, + ), + 1 => ( + scalar, + true, + 0, + MirPointerKind::UniqueRef, + MirPointerKindAuthorityAttr::Reborrow, + ), + 2 => ( + scalar, + false, + 3, + MirPointerKind::SharedRef, + MirPointerKindAuthorityAttr::Reborrow, + ), + 3 => ( + wrong, + false, + 0, + MirPointerKind::SharedRef, + MirPointerKindAuthorityAttr::Reborrow, + ), + _ => ( + scalar, + false, + 0, + MirPointerKind::SharedRef, + MirPointerKindAuthorityAttr::RustCast, + ), + }; + let pointer = + MirPtrType::get_with_kind(&mut ctx, pointee, mutable, address_space, kind).into(); + let (module, entry, body) = function(&mut ctx, vec![scalar]); + let value = entry.deref(&ctx).get_argument(0); + reference(&mut ctx, body, value, pointer, false, authority); + append::(&mut ctx, body, vec![], vec![]); + assert!( + translate(&ctx, &module).is_err(), + "accepted invalid reference case {case}" + ); + } +} + +/// A reference cannot bypass missing canonical bool or generic pointer storage adapters. +#[test] +fn references_reject_payloads_requiring_storage_adapters() { + let mut ctx = context(); + let scalar = FP32Type::get(&ctx).into(); + let flag = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let flags = MirArrayType::get(&mut ctx, flag, 2).into(); + let capture = MirStructType::get( + &mut ctx, + "BoolCapture".into(), + vec!["flag".into()], + vec![flag], + ) + .into(); + let shared = MirPtrType::get_shared(&mut ctx, scalar, false).into(); + for payload in [flag, flags, capture, shared] { + let pointer = + MirPtrType::get_with_kind(&mut ctx, payload, false, 0, MirPointerKind::SharedRef) + .into(); + let (module, entry, body) = function(&mut ctx, vec![payload]); + let value = entry.deref(&ctx).get_argument(0); + reference( + &mut ctx, + body, + value, + pointer, + false, + MirPointerKindAuthorityAttr::Reborrow, + ); + append::(&mut ctx, body, vec![], vec![]); + assert!( + translate(&ctx, &module).is_err(), + "accepted payload requiring a storage adapter" + ); + } +} diff --git a/cuda-oxide/crates/dialect-mir/src/ops/call.rs b/cuda-oxide/crates/dialect-mir/src/ops/call.rs index c4d48cbaac..8fde9fe530 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/call.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/call.rs @@ -91,6 +91,38 @@ impl MirCallOp { self.set_attr_external_callee_type(ctx, TypeAttr::new(signature)); self.set_attr_call_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::AbiBoundary); } + + /// Whether the call carries Rust's logical `()` result for a void callee. + /// Explicitly declared aggregate results, including empty tuples, are not + /// this bridge and must retain their ABI result when exported. + pub fn has_logical_unit_result_bridge(&self, ctx: &Context) -> bool { + let op = self.get_operation().deref(ctx); + if op.get_num_results() != 1 + || !op + .get_result(0) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|tuple| tuple.get_types().is_empty()) + { + return false; + } + let Some(name) = self.get_attr_callee(ctx) else { + return false; + }; + let signature = match resolve_call(ctx, self.get_operation(), name.as_str()) { + CallResolution::AttachedResolved(callee) => Some(callee.get_type(ctx).into()), + _ => self.get_attr_external_callee_type(ctx).map(|attr| { + attr_cast::(&*attr) + .expect("TypeAttr must implement TypedAttrInterface") + .get_type(ctx) + }), + }; + signature.is_some_and(|ty| { + type_cast::(&*ty.deref(ctx)) + .is_some_and(|function| function.res_types().is_empty()) + }) + } } impl Verify for MirCallOp { diff --git a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs index fc2d7f51ed..e2c99c62b9 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs @@ -141,10 +141,38 @@ impl Verify for MirReturnOp { pub struct MirGotoOp; impl MirGotoOp { + /// Operation attribute key marking a loop backedge whose unroll factor was + /// already applied by `mir-transforms`. + const LOOP_UNROLL_DISABLE_KEY: &str = "loop_unroll_disable"; + /// Create a new MirGotoOp wrapper. pub fn new(op: Ptr) -> Self { MirGotoOp { op } } + + /// Mark this branch as a loop backedge that downstream LLVM passes must + /// not unroll again. + pub fn set_loop_unroll_disable(&self, ctx: &mut Context) { + let key = + pliron::identifier::Identifier::try_new(Self::LOOP_UNROLL_DISABLE_KEY.to_string()) + .expect("valid loop_unroll_disable attribute key"); + self.get_operation() + .deref_mut(ctx) + .attributes + .set(key, UnrollAttr(1)); + } + + /// Whether downstream LLVM unrolling should be disabled on this backedge. + pub fn is_loop_unroll_disable(&self, ctx: &Context) -> bool { + let key = + pliron::identifier::Identifier::try_new(Self::LOOP_UNROLL_DISABLE_KEY.to_string()) + .expect("valid loop_unroll_disable attribute key"); + self.get_operation() + .deref(ctx) + .attributes + .get::(&key) + .is_some() + } } impl Verify for MirGotoOp { diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs index a9495a2ce2..de90e2d015 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for the generated active warp mask. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs index 725bccc422..59872c259e 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Cluster Launch Control operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs index a625b4e58a..0874df1091 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated operations for the closed cluster-barrier family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs index 21d02655d5..77e99469a9 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for cluster address mapping and remote shared reads. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs index c043e4082f..b377bcdd5c 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for classic `cp.async` instructions. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs index f5bd427647..d448faad5c 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated PTX debug controls. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs index 5149db0eff..59b9533e91 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated packed integer dot products. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs index 9fd056a5d0..e0a6a491ed 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated warp leader-election operation. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs index 1dc0242394..a7698ed8d2 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated counted-barrier, grid-dependency, and register-control operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs index 31808ec24c..2dc98f2cf0 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for extended floating-point min/max. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs index 956c4a7cab..4f27501c02 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated extended integer min/max. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs index 5c9e6433c4..d2d712a5f0 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation and compatibility carriers for generated `ldmatrix`. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs index 28e36eaf1d..98cd51e0c2 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the basic shared-memory mbarrier lifecycle. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs index 040d990e95..d33ad27ab5 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for extended mbarrier and async-proxy instructions. @@ -119,6 +119,35 @@ fn verify_mbarrier_extended( Ok(()) } +/// Acquires cluster-shared memory at cluster scope, without ordering global memory. +/// Lowers to `fence.acquire.sync_restrict::shared::cluster.cluster;`. +#[pliron_op( + name = "nvvm.fence_acquire_shared_cluster_cluster", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], +)] +pub struct FenceAcquireSharedClusterClusterOp; +impl FenceAcquireSharedClusterClusterOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context) -> Ptr { + Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0) + } +} + +impl Verify for FenceAcquireSharedClusterClusterOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_mbarrier_extended( + ctx, + self.get_operation(), + "nvvm.fence_acquire_shared_cluster_cluster", + MbarrierExtendedShape::ZeroToVoid, + ) + } +} + /// Releases mbarrier initialization at cluster scope. /// Lowers to `fence.mbarrier_init.release.cluster;`. #[pliron_op( @@ -271,6 +300,42 @@ impl Verify for MbarrierArriveClusterOp { } } +/// Arrives at a remote cluster-shared barrier with relaxed memory ordering. +/// Lowers to `mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [
];`. +#[pliron_op( + name = "nvvm.mbarrier_arrive_cluster_relaxed", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<0>], +)] +pub struct MbarrierArriveClusterRelaxedOp; +impl MbarrierArriveClusterRelaxedOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, address: Value) -> Ptr { + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![address], + vec![], + 0, + ) + } +} + +impl Verify for MbarrierArriveClusterRelaxedOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_mbarrier_extended( + ctx, + self.get_operation(), + "nvvm.mbarrier_arrive_cluster_relaxed", + MbarrierExtendedShape::RawU64ToVoid, + ) + } +} + /// Arrives at a CTA-shared barrier and adds expected transaction bytes. /// Lowers to `mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 , [
], ;`. #[pliron_op( @@ -456,6 +521,43 @@ impl Verify for MbarrierTryWaitParityClusterOp { } } +/// Tests barrier parity at cluster scope with relaxed memory ordering. +/// Lowers to `mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 , [
], ;`. +#[pliron_op( + name = "nvvm.mbarrier_try_wait_parity_cluster_relaxed", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], +)] +pub struct MbarrierTryWaitParityClusterRelaxedOp; +impl MbarrierTryWaitParityClusterRelaxedOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, barrier: Value, parity: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 1, Signedness::Signless); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![barrier, parity], + vec![], + 0, + ) + } +} + +impl Verify for MbarrierTryWaitParityClusterRelaxedOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + verify_mbarrier_extended( + ctx, + self.get_operation(), + "nvvm.mbarrier_try_wait_parity_cluster_relaxed", + MbarrierExtendedShape::PointerU32ToI1, + ) + } +} + /// Suspends the executing thread for approximately the requested nanoseconds. /// Lowers to `nanosleep.u32 ;`. #[pliron_op( @@ -493,15 +595,18 @@ impl Verify for NanosleepOp { } pub(super) fn register(ctx: &mut Context) { + FenceAcquireSharedClusterClusterOp::register(ctx); FenceMbarrierInitReleaseClusterOp::register(ctx); FenceProxyAsyncGenericAcquireSharedClusterClusterOp::register(ctx); FenceProxyAsyncGenericReleaseSharedCtaClusterOp::register(ctx); FenceProxyAsyncSharedCtaOp::register(ctx); MbarrierArriveClusterOp::register(ctx); + MbarrierArriveClusterRelaxedOp::register(ctx); MbarrierArriveExpectTxSharedOp::register(ctx); MbarrierArriveExpectTxClusterOp::register(ctx); MbarrierTryWaitSharedOp::register(ctx); MbarrierTryWaitParitySharedOp::register(ctx); MbarrierTryWaitParityClusterOp::register(ctx); + MbarrierTryWaitParityClusterRelaxedOp::register(ctx); NanosleepOp::register(ctx); } diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs index 1f009779b9..eb5c94e9da 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. mod active_mask; diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs index c36bee63cd..ee770d195b 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated in-register matrix transpose. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs index ff8f56f8a7..e97f6cea42 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated packed floating-point arithmetic. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs index c222cf85bb..fb56459682 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for the closed generated packed-atomic family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs index 4223b26248..ac9fdb5246 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated packed conversion. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs index 8c55ea87fb..d57fc2b7f1 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for the closed generated `prmt` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs index f66d2dfe03..b2fbd5c62f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural NVVM operations for the closed generated `redux.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs index b6bf65886f..8839d584fa 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated register-only `mma.sync` variants. @@ -55,6 +55,7 @@ pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4, + Mxf4, } #[pliron_attr(name = "nvvm.register_mma_accumulator", format, verifier = "succ")] @@ -3730,6 +3731,42 @@ impl Verify for RegisterMmaOp { MmaCarrier::F32, ], ), + ( + Some(&RegisterMmaShapeAttr::M16n8k64), + RegisterMmaOperationAttr::Multiply, + RegisterMmaKindAttr::Mxf4, + Some(&RegisterMmaAccumulatorAttr::F32), + Some(&RegisterMmaElementAttr::E2m1), + Some(&RegisterMmaElementAttr::E2m1), + Some(&RegisterMmaLayoutAttr::Row), + Some(&RegisterMmaLayoutAttr::Col), + Some(&RegisterMmaOverflowAttr::NotApplicable), + ) => ( + &[ + MmaCarrier::F32, + MmaCarrier::F32, + MmaCarrier::F32, + MmaCarrier::F32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U32, + MmaCarrier::U16, + MmaCarrier::U16, + MmaCarrier::U32, + MmaCarrier::U16, + MmaCarrier::U16, + ], + &[ + MmaCarrier::F32, + MmaCarrier::F32, + MmaCarrier::F32, + MmaCarrier::F32, + ], + ), ( Some(&RegisterMmaShapeAttr::M16n8k64), RegisterMmaOperationAttr::Multiply, diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs index 1fc9f83f5d..5bf9def0af 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar arithmetic. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs index 35358e18bc..120190439b 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated `f32` to TF32 conversions. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs index bcfd1655b6..be157d8c4f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar math. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs index df9659d2f9..342af76a23 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated sparse MMA variants. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs index dd514a3e42..b42167602b 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural NVVM operations for generated special-register reads. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs index fbde0f174f..cc11851956 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the four generated `stmatrix` stores. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs index 6e0f7527db..ef627dd796 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated CTA synchronization. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs index 95dd733aad..71606692c4 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Core Generation 5 operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs index 2a6454bc62..2f3c2597ce 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Memory Accelerator operations. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs index 61fe550adf..e9246ec9da 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `vote.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs index 1c280bfe18..9cbc2ab8e9 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated warp synchronization. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs index d10423c5c2..5832cc6834 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `match.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs index 4a33bd74f5..e6e3522f3f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for the generated `shfl.sync` family. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs index 8de8f3f129..5edbefc59b 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated WGMMA controls. diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs index 9a085ca4b9..fd60ff734a 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs @@ -74,11 +74,11 @@ impl Verify for CvtaGenericToSharedOffsetOp { }; if !matches!( pointer_ty.address_space, - address_space::GENERIC | address_space::SHARED + address_space::GENERIC | address_space::SHARED | 7 ) { return verify_err!( op.loc(), - "nvvm.cvta_generic_to_shared_offset operand must point to generic or shared memory" + "nvvm.cvta_generic_to_shared_offset operand must point to generic, shared or cluster shared memory" ); } if !is_shared_offset_integer(ctx, op.get_result(0).get_type(ctx)) { @@ -90,3 +90,55 @@ impl Verify for CvtaGenericToSharedOffsetOp { Ok(()) } } + +/// Convert generic, CTA-shared or cluster-shared pointers to cluster offsets. +/// Generic and CTA-shared pointers first cast to AS7; AS7 inputs remain intact. +#[pliron_op( + name = "nvvm.cvta_generic_to_cluster_shared_offset", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], +)] +pub struct CvtaGenericToClusterSharedOffsetOp; + +impl CvtaGenericToClusterSharedOffsetOp { + /// Wrap an existing operation pointer. + pub fn new(op: Ptr) -> Self { + CvtaGenericToClusterSharedOffsetOp { op } + } +} + +impl Verify for CvtaGenericToClusterSharedOffsetOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "nvvm.cvta_generic_to_cluster_shared_offset requires one operand and one result" + ); + } + let pointer_ty = op.get_operand(0).get_type(ctx); + let pointer_ty_obj = pointer_ty.deref(ctx); + let Some(pointer_ty) = pointer_ty_obj.downcast_ref::() else { + return verify_err!( + op.loc(), + "nvvm.cvta_generic_to_cluster_shared_offset operand must be a MIR pointer" + ); + }; + if !matches!( + pointer_ty.address_space, + address_space::GENERIC | address_space::SHARED | 7 + ) { + return verify_err!( + op.loc(), + "nvvm.cvta_generic_to_cluster_shared_offset operand must point to generic, shared or cluster shared memory" + ); + } + if !is_shared_offset_integer(ctx, op.get_result(0).get_type(ctx)) { + return verify_err!( + op.loc(), + "nvvm.cvta_generic_to_cluster_shared_offset result must be u32 or u64" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/crates/dialect-nvvm/tests/ops_test/allowlist.rs b/cuda-oxide/crates/dialect-nvvm/tests/ops_test/allowlist.rs index bc48d66501..fe3fe3fcd9 100644 --- a/cuda-oxide/crates/dialect-nvvm/tests/ops_test/allowlist.rs +++ b/cuda-oxide/crates/dialect-nvvm/tests/ops_test/allowlist.rs @@ -47,6 +47,8 @@ fn handwritten_ops_match_reviewed_allowlist() { ("debug.rs", "AssertFailOp"), ("debug.rs", "VprintfOp"), ("grid.rs", "GridSyncOp"), + // Address-domain conversions are handwritten, like the CTA CVTA. + ("memory.rs", "CvtaGenericToClusterSharedOffsetOp"), ("memory.rs", "CvtaGenericToSharedOffsetOp"), ("wgmma.rs", "WgmmaMakeSmemDescOp"), ("wgmma.rs", "WgmmaMmaM64N64K16F32Bf16Op"), diff --git a/cuda-oxide/crates/llvm-export/src/export/config.rs b/cuda-oxide/crates/llvm-export/src/export/config.rs index 8cd2e86c6e..69b3f41e87 100644 --- a/cuda-oxide/crates/llvm-export/src/export/config.rs +++ b/cuda-oxide/crates/llvm-export/src/export/config.rs @@ -5,8 +5,11 @@ //! Backend configuration traits and built-in implementations. -/// Minimal data layout for PTX mode (default behavior). -pub(super) const NVPTX_DATALAYOUT_PTX: &str = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"; +/// Data layout for PTX mode (default behavior). +/// +/// Matches `nvcc` and `libNVVM` (`NVPTX_DATALAYOUT_FULL`), declaring 32-bit +/// shared-memory pointers (`p3:32:32:32`) and 64-bit generic/global pointers. +pub(super) const NVPTX_DATALAYOUT_PTX: &str = NVPTX_DATALAYOUT_FULL; /// Full NVPTX data layout for libNVVM/LTOIR mode (Blackwell+, modern dialect). /// diff --git a/cuda-oxide/crates/llvm-export/src/export/debug.rs b/cuda-oxide/crates/llvm-export/src/export/debug.rs index ca0bfcfc31..4feb86212a 100644 --- a/cuda-oxide/crates/llvm-export/src/export/debug.rs +++ b/cuda-oxide/crates/llvm-export/src/export/debug.rs @@ -177,9 +177,12 @@ impl<'a> ModuleExportState<'a> { "!{dwarf_version_id} = !{{i32 2, !\"Dwarf Version\", i32 2}}" ) .unwrap(); + // NVVM requires Error (1), including when linking native CUDA LTO-IR. + // Ordinary LLVM debug metadata keeps its Warning (2) merge behavior. + let debug_version_behavior = if self.nvvm_ir_dialect.is_some() { 1 } else { 2 }; writeln!( output, - "!{debug_info_version_id} = !{{i32 2, !\"Debug Info Version\", i32 3}}" + "!{debug_info_version_id} = !{{i32 {debug_version_behavior}, !\"Debug Info Version\", i32 3}}" ) .unwrap(); diff --git a/cuda-oxide/crates/llvm-export/src/export/function.rs b/cuda-oxide/crates/llvm-export/src/export/function.rs index db8900f9a2..911961c53c 100644 --- a/cuda-oxide/crates/llvm-export/src/export/function.rs +++ b/cuda-oxide/crates/llvm-export/src/export/function.rs @@ -1041,7 +1041,7 @@ impl<'a> ModuleExportState<'a> { // A struct or array parameter reaches the kernel as the bytes the // host wrote, so a shared or local pointer inside one is as // impossible to supply as a bare one. Before this recursion a - // `struct { p: SharedPtr, n: u32 }` parameter compiled and the + // `struct { p: *mut Barrier, n: u32 }` parameter compiled and the // kernel applied `cvta.shared.u64` to host bytes. if is_kernel { for (index, arg_ty) in func_ty.arg_types().iter().enumerate() { diff --git a/cuda-oxide/crates/llvm-export/src/export/module.rs b/cuda-oxide/crates/llvm-export/src/export/module.rs index 1a9e934851..89ade3d557 100644 --- a/cuda-oxide/crates/llvm-export/src/export/module.rs +++ b/cuda-oxide/crates/llvm-export/src/export/module.rs @@ -729,6 +729,22 @@ pub(super) fn export_module_with_externs_impl( emit_nvvmir_version(&mut output, &mut state, config.nvvmir_version()); } + if let Some(attr_id) = state.loop_unroll_disable_attr_id { + writeln!(&mut output).unwrap(); + writeln!( + &mut output, + "!{attr_id} = !{{!\"llvm.loop.unroll.disable\"}}" + ) + .unwrap(); + for &loop_id in &state.loop_unroll_disable_loop_ids { + writeln!( + &mut output, + "!{loop_id} = distinct !{{!{loop_id}, !{attr_id}}}" + ) + .unwrap(); + } + } + // 9. DWARF metadata (if requested and source locations exist) if state.has_debug_metadata() { writeln!(&mut output).unwrap(); @@ -876,6 +892,22 @@ pub(super) fn export_module_to_string_with_config( emit_nvvmir_version(&mut output, &mut state, config.nvvmir_version()); } + if let Some(attr_id) = state.loop_unroll_disable_attr_id { + writeln!(&mut output).unwrap(); + writeln!( + &mut output, + "!{attr_id} = !{{!\"llvm.loop.unroll.disable\"}}" + ) + .unwrap(); + for &loop_id in &state.loop_unroll_disable_loop_ids { + writeln!( + &mut output, + "!{loop_id} = distinct !{{!{loop_id}, !{attr_id}}}" + ) + .unwrap(); + } + } + // Emit DWARF line-table metadata if debug export requested it and at least // one function had a real source location. if state.has_debug_metadata() { diff --git a/cuda-oxide/crates/llvm-export/src/export/ops.rs b/cuda-oxide/crates/llvm-export/src/export/ops.rs index 1c1216bc43..b9ba63d704 100644 --- a/cuda-oxide/crates/llvm-export/src/export/ops.rs +++ b/cuda-oxide/crates/llvm-export/src/export/ops.rs @@ -668,15 +668,21 @@ impl<'a> ModuleExportState<'a> { } fn emit_br( - &self, + &mut self, op: &ops::BrOp, block_labels: &FxHashMap, String>, output: &mut String, ) -> Result<(), String> { - let op_ref = op.get_operation().deref(self.ctx); + let op_ptr = op.get_operation(); + let op_ref = op_ptr.deref(self.ctx); let dest = op_ref.successors().next().unwrap(); let label = block_labels.get(&dest).ok_or("Missing block label")?; - writeln!(output, " br label %{label}").unwrap(); + if !self.legacy_typed_pointers() && ops::is_loop_unroll_disable(self.ctx, op_ptr) { + let loop_id = self.alloc_loop_unroll_disable_id(); + writeln!(output, " br label %{label}, !llvm.loop !{loop_id}").unwrap(); + } else { + writeln!(output, " br label %{label}").unwrap(); + } Ok(()) } diff --git a/cuda-oxide/crates/llvm-export/src/export/state.rs b/cuda-oxide/crates/llvm-export/src/export/state.rs index c9fe47022a..9a3a918499 100644 --- a/cuda-oxide/crates/llvm-export/src/export/state.rs +++ b/cuda-oxide/crates/llvm-export/src/export/state.rs @@ -216,6 +216,10 @@ pub(super) struct ModuleExportState<'a> { pub(super) debug_declare_used: bool, /// Whether any function emitted `llvm.dbg.value`. pub(super) debug_value_used: bool, + /// Shared `!{!"llvm.loop.unroll.disable"}` metadata id, allocated on first use. + pub(super) loop_unroll_disable_attr_id: Option, + /// Per-backedge `distinct !{!id, !attr_id}` loop metadata ids, in allocation order. + pub(super) loop_unroll_disable_loop_ids: Vec, } #[derive(Clone, Copy, Debug)] @@ -284,6 +288,8 @@ impl<'a> ModuleExportState<'a> { debug_nodes: Vec::new(), debug_declare_used: false, debug_value_used: false, + loop_unroll_disable_attr_id: None, + loop_unroll_disable_loop_ids: Vec::new(), } } @@ -309,6 +315,15 @@ impl<'a> ModuleExportState<'a> { id } + pub(super) fn alloc_loop_unroll_disable_id(&mut self) -> usize { + if self.loop_unroll_disable_attr_id.is_none() { + self.loop_unroll_disable_attr_id = Some(self.alloc_metadata_id()); + } + let loop_id = self.alloc_metadata_id(); + self.loop_unroll_disable_loop_ids.push(loop_id); + loop_id + } + #[cfg(test)] pub(super) fn next_metadata_id(&self) -> usize { self.next_metadata_id diff --git a/cuda-oxide/crates/llvm-export/src/lib.rs b/cuda-oxide/crates/llvm-export/src/lib.rs index ba5dc7d272..11c5aa9157 100644 --- a/cuda-oxide/crates/llvm-export/src/lib.rs +++ b/cuda-oxide/crates/llvm-export/src/lib.rs @@ -1985,6 +1985,18 @@ pub mod ops { /// Op-attribute key for [`LocalMemoryProvenanceAttr`] on `mir.alloca` and /// `llvm.alloca`. const LOCAL_MEMORY_PROVENANCE_KEY: &str = "cuda_oxide_local_memory_provenance"; + const LOOP_UNROLL_DISABLE_KEY: &str = "cuda_oxide_loop_unroll_disable"; + + /// Mark an `llvm.br` backedge so textual export attaches `!llvm.loop` + /// with `!"llvm.loop.unroll.disable"`. + pub fn set_loop_unroll_disable(ctx: &mut Context, op: Ptr) { + set_string_attr(ctx, op, LOOP_UNROLL_DISABLE_KEY, "1".to_string()); + } + + /// Whether an `llvm.br` backedge carries `!"llvm.loop.unroll.disable"`. + pub fn is_loop_unroll_disable(ctx: &Context, op: Ptr) -> bool { + get_string_attr(ctx, op, LOOP_UNROLL_DISABLE_KEY).is_some() + } /// Attach Rust-local provenance to a stack-slot op. pub fn set_local_memory_provenance( diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/debug_info.rs b/cuda-oxide/crates/llvm-export/tests/export_test/debug_info.rs index abc11b7af6..bbb601856c 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/debug_info.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/debug_info.rs @@ -90,6 +90,20 @@ fn line_table_debug_metadata_emits_function_scope_and_instruction_locations() { let ir = export_module_to_string_with_config(&ctx, &module, &config).expect("debug export succeeds"); + assert!(ir.contains("!{i32 2, !\"Debug Info Version\", i32 3}")); + for debug_kind in [DebugKind::LineTables, DebugKind::Full] { + let nvvm_config = DebugConfig { + inner: NvvmExportConfig::new(NvvmIrDialect::Modern), + debug_kind, + }; + let nvvm_ir = export_module_to_string_with_config(&ctx, &module, &nvvm_config) + .expect("NVVM debug export succeeds"); + assert!( + nvvm_ir.contains("!{i32 1, !\"Debug Info Version\", i32 3}"), + "NVVM requires Error merge behavior for native LTO linkage:\n{nvvm_ir}" + ); + } + let define_line = ir .lines() .find(|line| line.starts_with("define ")) diff --git a/cuda-oxide/crates/mir-importer/Cargo.toml b/cuda-oxide/crates/mir-importer/Cargo.toml index 3f32e242a4..ff7cc7d5b1 100644 --- a/cuda-oxide/crates/mir-importer/Cargo.toml +++ b/cuda-oxide/crates/mir-importer/Cargo.toml @@ -16,6 +16,8 @@ llvm-export = { workspace = true } dialect-mir = { workspace = true } dialect-iket = { workspace = true } dialect-nvvm = { workspace = true } +dialect-cute = { workspace = true } +cute-layout = { workspace = true } mir-lower = { workspace = true } oxide-artifacts = { workspace = true } cuda-oxide-codegen = { workspace = true } diff --git a/cuda-oxide/crates/mir-importer/src/lib.rs b/cuda-oxide/crates/mir-importer/src/lib.rs index 2909011287..796f34e5cc 100644 --- a/cuda-oxide/crates/mir-importer/src/lib.rs +++ b/cuda-oxide/crates/mir-importer/src/lib.rs @@ -86,6 +86,7 @@ extern crate rustc_middle; extern crate rustc_public; extern crate rustc_public_bridge; extern crate rustc_span; +extern crate rustc_target; /// Value used for every rustc `Operand::RuntimeChecks` query in device MIR. /// @@ -101,6 +102,7 @@ pub mod pipeline; // under it, since every item in it would count as reachable API. pub(crate) mod translator; +pub use cuda_oxide_codegen::__private::{CutlassBackendConfig, DeviceBackend}; pub use error::{TranslationErr, TranslationResult}; pub use pipeline::{ CollectedFunction, CompilationArtifactKind, CompilationResult, DebugGlobalVariableIdentity, @@ -109,6 +111,7 @@ pub use pipeline::{ build_debug_global_variable_info, build_debug_shared_array_variable_info, device_static_global_key, run_pipeline, }; +pub use translator::cute::is_cute_compiler_hook_path; pub use translator::facts::KnownDefs; pub use translator::terminator::drop_glue::{drop_glue_is_noop, drop_instance_is_noop}; pub use translator::terminator::is_panic_entry_path; diff --git a/cuda-oxide/crates/mir-importer/src/pipeline.rs b/cuda-oxide/crates/mir-importer/src/pipeline.rs index 12d74984ef..9b587d233e 100644 --- a/cuda-oxide/crates/mir-importer/src/pipeline.rs +++ b/cuda-oxide/crates/mir-importer/src/pipeline.rs @@ -38,8 +38,8 @@ //! Override with `CUDA_OXIDE_TARGET=` environment variable. use cuda_oxide_codegen::__private::{ - BackendOptions, ModuleArtifactKind, ModulePipelineRequest, OutputFiles, PipelineTrace, - append_to_module, compile_translated_module, verify_operation, + BackendOptions, DeviceBackend, ModuleArtifactKind, ModulePipelineRequest, OutputFiles, + PipelineTrace, append_to_module, compile_translated_module, verify_operation, }; pub use cuda_oxide_codegen::__private::{ DeviceArchHint, DeviceExternAttrs, DeviceExternDecl, PipelineError, @@ -189,6 +189,13 @@ pub struct PipelineConfig { pub output_dir: std::path::PathBuf, /// Base name for output files (e.g., `"kernel"` → `kernel.ll`, `kernel.ptx`). pub output_name: String, + /// Optional exact path for a textual CUTLASS MLIR observation. + /// + /// The export observes the prepared semantic module before the selected + /// backend continuation; failures are compilation errors. + pub mlir_output: Option, + /// Device artifact backend. The ordinary backend remains the default. + pub device_backend: DeviceBackend, /// Print progress messages to stdout. pub verbose: bool, /// Dump the `dialect-mir` module after translation (for debugging). @@ -263,6 +270,8 @@ impl Default for PipelineConfig { Self { output_dir: std::env::current_dir().unwrap_or_else(|_| ".".into()), output_name: "kernel".to_string(), + mlir_output: None, + device_backend: DeviceBackend::Native, verbose: true, show_mir_dialect: false, show_llvm_dialect: false, @@ -521,6 +530,7 @@ pub fn run_pipeline( let mut legaliser = Legaliser::default(); let mut kernel_launch_bounds = BTreeMap::new(); + let mut expected_kernels = Vec::new(); // Step 3: Translate all functions for func in functions { @@ -583,6 +593,9 @@ pub fn run_pipeline( { kernel_launch_bounds.insert(func.export_name.clone(), bounds); } + if func.is_kernel { + expected_kernels.push(func.export_name.clone()); + } // Append to module append_to_module(&ctx, module_op_ptr, func_op_ptr); @@ -592,6 +605,9 @@ pub fn run_pipeline( let ptx_path = config .output_dir .join(format!("{}.ptx", config.output_name)); + let cubin_path = config + .output_dir + .join(format!("{}.cubin", config.output_name)); let stale_artifacts = stale_compilation_artifact_paths(&config.output_dir, &config.output_name); let backend_options = backend_options_for(config); @@ -601,9 +617,16 @@ pub fn run_pipeline( config.emit_nvvm_ir, &backend_options, config.debug_kind, + match &config.device_backend { + DeviceBackend::Native => None, + DeviceBackend::CutlassMlir(config) => Some(config), + }, + &expected_kernels, OutputFiles { llvm_ir: &ll_path, ptx: &ptx_path, + cubin: &cubin_path, + mlir: config.mlir_output.as_deref(), stale_before_export: &stale_artifacts, }, PipelineTrace { @@ -645,6 +668,15 @@ pub fn run_pipeline( allow_fma_contraction: config.allow_fma_contraction, kernel_launch_bounds, }), + ModuleArtifactKind::Cubin => Ok(CompilationResult { + artifact_path: cubin_path, + artifact_kind: CompilationArtifactKind::Cubin, + ll_path, + ptx_path, + target: generated.target, + allow_fma_contraction: config.allow_fma_contraction, + kernel_launch_bounds, + }), } } @@ -769,6 +801,8 @@ mod tests { let config = PipelineConfig::default(); assert_eq!(config.output_name, "kernel"); + assert_eq!(config.mlir_output, None); + assert_eq!(config.device_backend, DeviceBackend::Native); assert!(config.verbose); assert!(!config.show_mir_dialect); assert!(!config.show_llvm_dialect); @@ -1299,6 +1333,8 @@ fn adversarial_export_name() -> u64 { let config = PipelineConfig { output_dir: root.clone(), output_name: "kernel".to_string(), + mlir_output: None, + device_backend: DeviceBackend::Native, verbose: false, show_mir_dialect: false, show_llvm_dialect: false, @@ -1353,6 +1389,8 @@ fn adversarial_export_name() -> u64 { let config = PipelineConfig { output_dir: output_dir.clone(), output_name: "empty".to_string(), + mlir_output: None, + device_backend: DeviceBackend::Native, verbose: false, show_mir_dialect: false, show_llvm_dialect: false, @@ -1445,6 +1483,8 @@ fn adversarial_export_name() -> u64 { let config = PipelineConfig { output_dir: root.clone(), output_name: "extern_only".to_string(), + mlir_output: None, + device_backend: DeviceBackend::Native, verbose: false, show_mir_dialect: false, show_llvm_dialect: false, diff --git a/cuda-oxide/crates/mir-importer/src/translator/body.rs b/cuda-oxide/crates/mir-importer/src/translator/body.rs index cd49e86c91..7149aac141 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/body.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/body.rs @@ -2154,8 +2154,9 @@ pub fn translate_body( for arg_idx in 0..num_args { let local = mir::Local::from(arg_idx + 1); let source_ty = &body.locals()[local].ty; - let Some(validity) = facts::reference_param_validity(source_ty, &fn_abi, arg_idx) - .map_err(|error| input_error_noloc!(TranslationErr::invalid_op(error)))? + let Some(validity) = + facts::reference_param_validity(instance, source_ty, &fn_abi, arg_idx) + .map_err(|error| input_error_noloc!(TranslationErr::invalid_op(error)))? else { continue; }; @@ -2375,7 +2376,7 @@ pub fn translate_body( } if let Some(scope_map) = debug_source_scopes - && debug_kind.variables_enabled() + && debug_kind.line_tables_enabled() { llvm_export::ops::set_debug_source_scope_map(ctx, op_ptr, scope_map); } @@ -2742,6 +2743,138 @@ pub fn cuda_oxide_device_generated_kernel(mut wrapped: Wrapper) -> u32 { } } + /// Import a real MIR body with the source-scope sidecar supplied by the + /// codegen bridge, then verify its inline provenance survives MIR lowering. + /// Line tables need these scopes even when variable debug info is disabled. + #[test] + fn imported_inline_scopes_survive_lowering_in_line_table_and_full_modes() { + use llvm_export::ops::{DebugInlinedScope, DebugSourceScope}; + + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "cuda_oxide_inline_scopes_{}_{}", + std::process::id(), + unique + )); + std::fs::create_dir_all(&root).unwrap(); + let fixture = root.join("entry.rs"); + std::fs::write(&fixture, "pub fn entry() {}\n").unwrap(); + let rustc = std::env::var_os("RUSTC").unwrap_or_else(|| "rustc".into()); + let sysroot = std::process::Command::new(rustc) + .args(["--print", "sysroot"]) + .output() + .unwrap(); + assert!(sysroot.status.success()); + let args = vec![ + "rustc".to_string(), + "--edition=2024".to_string(), + "--crate-type=rlib".to_string(), + "--emit=metadata".to_string(), + format!("--out-dir={}", root.display()), + format!( + "--sysroot={}", + String::from_utf8(sysroot.stdout).unwrap().trim() + ), + fixture.display().to_string(), + ]; + let worker_root = root.clone(); + std::thread::Builder::new() + .stack_size(64 * 1024 * 1024) + .spawn(move || { + rustc_public::run!(&args, || { + let item = rustc_public::all_local_items() + .into_iter() + .find(|item| item.name().ends_with("::entry")) + .unwrap(); + let instance = mono::Instance::try_from(item).unwrap(); + let body = instance.body().unwrap(); + let successors: Vec<_> = body + .blocks + .iter() + .map(|block| block.terminator.successors()) + .collect(); + let caller = DebugSourcePosition { + file: fixture, + line: 1, + column: 1, + }; + let scopes = DebugSourceScopeMap { + scopes: vec![ + DebugSourceScope { + id: 0, + parent: None, + span: Some(caller.clone()), + inlined: None, + }, + DebugSourceScope { + id: 1, + parent: Some(0), + span: Some(DebugSourcePosition { + file: worker_root.join("helper.rs"), + line: 279, + column: 5, + }), + inlined: Some(DebugInlinedScope { + callee_name: "helper::group_rank".to_string(), + callsite: Some(caller), + }), + }, + ], + locations: vec![], + }; + for kind in [DebugKind::Off, DebugKind::LineTables, DebugKind::Full] { + let mut ctx = Context::new(); + crate::translator::register_dialects(&mut ctx); + let module = ModuleOp::new(&mut ctx, "test".try_into().unwrap()); + let region = module.get_operation().deref(&ctx).get_region(0); + let block = region.deref(&ctx).iter(&ctx).next().unwrap(); + let function = translate_body( + &mut ctx, + &body, + &instance, + body.blocks.len(), + &successors, + false, + None, + Some("entry"), + &mut Legaliser::default(), + kind, + Some(&scopes), + None, + ) + .unwrap(); + function.insert_at_back(block, &ctx); + mir_lower::register(&mut ctx); + mir_lower::lower_mir_to_llvm(&mut ctx, module.get_operation()).unwrap(); + let lowered = block + .deref(&ctx) + .iter(&ctx) + .find(|op| { + Operation::get_op::(*op, &ctx).is_some() + }) + .unwrap(); + let expected = match kind { + DebugKind::Off => None, + DebugKind::LineTables | DebugKind::Full => Some(scopes.clone()), + }; + assert_eq!( + llvm_export::ops::debug_source_scope_map(&ctx, lowered), + expected + ); + } + std::ops::ControlFlow::<(), _>::Continue(()) + }) + .expect("fixture compilation succeeds"); + }) + .unwrap() + .join() + .unwrap(); + std::fs::remove_dir_all(root).ok(); + } + #[test] fn inline_intent_reaches_llvm_func_attr_before_export() { for (intent, keyword) in [ diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/emit.rs new file mode 100644 index 0000000000..557bece319 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/emit.rs @@ -0,0 +1,484 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the high-level block-scaled GEMV chain. +//! +//! Pointer arithmetic and packed-number conversion stay out of the importer: +//! +//! ```text +//! slices -> scaled view -> row -> K64 tile -> fragment -> dot +//! ``` + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::CuteTensorFormatAttr; +use dialect_cute::gemv_ops::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, +}; +use dialect_mir::attributes::FieldIndexAttr; +use dialect_mir::ops::MirExtractFieldOp; +use dialect_mir::types::{MirPtrType, MirSliceType}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::types::{IntegerType, Signedness}; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::block_scaled::{BlockScaledRust, BlockScaledStage, decode_block_scaled}; + +/// Insert an operation after the current anchor, or at the start of the block. +fn insert( + ctx: &mut Context, + op: Ptr, + block: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block, ctx); + } +} + +/// Branch from the last emitted effect to the call continuation; reject diverging calls. +fn finish( + ctx: &mut Context, + effect: Ptr, + target: &Option, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target is not supported")) + ); + }; + Ok(emit_goto(ctx, *target, effect, block_map, loc)) +} + +/// Store a result in the unprojected MIR destination, then emit the continuation branch. +#[allow(clippy::too_many_arguments)] +fn finish_value( + ctx: &mut Context, + producer: Ptr, + value: Value, + destination: &mir::Place, + target: &Option, + block: Ptr, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "{what} destination with projections is not supported" + )) + ); + } + let after_store = value_map + .store_local(ctx, destination.local, value, block, Some(producer)) + .unwrap_or(producer); + finish(ctx, after_store, target, block_map, loc, what) +} + +/// Emit one aggregate-field extraction and return its value and insertion anchor. +fn extract_field( + ctx: &mut Context, + aggregate: Value, + index: u32, + result_ty: TypeHandle, + block: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![aggregate], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(index)); + insert(ctx, op, block, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Split a verified byte slice into its data pointer and element count. +fn split_u8_slice( + ctx: &mut Context, + slice: Value, + block: Ptr, + prev: Option>, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, Value, Ptr)> { + let u8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let element_type = { + let slice_ty = slice.get_type(ctx); + let slice_ty = slice_ty.deref(ctx); + let Some(slice_ty) = slice_ty.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} must be a slice")) + ); + }; + slice_ty.element_type() + }; + if element_type != u8_ty { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} must contain u8 storage")) + ); + } + + let data_ty: TypeHandle = MirPtrType::get_generic(ctx, u8_ty, false).into(); + let (data, data_op) = extract_field(ctx, slice, 0, data_ty, block, prev, loc); + let len_ty: TypeHandle = crate::translator::types::get_usize_type(ctx).into(); + let (len, len_op) = extract_field(ctx, slice, 1, len_ty, block, Some(data_op), loc); + Ok((data, len, len_op)) +} + +/// Require the destination's scaled-view stage and return its matrix/vector role and scale format. +fn destination_facts( + body: &mir::Body, + destination: &mir::Place, + expected_stage: BlockScaledStage, + loc: &Location, + what: &str, +) -> TranslationResult { + let ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} destination type: {error:?}" + ))) + })?; + let decoded = decode_block_scaled(&ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what} destination: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} destination is not a block-scaled view" + ))) + })?; + if decoded.stage != expected_stage { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} destination has the wrong view stage")) + ); + } + Ok(decoded) +} + +/// Pair packed-data and scale views into a typed block-scaled tensor. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_make( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 4 { + return input_err!( + loc, + TranslationErr::unsupported( + "block_scaled_make expects value/scale slices, rows, and K".to_string() + ) + ); + } + let facts = destination_facts( + body, + destination, + BlockScaledStage::Full, + &loc, + "block_scaled_make", + )?; + let role = facts.role; + + let (values_slice, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (values_data, values_len, after) = + split_u8_slice(ctx, values_slice, block, after, &loc, "block-scaled values")?; + let (scales_slice, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (scales_data, scales_len, after) = + split_u8_slice(ctx, scales_slice, block, after, &loc, "block-scaled scales")?; + let (rows, after) = translate_operand( + ctx, + body, + &args[2], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (k, after) = translate_operand(ctx, body, &args[3], value_map, block, after, loc.clone())?; + + let values = CuteTensorMake2DOp::new_e2m1(ctx, values_data, values_len, rows, k, role, 1); + let values_op = values.get_operation(); + values_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, values_op, block, after.or(prev)); + + let scales = match facts.scale_format { + CuteTensorFormatAttr::E4M3 => { + CuteTensorMake2DOp::new_e4m3(ctx, scales_data, scales_len, rows, k, role, 1, 16) + } + _ => CuteTensorMake2DOp::new_ue8m0(ctx, scales_data, scales_len, rows, k, role, 1, 16), + }; + let scales_op = scales.get_operation(); + scales_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, scales_op, block, Some(values_op)); + + let values_view = values_op.deref(ctx).get_result(0); + let scales_view = scales_op.deref(ctx).get_result(0); + let scaled = CuteScaledViewMakeOp::new(ctx, values_view, scales_view); + let op = scaled.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(scales_op)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "block_scaled_make", + ) +} + +/// Select one matrix row or vector batch while retaining its scale view. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_row( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported( + "block_scaled_thread_row expects a view, batch, and row".to_string() + ) + ); + } + destination_facts( + body, + destination, + BlockScaledStage::Row, + &loc, + "block_scaled_thread_row", + )?; + let (scaled, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (batch, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + let (row, after) = + translate_operand(ctx, body, &args[2], value_map, block, after, loc.clone())?; + let row_op = CuteScaledViewRowOp::new(ctx, scaled, batch, row); + let op = row_op.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "block_scaled_thread_row", + ) +} + +/// Select a 64-element K tile from a block-scaled row. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_k_tile( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "block_scaled_k_tile expects a row and tile index".to_string() + ) + ); + } + destination_facts( + body, + destination, + BlockScaledStage::KTile64, + &loc, + "block_scaled_k_tile", + )?; + let (row, after) = translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (tile, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + let tile_op = CuteScaledViewKTileOp::new(ctx, row, tile); + let op = tile_op.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "block_scaled_k_tile", + ) +} + +/// Load packed values and scales from a selected K64 tile into a semantic fragment. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("block_scaled_load_k64 expects one K tile".to_string()) + ); + } + destination_facts( + body, + destination, + BlockScaledStage::Fragment64, + &loc, + "block_scaled_load_k64", + )?; + let (tile, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + // The safe slice constructor knows only byte alignment. These stronger + // selected-address promises belong to the unsafe Rust load boundary. + let load = CuteScaledViewLoadOp::new(ctx, tile, 16, 4); + let op = load.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "block_scaled_load_k64", + ) +} + +/// Accumulate the scaled K64 matrix/vector dot product into an FP32 result. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_dot( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported( + "block_scaled_dot_k64 expects matrix/vector fragments and acc".to_string() + ) + ); + } + let (matrix, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (vector, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + let (acc, after) = + translate_operand(ctx, body, &args[2], value_map, block, after, loc.clone())?; + let dot = CuteDotOp::new(ctx, matrix, vector, acc); + let op = dot.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "block_scaled_dot_k64", + ) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs new file mode 100644 index 0000000000..cd2badaa92 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/block_scaled/mod.rs @@ -0,0 +1,351 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Rust type decoding for the block-scaled GEMV flow. +//! +//! The four Rust view types keep one selection story visible: +//! +//! ```text +//! BlockScaledTensor -> BlockScaledThreadRow -> BlockScaledTile64 +//! | +//! v load +//! LoadedBlockScaledTile64 +//! ``` +//! +//! Only the exact E2M1 + UE8M0, 16-values-per-scale path becomes a ghost CuTe +//! type. Other block-scaled types keep their ordinary Rust representation. + +pub(crate) mod emit; + +use crate::error::{TranslationErr, TranslationResult}; +use dialect_cute::attributes::{ + CuteScaledLayoutAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, + CuteTensorLayoutAttr, CuteTensorRoleAttr, +}; +use dialect_cute::types::{CuteFragmentType, CuteScaledViewType, CuteTensorViewType}; +use pliron::builtin::types::{FP32Type, IntegerType, Signedness}; +use pliron::context::Context; +use pliron::input_error_noloc; +use pliron::r#type::TypeHandle; +use rustc_public::ty::{GenericArgs, RigidTy, Ty, TyKind}; + +use super::tensor::{canonical_adt_path, const_arg, exact_schema, type_arg, unsigned_const}; + +/// Exact defining path of `cute_rs::block_scaled::tensor::BlockScaledTensor`; public reexports do not change it. +const BLOCK_SCALED_TENSOR: &str = "cute_rs::block_scaled::tensor::BlockScaledTensor"; +/// Exact defining path of `cute_rs::block_scaled::gemv::BlockScaledThreadRow`; public reexports do not change it. +const BLOCK_SCALED_THREAD_ROW: &str = "cute_rs::block_scaled::gemv::BlockScaledThreadRow"; +/// Exact defining path of `cute_rs::block_scaled::gemv::BlockScaledTile64`; public reexports do not change it. +const BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::gemv::BlockScaledTile64"; +/// Exact defining path of `cute_rs::block_scaled::gemv::LoadedBlockScaledTile64`; public reexports do not change it. +const LOADED_BLOCK_SCALED_TILE_64: &str = "cute_rs::block_scaled::gemv::LoadedBlockScaledTile64"; +/// Exact defining path of `cute_rs::block_scaled::gemv::Nvfp4ThreadRow`; public reexports do not change it. +const NVFP4_THREAD_ROW: &str = "cute_rs::block_scaled::gemv::Nvfp4ThreadRow"; +/// Exact defining path of `cute_rs::block_scaled::gemv::Nvfp4Tile64`; public reexports do not change it. +const NVFP4_TILE_64: &str = "cute_rs::block_scaled::gemv::Nvfp4Tile64"; +/// Exact defining path of `cute_rs::block_scaled::gemv::LoadedNvfp4Tile64`; public reexports do not change it. +const LOADED_NVFP4_TILE_64: &str = "cute_rs::block_scaled::gemv::LoadedNvfp4Tile64"; +/// Exact defining path of `cute_rs::block_scaled::layouts::KMajor`; public reexports do not change it. +const K_MAJOR: &str = "cute_rs::block_scaled::layouts::KMajor"; +/// Exact defining path of `cute_rs::block_scaled::Mkl`; public reexports do not change it. +const MKL: &str = "cute_rs::block_scaled::Mkl"; +/// Exact defining path of `cute_rs::block_scaled::Nkl`; public reexports do not change it. +const NKL: &str = "cute_rs::block_scaled::Nkl"; +/// Exact defining path of `cute_rs::numeric::Float4E2M1FN`; public reexports do not change it. +const FLOAT4_E2M1FN: &str = "cute_rs::numeric::Float4E2M1FN"; +/// Exact defining path of `cute_rs::numeric::Float8E8M0FNU`; public reexports do not change it. +const FLOAT8_E8M0FNU: &str = "cute_rs::numeric::Float8E8M0FNU"; +/// Exact defining path of `cute_rs::numeric::Float8E4M3FN`; public reexports do not change it. +const FLOAT8_E4M3FN: &str = "cute_rs::numeric::Float8E4M3FN"; + +/// Which stage of the block-scaled view chain one Rust ADT represents. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum BlockScaledStage { + /// The complete scaled matrix/vector view before row selection. + Full, + /// One selected output row or vector batch. + Row, + /// A selected 64-element K tile that has not been loaded. + KTile64, + /// Loaded packed values and scales for one K64 dot product. + Fragment64, +} + +/// Static facts used to classify calls and build ghost types. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct BlockScaledRust { + /// Matrix or vector role carried by the source layout marker. + pub role: CuteTensorRoleAttr, + /// Scale format (`UE8M0` or `E4M3`). + pub scale_format: CuteTensorFormatAttr, + /// Selection stage represented by the scaled source view. + pub stage: BlockScaledStage, +} + +/// Return the defining path and generic arguments of an ADT, or no match for other types. +fn adt_path(ty: &Ty) -> Option<(String, GenericArgs)> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return None; + }; + Some((canonical_adt_path(&def), args)) +} + +/// Match the exact marker definition and require its expected empty generic schema. +fn exact_marker(ty: &Ty, expected: &str) -> Result<(), String> { + let Some((path, args)) = adt_path(ty) else { + return Err(format!("`{expected}` must be an ADT marker")); + }; + if path != expected { + return Err(format!("expected `{expected}`, found `{path}`")); + } + exact_schema(&args, &[], expected) +} + +/// Decode the exact role marker; reject unsupported or look-alike role types. +fn decode_role(ty: &Ty) -> Result { + let Some((path, args)) = adt_path(ty) else { + return Err("block-scaled role must be an ADT marker".to_string()); + }; + exact_schema(&args, &[], &path)?; + match path.as_str() { + MKL => Ok(CuteTensorRoleAttr::Mkl), + NKL => Ok(CuteTensorRoleAttr::Nkl), + _ => Err(format!( + "block-scaled role must be `{MKL}` or `{NKL}`, found `{path}`" + )), + } +} + +/// Require a K-major layout and decode its matrix or vector role. +fn decode_k_major_role(ty: &Ty) -> Result { + let Some((path, args)) = adt_path(ty) else { + return Err("block-scaled value layout must be KMajor".to_string()); + }; + if path != K_MAJOR { + return Err(format!("expected `{K_MAJOR}`, found `{path}`")); + } + exact_schema(&args, &["Type"], K_MAJOR)?; + decode_role(&type_arg(&args, 0, K_MAJOR)?) +} + +/// Decode one exact Rust ADT in the GEMV block-scaled chain. +pub(crate) fn decode_block_scaled(ty: &Ty) -> Result, String> { + let Some((path, args)) = adt_path(ty) else { + return Ok(None); + }; + + let decoded = match path.as_str() { + BLOCK_SCALED_TENSOR => { + exact_schema( + &args, + &["Lifetime", "Type", "Type", "Const", "Type"], + BLOCK_SCALED_TENSOR, + )?; + exact_marker(&type_arg(&args, 1, BLOCK_SCALED_TENSOR)?, FLOAT4_E2M1FN)?; + let scale_ty = type_arg(&args, 2, BLOCK_SCALED_TENSOR)?; + let scale_format = if exact_marker(&scale_ty, FLOAT8_E8M0FNU).is_ok() { + CuteTensorFormatAttr::UE8M0 + } else { + exact_marker(&scale_ty, FLOAT8_E4M3FN)?; + CuteTensorFormatAttr::E4M3 + }; + let values_per_scale = unsigned_const( + &const_arg(&args, 3, BLOCK_SCALED_TENSOR)?, + "block-scaled values per scale", + )?; + if values_per_scale != 16 { + return Err(format!( + "GEMV block-scaled views require 16 values per scale, found {values_per_scale}" + )); + } + BlockScaledRust { + role: decode_k_major_role(&type_arg(&args, 4, BLOCK_SCALED_TENSOR)?)?, + scale_format, + stage: BlockScaledStage::Full, + } + } + BLOCK_SCALED_THREAD_ROW | BLOCK_SCALED_TILE_64 => { + exact_schema(&args, &["Lifetime", "Type"], &path)?; + BlockScaledRust { + role: decode_role(&type_arg(&args, 1, &path)?)?, + scale_format: CuteTensorFormatAttr::UE8M0, + stage: if path == BLOCK_SCALED_THREAD_ROW { + BlockScaledStage::Row + } else { + BlockScaledStage::KTile64 + }, + } + } + NVFP4_THREAD_ROW | NVFP4_TILE_64 => { + exact_schema(&args, &["Lifetime", "Type"], &path)?; + BlockScaledRust { + role: decode_role(&type_arg(&args, 1, &path)?)?, + scale_format: CuteTensorFormatAttr::E4M3, + stage: if path == NVFP4_THREAD_ROW { + BlockScaledStage::Row + } else { + BlockScaledStage::KTile64 + }, + } + } + LOADED_BLOCK_SCALED_TILE_64 => { + exact_schema(&args, &["Type"], LOADED_BLOCK_SCALED_TILE_64)?; + BlockScaledRust { + role: decode_role(&type_arg(&args, 0, LOADED_BLOCK_SCALED_TILE_64)?)?, + scale_format: CuteTensorFormatAttr::UE8M0, + stage: BlockScaledStage::Fragment64, + } + } + LOADED_NVFP4_TILE_64 => { + exact_schema(&args, &["Type"], LOADED_NVFP4_TILE_64)?; + BlockScaledRust { + role: decode_role(&type_arg(&args, 0, LOADED_NVFP4_TILE_64)?)?, + scale_format: CuteTensorFormatAttr::E4M3, + stage: BlockScaledStage::Fragment64, + } + } + _ => return Ok(None), + }; + Ok(Some(decoded)) +} + +/// Decode a block-scaled value behind exactly one reference or raw pointer. +pub(crate) fn decode_block_scaled_receiver(ty: &Ty) -> Result, String> { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode_block_scaled(&pointee) + } + _ => decode_block_scaled(ty), + } +} + +/// Build the IR scaled-view type for the decoded role, scale format, and selection stage. +fn scaled_view_type( + ctx: &mut Context, + role: CuteTensorRoleAttr, + scale_format: CuteTensorFormatAttr, + layout: CuteScaledLayoutAttr, +) -> TypeHandle { + let logical: TypeHandle = FP32Type::get(ctx).into(); + let storage: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let values: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 1, + CuteTensorFormatAttr::E2M1, + role, + CuteTensorLayoutAttr::KMajor, + ) + .into(); + let scales: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 1, + scale_format, + role, + CuteTensorLayoutAttr::BlockScaleKMajor(16), + ) + .into(); + CuteScaledViewType::get(ctx, values, scales, role, layout).into() +} + +/// Type-translation hook called before generic ADT lowering. +pub(crate) fn try_translate_block_scaled_type( + ctx: &mut Context, + rust_ty: &Ty, +) -> Option> { + let decoded = match decode_block_scaled(rust_ty) { + Ok(Some(decoded)) => decoded, + Ok(None) => return None, + Err(error) => { + return Some(Err(input_error_noloc!(TranslationErr::unsupported( + format!("invalid cute-rs block-scaled type: {error}") + )))); + } + }; + + Some(Ok(match decoded.stage { + BlockScaledStage::Full => scaled_view_type( + ctx, + decoded.role, + decoded.scale_format, + CuteScaledLayoutAttr::Full, + ), + BlockScaledStage::Row => scaled_view_type( + ctx, + decoded.role, + decoded.scale_format, + CuteScaledLayoutAttr::Row, + ), + BlockScaledStage::KTile64 => scaled_view_type( + ctx, + decoded.role, + decoded.scale_format, + CuteScaledLayoutAttr::KTile(64), + ), + BlockScaledStage::Fragment64 => { + let source = scaled_view_type( + ctx, + decoded.role, + decoded.scale_format, + CuteScaledLayoutAttr::KTile(64), + ); + CuteFragmentType::get(ctx, source).into() + } + })) +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that each source view maps to its expected block-scaled stage. + #[test] + fn stage_names_keep_the_source_flow_visible() { + assert_ne!(BlockScaledStage::Full, BlockScaledStage::Row); + assert_ne!(BlockScaledStage::KTile64, BlockScaledStage::Fragment64); + } + + /// Check that scaled IR views retain their role and selection stage. + #[test] + fn scaled_types_keep_role_and_selection_stage() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + + let matrix = scaled_view_type( + &mut ctx, + CuteTensorRoleAttr::Mkl, + CuteTensorFormatAttr::UE8M0, + CuteScaledLayoutAttr::KTile(64), + ); + let nvfp4 = scaled_view_type( + &mut ctx, + CuteTensorRoleAttr::Nkl, + CuteTensorFormatAttr::E4M3, + CuteScaledLayoutAttr::KTile(64), + ); + let matrix_type = matrix.deref(&ctx); + let matrix = matrix_type + .downcast_ref::() + .expect("scaled view"); + assert_eq!(matrix.role, CuteTensorRoleAttr::Mkl); + assert_eq!(matrix.k_width(), Some(64)); + + let nvfp4_type = nvfp4.deref(&ctx); + let nvfp4 = nvfp4_type + .downcast_ref::() + .expect("nvfp4 scaled view"); + assert_eq!(nvfp4.role, CuteTensorRoleAttr::Nkl); + assert_eq!(nvfp4.k_width(), Some(64)); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/classify.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/classify.rs new file mode 100644 index 0000000000..1a7750de6f --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/classify.rs @@ -0,0 +1,629 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Recognize ergonomic methods using exact defining modules and receiver types. + +use super::registry::CuteFn; +use super::{block_scaled, epilogue, pipeline, scheduler, smem_mma, tensor}; +use rustc_public::mir; + +/// Match an inherent method in one exact defining module. +/// +/// The printed impl owner varies across rustc revisions; callers separately +/// validate the receiver/result ADT. Neither extra module nesting nor a +/// public reexport path is accepted as the defining identity. +fn method_in_module<'a>(path: &'a str, module: &str) -> Option<&'a str> { + let (crate_name, rest) = path.split_once("::")?; + if !matches!(crate_name, "cute_rs" | "cute-rs") { + return None; + } + let rest = rest.strip_prefix(module)?.strip_prefix("::")?; + let (impl_owner, method) = rest.split_once("::")?; + if impl_owner.is_empty() || method.contains("::") { + return None; + } + Some(method) +} + +/// Extract a method only from the exact tensor implementation module. +pub(super) fn tensor_method_name_from_path(path: &str) -> Option<&str> { + method_in_module(path, "tensor") +} + +/// Extract a method from the supported block-scaled tensor or GEMV implementation. +pub(super) fn block_scaled_method_name_from_path(path: &str) -> Option<&str> { + method_in_module(path, "block_scaled::tensor") + .or_else(|| method_in_module(path, "block_scaled::gemv")) +} + +/// Extract a method from the supported tiled-MMA or fragment implementation. +pub(super) fn smem_mma_method_name_from_path<'a>( + path: &'a str, + expected_module: &str, +) -> Option<&'a str> { + method_in_module(path, expected_module) +} + +/// Recognize shared-MMA methods by defining module and receiver/result ADT. +pub(super) fn classify_smem_mma_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + let destination_ty = destination.ty(body.locals()).ok(); + let operand_ty = |index: usize| args.get(index)?.ty(body.locals()).ok(); + + if smem_mma_method_name_from_path(path, "copy::tiled") == Some("from_raw_parts") + && args.len() == 2 + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::decode_shared_tensor(ty).ok().flatten().is_some()) + { + return Some(CuteFn::SharedTensorOverlay); + } + + let method = smem_mma_method_name_from_path(path, "mma::block_scaled::tiled") + .filter(|method| { + matches!( + *method, + "get_slice" | "load_scale_atom_128" | "load_a_128" | "get_b_tile_k64" + ) + }) + .or_else(|| { + smem_mma_method_name_from_path(path, "mma::block_scaled::fragments") + .filter(|method| matches!(*method, "zero" | "pairs_at_unchecked" | "gemm")) + })?; + match method { + "get_slice" + if args.len() == 1 + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::decode_tiled_mma(ty).ok().flatten().is_some()) => + { + Some(CuteFn::TiledMmaSlice) + } + "zero" + if args.is_empty() + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::is_accumulator(ty).unwrap_or(false)) => + { + Some(CuteFn::MmaFragmentFill) + } + "load_scale_atom_128" + if args.len() == 5 + && operand_ty(0).is_some_and(|ty| { + smem_mma::decode_tiled_mma_receiver(&ty) + .ok() + .flatten() + .is_some() + }) + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::is_scale_stage(ty).unwrap_or(false)) => + { + Some(CuteFn::MmaLoadScales) + } + "pairs_at_unchecked" + if args.len() == 2 + && operand_ty(0) + .is_some_and(|ty| smem_mma::is_scale_stage(&ty).unwrap_or(false)) + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::is_scale_k64(ty).unwrap_or(false)) => + { + Some(CuteFn::MmaFragmentSliceK) + } + "load_a_128" + if args.len() == 4 + && operand_ty(0).is_some_and(|ty| { + smem_mma::decode_tiled_mma_receiver(&ty) + .ok() + .flatten() + .is_some() + }) => + { + Some(CuteFn::MmaLoadA) + } + "get_b_tile_k64" + if args.len() == 4 + && destination_ty + .as_ref() + .is_some_and(|ty| smem_mma::decode_b_tile(ty).ok().flatten().is_some()) => + { + Some(CuteFn::MmaPartitionB) + } + "gemm" + if args.len() == 5 + && operand_ty(0) + .is_some_and(|ty| smem_mma::is_accumulator(&ty).unwrap_or(false)) => + { + Some(CuteFn::TiledGemm) + } + _ => None, + } +} + +/// Recognize epilogue and store-pipeline methods with their exact source types. +pub(super) fn classify_epilogue_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + let operand_ty = |index: usize| args.get(index)?.ty(body.locals()).ok(); + let destination_ty = destination.ty(body.locals()).ok(); + + if let Some(method) = smem_mma_method_name_from_path(path, "epilogue::sm120") { + let tile_receiver = operand_ty(0) + .as_ref() + .is_some_and(|ty| epilogue::is_epilogue_tile_receiver(ty).unwrap_or(false)); + return match method { + "from_raw" + if args.len() == 1 + && destination_ty + .as_ref() + .is_some_and(|ty| epilogue::is_epilogue_tile(ty).unwrap_or(false)) => + { + Some(CuteFn::EpilogueOverlay) + } + "get_slice" + if args.len() == 3 + && tile_receiver + && destination_ty.as_ref().is_some_and(|ty| { + epilogue::is_epilogue_warp_slice(ty).unwrap_or(false) + }) => + { + Some(CuteFn::EpilogueWarpSlice) + } + "sync_reusable" if args.len() == 1 && tile_receiver => { + Some(CuteFn::EpilogueSyncReusable) + } + "sync_ready_for_tma" if args.len() == 1 && tile_receiver => { + Some(CuteFn::EpilogueSyncReadyForTma) + } + "tma_half" if args.len() == 1 && tile_receiver => Some(CuteFn::EpilogueHalf), + _ => None, + }; + } + + if smem_mma_method_name_from_path(path, "mma::block_scaled::epilogue") == Some("store_tile") + && args.len() == 2 + && operand_ty(0) + .as_ref() + .is_some_and(|ty| epilogue::is_epilogue_warp_slice_receiver(ty).unwrap_or(false)) + && operand_ty(1) + .as_ref() + .is_some_and(|ty| epilogue::is_accumulator(ty).unwrap_or(false)) + { + return Some(CuteFn::EpilogueStoreFragment); + } + + let method = pipeline_method_name_from_path(path)?; + let store_pipeline = operand_ty(0) + .as_ref() + .and_then(|ty| epilogue::decode_store_pipeline(ty).ok().flatten()); + match method { + "producer_acquire" if args.len() == 1 && store_pipeline == Some(1) => { + Some(CuteFn::TmaStoreAcquire) + } + "producer_commit" if args.len() == 1 && store_pipeline == Some(1) => { + Some(CuteFn::TmaStoreCommit) + } + "producer_tail" if args.len() == 1 && store_pipeline == Some(1) => { + Some(CuteFn::TmaStoreTail) + } + _ => None, + } +} + +/// Return an inherent method name from the exact scheduler module. +/// +/// The impl-block spelling changes between rustc revisions. Exact scheduler +/// and work-tile receiver types below provide the stable identity. +pub(super) fn scheduler_method_name_from_path(path: &str) -> Option<&str> { + method_in_module(path, "scheduler") +} + +/// Return an inherent method name from the exact pipeline module. +/// +/// The impl-block spelling is not stable. Exact pipeline/state Rust types in +/// `classify_pipeline_method_call` provide the method identity. +pub(super) fn pipeline_method_name_from_path(path: &str) -> Option<&str> { + method_in_module(path, "pipeline") +} + +/// Recognize load-pipeline and cursor methods without treating other helpers as primitives. +pub(super) fn classify_pipeline_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + let method = pipeline_method_name_from_path(path)?; + let destination_ty = destination.ty(body.locals()).ok(); + let destination_pipeline = destination_ty + .as_ref() + .and_then(|ty| pipeline::decode_tma_load_pipeline(ty).ok().flatten()); + let destination_state = destination_ty + .as_ref() + .and_then(|ty| pipeline::decode_pipeline_state(ty).ok().flatten()); + let receiver_pipeline = args.first().and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| pipeline::decode_tma_load_pipeline(&ty).ok().flatten()) + }); + let receiver_state = args.first().and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| pipeline::decode_pipeline_state_receiver(&ty).ok().flatten()) + }); + let lifecycle_state = args.get(1).and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| pipeline::decode_pipeline_state_receiver(&ty).ok().flatten()) + }); + + match method { + "from_raw_base" if args.len() == 1 && destination_pipeline.is_some() => { + Some(CuteFn::TmaLoadPipelineMake) + } + "from_raw_parts" if args.len() == 2 && destination_pipeline.is_some() => { + Some(CuteFn::TmaLoadPipelineFromRawPartsRejected) + } + "init" if args.len() == 2 && receiver_pipeline.is_some() => { + Some(CuteFn::TmaLoadPipelineInit) + } + "new" if args.is_empty() && destination_state.is_some() => Some(CuteFn::PipelineStateNew), + "index" if args.len() == 1 && receiver_state.is_some() => Some(CuteFn::PipelineStateSlot), + "advance" if args.len() == 1 && receiver_state.is_some() => { + Some(CuteFn::PipelineStateAdvance) + } + "producer_acquire" + if args.len() == 2 + && receiver_pipeline.is_some() + && lifecycle_state.is_some_and(|state| { + state.role == dialect_cute::attributes::CutePipelineRoleAttr::Producer + }) => + { + Some(CuteFn::PipelineProducerAcquire) + } + "producer_expect_tx" + if args.len() == 2 + && receiver_pipeline.is_some() + && lifecycle_state.is_some_and(|state| { + state.role == dialect_cute::attributes::CutePipelineRoleAttr::Producer + }) => + { + Some(CuteFn::PipelineProducerExpectTx) + } + "consumer_wait" + if args.len() == 2 + && receiver_pipeline.is_some() + && lifecycle_state.is_some_and(|state| { + state.role == dialect_cute::attributes::CutePipelineRoleAttr::Consumer + }) => + { + Some(CuteFn::PipelineConsumerWait) + } + "consumer_release" + if args.len() == 2 + && receiver_pipeline.is_some() + && lifecycle_state.is_some_and(|state| { + state.role == dialect_cute::attributes::CutePipelineRoleAttr::Consumer + }) => + { + Some(CuteFn::PipelineConsumerRelease) + } + "producer_tail" + if args.len() == 2 + && receiver_pipeline.is_some() + && lifecycle_state.is_some_and(|state| { + state.role == dialect_cute::attributes::CutePipelineRoleAttr::Producer + }) => + { + Some(CuteFn::PipelineProducerTail) + } + _ => None, + } +} + +/// Recognize persistent-scheduler methods and validate receiver/result ownership. +pub(super) fn classify_scheduler_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + let method = scheduler_method_name_from_path(path)?; + let destination_ty = destination.ty(body.locals()).ok(); + let destination_scheduler = destination_ty + .as_ref() + .and_then(|ty| scheduler::decode_scheduler(ty).ok().flatten()); + let destination_tile = destination_ty + .as_ref() + .and_then(|ty| scheduler::decode_work_tile(ty).ok().flatten()); + let receiver_scheduler = args.first().and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| scheduler::decode_scheduler_receiver(&ty).ok().flatten()) + }); + let receiver_tile = args.first().and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| scheduler::decode_work_tile_receiver(&ty).ok().flatten()) + }); + + match method { + "new_1d" if args.is_empty() && destination_scheduler.is_some() => { + Some(CuteFn::SchedulerNew1d) + } + "has_work" if args.len() == 1 && receiver_scheduler.is_some() => { + Some(CuteFn::SchedulerHasWork) + } + "get_current_work" + if args.len() == 1 + && matches!((receiver_scheduler, destination_tile), (Some(source), Some(result)) if source == result) => + { + Some(CuteFn::SchedulerCurrentTile) + } + "tile_idx" if args.len() == 1 && receiver_tile.is_some() => { + Some(CuteFn::WorkTileCoordinates) + } + "advance_to_next_work" if args.len() == 1 && receiver_scheduler.is_some() => { + Some(CuteFn::SchedulerAdvance) + } + _ => None, + } +} + +/// Recognize scaled-tensor methods while checking their selection stage and role. +pub(super) fn classify_block_scaled_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + use block_scaled::BlockScaledStage; + use dialect_cute::attributes::CuteTensorRoleAttr; + + let method = block_scaled_method_name_from_path(path)?; + let destination = destination + .ty(body.locals()) + .ok() + .and_then(|ty| block_scaled::decode_block_scaled(&ty).ok().flatten()); + let receiver = args.first().and_then(|operand| { + operand.ty(body.locals()).ok().and_then(|ty| { + block_scaled::decode_block_scaled_receiver(&ty) + .ok() + .flatten() + }) + }); + + match method { + "from_slices" + if args.len() == 4 + && destination.is_some_and(|view| view.stage == BlockScaledStage::Full) => + { + Some(CuteFn::BlockScaledMake) + } + "thread_row" + if args.len() == 3 + && matches!((receiver, destination), (Some(source), Some(result)) + if source.stage == BlockScaledStage::Full + && result.stage == BlockScaledStage::Row + && source.role == result.role) => + { + Some(CuteFn::BlockScaledThreadRow) + } + "k_tile" + if args.len() == 2 + && matches!((receiver, destination), (Some(source), Some(result)) + if source.stage == BlockScaledStage::Row + && result.stage == BlockScaledStage::KTile64 + && source.role == result.role) => + { + Some(CuteFn::BlockScaledKTile) + } + "load" + if args.len() == 1 + && matches!((receiver, destination), (Some(source), Some(result)) + if source.stage == BlockScaledStage::KTile64 + && result.stage == BlockScaledStage::Fragment64 + && source.role == result.role) => + { + Some(CuteFn::BlockScaledLoadK64) + } + "dot_accumulate" + if args.len() == 3 + && receiver.is_some_and(|lhs| { + lhs.stage == BlockScaledStage::Fragment64 && lhs.role == CuteTensorRoleAttr::Mkl + }) + && args.get(1).is_some_and(|operand| { + operand.ty(body.locals()).ok().is_some_and(|ty| { + block_scaled::decode_block_scaled_receiver(&ty) + .ok() + .flatten() + .is_some_and(|rhs| { + rhs.stage == BlockScaledStage::Fragment64 + && rhs.role == CuteTensorRoleAttr::Nkl + }) + }) + }) => + { + Some(CuteFn::BlockScaledDotK64) + } + _ => None, + } +} + +/// Recognize tensor-view methods from exact types and access mode. +pub(super) fn classify_tensor_method_call( + path: &str, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, +) -> Option { + let method = tensor_method_name_from_path(path)?; + + let destination_ty = destination.ty(body.locals()).ok(); + let destination_view = destination_ty + .as_ref() + .and_then(|ty| tensor::decode_tensor_view(ty).ok().flatten()); + let receiver_view = args.first().and_then(|operand| { + operand + .ty(body.locals()) + .ok() + .and_then(|ty| tensor::decode_tensor_view_receiver(&ty).ok().flatten()) + }); + + use dialect_cute::attributes::{CuteTensorAccessAttr, CuteTensorLayoutAttr}; + match method { + "from_slice" + if destination_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadOnly + && view.layout == CuteTensorLayoutAttr::Contiguous1D + }) => + { + Some(CuteFn::MakeTensorRead) + } + "from_disjoint_slice" + if destination_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadWrite + && view.layout == CuteTensorLayoutAttr::Contiguous1D + }) => + { + Some(CuteFn::MakeTensorWrite) + } + "zipped_divide" => match (receiver_view, destination_view) { + (Some(source), Some(result)) + if source.layout == CuteTensorLayoutAttr::Contiguous1D + && matches!(result.layout, CuteTensorLayoutAttr::Zipped1D(_)) + && source.access == result.access => + { + Some(if source.access == CuteTensorAccessAttr::ReadOnly { + CuteFn::ZippedDivideRead + } else { + CuteFn::ZippedDivideWrite + }) + } + _ => None, + }, + "slice" => match (receiver_view, destination_view) { + (Some(source), Some(result)) + if matches!(source.layout, CuteTensorLayoutAttr::Zipped1D(_)) + && matches!(result.layout, CuteTensorLayoutAttr::Tile1D(_)) + && source.access == result.access => + { + Some(if source.access == CuteTensorAccessAttr::ReadOnly { + CuteFn::SliceRead + } else { + CuteFn::SliceWrite + }) + } + _ => None, + }, + // CuTe's `local_tile` for a one-dimensional tiler: one call whose + // emitter produces the same zipped_divide and slice pair as the chain. + "local_tile" => match (receiver_view, destination_view) { + (Some(source), Some(result)) + if source.layout == CuteTensorLayoutAttr::Contiguous1D + && matches!(result.layout, CuteTensorLayoutAttr::Tile1D(_)) + && source.access == result.access => + { + Some(if source.access == CuteTensorAccessAttr::ReadOnly { + CuteFn::LocalTileRead + } else { + CuteFn::LocalTileWrite + }) + } + _ => None, + }, + "is_full" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadOnly + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorTileIsFull) + } + "base" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadOnly + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorTileBase) + } + "load" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadOnly + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorLoadTile) + } + "store" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadWrite + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorStoreTile) + } + // The Tile1D guard keeps `TensorMut::store_at`, ordinary + // Rust with the same method name, on the body-translation path. + "store_at" + if receiver_view.is_some_and(|view| { + view.access == CuteTensorAccessAttr::ReadWrite + && matches!(view.layout, CuteTensorLayoutAttr::Tile1D(_)) + }) => + { + Some(CuteFn::TensorStoreElementAbs) + } + _ => None, + } +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::method_in_module; + + /// Check that recognition uses exact defining modules rather than public reexport paths. + #[test] + fn nested_defining_modules_are_exact_not_reexport_paths() { + for module in [ + "block_scaled::tensor", + "block_scaled::gemv", + "mma::block_scaled::tiled", + "mma::block_scaled::fragments", + "mma::block_scaled::epilogue", + ] { + assert_eq!( + method_in_module(&format!("cute_rs::{module}::{{impl#7}}::method"), module), + Some("method") + ); + for path in [ + format!("other::{module}::{{impl#7}}::method"), + format!("cute_rs::{module}_extra::{{impl#7}}::method"), + format!("cute_rs::{module}::nested::{{impl#7}}::method"), + format!("cute_rs::{module}::method"), + format!( + "cute_rs::{}::{{impl#7}}::method", + module.split("::").next().unwrap() + ), + ] { + assert_eq!(method_in_module(&path, module), None, "{path}"); + } + } + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/dispatch.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/dispatch.rs new file mode 100644 index 0000000000..1eca3961ac --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/dispatch.rs @@ -0,0 +1,763 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Validate recognized calls and dispatch to operation-family emitters. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::values::ValueMap; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::input_error_noloc; +use pliron::location::Location; +use pliron::operation::Operation; +use rustc_public::CrateDef; +use rustc_public::mir; + +use super::classify::{ + classify_block_scaled_method_call, classify_epilogue_method_call, + classify_pipeline_method_call, classify_scheduler_method_call, classify_smem_mma_method_call, + classify_tensor_method_call, +}; +use super::identity::{CuteIdentity, canonical_path, classify_cute_call}; +use super::registry::{CuteFn, CuteFnSpec, SubstKind}; +use super::{ + block_scaled, emit, epilogue, pipeline, scheduler, sm100, sm100_attention, smem_mma, tensor, +}; + +/// Validate the callee's generics against the spec's schema BEFORE any +/// positional extraction, so a cute-rs signature refactor fails at the +/// boundary with a message naming the drift instead of misreading deep in +/// the importer. +fn validate_schema( + spec: &CuteFnSpec, + substs: &rustc_public::ty::GenericArgs, + loc: &Location, +) -> TranslationResult<()> { + let found: Vec<&'static str> = substs + .0 + .iter() + .map(|arg| match arg { + rustc_public::ty::GenericArgKind::Type(_) => "Type", + rustc_public::ty::GenericArgKind::Const(_) => "Const", + rustc_public::ty::GenericArgKind::Lifetime(_) => "Lifetime", + }) + .collect(); + let expected: Vec<&'static str> = spec + .schema + .iter() + .map(|k| match k { + SubstKind::Type => "Type", + SubstKind::Const => "Const", + SubstKind::Lifetime => "Lifetime", + }) + .collect(); + if found != expected { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{}: generics schema mismatch: the importer expects [{}] but the call \ + carries [{}]; cute-rs and the importer are out of sync (was the \ + function's signature refactored without updating CUTE_FNS?)", + spec.canonical, + expected.join(", "), + found.join(", ") + )) + ); + } + Ok(()) +} + +/// Dispatch hook called from `translate_call`. +/// +/// Returns `None` for other crates or an approved source abstraction that has +/// no recognized ergonomic method (the caller uses normal body translation). +/// Returns `Some(result)` when this module owns the call, in which case the +/// shared path must not touch it. Registered primitives validate their generic +/// schema; ergonomic methods instead use their family-specific ADT decoders. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_translate_cute_call( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: &Location, +) -> Option>> { + // Structural callee extraction: only direct calls to a constant FnDef + // can require CuTe recognition. Compiler hooks remain direct-call-only; + // audited source-only modules may also be imported as Fn-trait callbacks. + let mir::Operand::Constant(const_op) = func else { + return None; + }; + let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::FnDef(fn_def, substs)) = + const_op.const_.ty().kind() + else { + return None; + }; + + let (cute_fn, spec) = match classify_cute_call(&fn_def) { + CuteIdentity::NotCuteCrate => return None, + CuteIdentity::SourceAbstraction => { + let path = canonical_path(&fn_def, &fn_def.krate().name.to_string()); + let method = classify_tensor_method_call(&path, body, args, destination) + .or_else(|| classify_block_scaled_method_call(&path, body, args, destination)) + .or_else(|| classify_smem_mma_method_call(&path, body, args, destination)) + .or_else(|| classify_epilogue_method_call(&path, body, args, destination)) + .or_else(|| classify_scheduler_method_call(&path, body, args, destination)) + .or_else(|| classify_pipeline_method_call(&path, body, args, destination))?; + (method, None) + } + CuteIdentity::Known(spec) => (spec.func, Some(spec)), + CuteIdentity::Unknown(path) => { + return Some(Err(input_error_noloc!(TranslationErr::unsupported( + format!( + "`{path}` is in the recognized cute-rs namespace but has no importer \ + handler and is not in an approved source-abstraction module. Add an \ + importer-owned primitive to CUTE_FNS, or deliberately classify its \ + device-safe Rust module as a source abstraction." + ) + )))); + } + }; + + if let Some(spec) = spec + && let Err(e) = validate_schema(spec, &substs, loc) + { + return Some(Err(e)); + } + + Some(match cute_fn { + CuteFn::Attention(kind) => sm100_attention::emit(ctx, body, func, args, kind, target, block_ptr, prev_op, value_map, block_map, loc.clone()), + CuteFn::Sm100(kind) => sm100::emit( + ctx, body, func, args, kind, target, block_ptr, prev_op, value_map, block_map, loc.clone(), + ), + CuteFn::MakeTensorRead => tensor::emit::emit_make( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteTensorAccessAttr::ReadOnly, + ), + CuteFn::MakeTensorWrite => tensor::emit::emit_make( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteTensorAccessAttr::ReadWrite, + ), + CuteFn::ZippedDivideRead | CuteFn::ZippedDivideWrite => tensor::emit::emit_zipped_divide( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SliceRead | CuteFn::SliceWrite => tensor::emit::emit_slice( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::LocalTileRead | CuteFn::LocalTileWrite => tensor::emit::emit_local_tile( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TensorTileIsFull | CuteFn::TensorTileBase => tensor::emit::emit_query( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + matches!(cute_fn, CuteFn::TensorTileIsFull), + ), + CuteFn::TensorLoadTile => tensor::emit::emit_load( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TensorStoreTile => tensor::emit::emit_store( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TensorStoreElementAbs => tensor::emit::emit_store_element_abs( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::BlockScaledMake => block_scaled::emit::emit_make( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::BlockScaledThreadRow => block_scaled::emit::emit_row( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::BlockScaledKTile => block_scaled::emit::emit_k_tile( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::BlockScaledLoadK64 => block_scaled::emit::emit_load( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::BlockScaledDotK64 => block_scaled::emit::emit_dot( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SchedulerNew1d => scheduler::emit::emit_new_1d( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SchedulerHasWork => scheduler::emit::emit_has_work( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SchedulerCurrentTile => scheduler::emit::emit_current_tile( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::WorkTileCoordinates => scheduler::emit::emit_coordinates( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SchedulerAdvance => scheduler::emit::emit_advance( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TmaLoadPipelineMake => pipeline::emit::emit_make( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TmaLoadPipelineFromRawPartsRejected => input_err!( + loc.clone(), + TranslationErr::unsupported( + "PipelineTmaAsync::from_raw_parts is not supported by the v0 semantic pipeline; use one contiguous [full][empty] ring with from_raw_base" + .to_owned() + ) + ), + CuteFn::TmaLoadPipelineInit => pipeline::emit::emit_init( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineStateNew => pipeline::emit::emit_state_new( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineStateSlot => pipeline::emit::emit_state_slot( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineStateAdvance => pipeline::emit::emit_state_advance( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineProducerAcquire => pipeline::emit::emit_producer_acquire( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineProducerExpectTx => pipeline::emit::emit_producer_expect_tx( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineConsumerWait => pipeline::emit::emit_consumer_wait( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineConsumerRelease => pipeline::emit::emit_consumer_release( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::PipelineProducerTail => pipeline::emit::emit_producer_tail( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::SharedTensorOverlay => smem_mma::emit::emit_overlay( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::TiledMmaSlice => smem_mma::emit::emit_mma_slice( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::MmaFragmentFill => smem_mma::emit::emit_fragment_fill( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::MmaLoadScales => smem_mma::emit::emit_load_scales( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::MmaFragmentSliceK => smem_mma::emit::emit_slice_k( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::MmaLoadA => smem_mma::emit::emit_load_a( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::MmaPartitionB => smem_mma::emit::emit_partition_b( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::TiledGemm => smem_mma::emit::emit_tiled_gemm( + ctx, body, args, destination, target, block_ptr, prev_op, value_map, block_map, + loc.clone(), + ), + CuteFn::EpilogueOverlay => epilogue::emit::emit_overlay( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::EpilogueWarpSlice => epilogue::emit::emit_warp_slice( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::EpilogueStoreFragment => epilogue::emit::emit_store_fragment( + ctx, + body, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::EpilogueSyncReusable => epilogue::emit::emit_sync( + ctx, + body, + args, + dialect_cute::attributes::CuteEpilogueSyncPhaseAttr::Reusable, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::EpilogueSyncReadyForTma => epilogue::emit::emit_sync( + ctx, + body, + args, + dialect_cute::attributes::CuteEpilogueSyncPhaseAttr::ReadyForTma, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::EpilogueHalf => epilogue::emit::emit_half( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::TmaStoreAcquire => epilogue::emit::emit_store_pipeline_effect( + ctx, + func, + args, + epilogue::emit::StorePipelineEffect::Acquire, + target, + block_ptr, + prev_op, + block_map, + loc.clone(), + ), + CuteFn::TmaStoreCommit => epilogue::emit::emit_store_pipeline_effect( + ctx, + func, + args, + epilogue::emit::StorePipelineEffect::Commit, + target, + block_ptr, + prev_op, + block_map, + loc.clone(), + ), + CuteFn::TmaStoreTail => epilogue::emit::emit_store_pipeline_effect( + ctx, + func, + args, + epilogue::emit::StorePipelineEffect::Tail, + target, + block_ptr, + prev_op, + block_map, + loc.clone(), + ), + CuteFn::LoadTile => emit::emit_load_tile( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::StoreTile => emit::emit_store_tile( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::StoreGlobalNoAllocate => emit::emit_store_global_no_allocate( + ctx,body,func,args,target,block_ptr,prev_op,value_map,block_map,loc.clone(), + ), + CuteFn::AssumeDiv => emit::emit_assume_div( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::LoadMatrixA => emit::emit_load_matrix( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteMatrixRoleAttr::A, + ), + CuteFn::LoadMatrixB => emit::emit_load_matrix( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteMatrixRoleAttr::B, + ), + CuteFn::LoadMatrixBPair => emit::emit_load_matrix( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteMatrixRoleAttr::BPair, + ), + CuteFn::LoadMatrixBPairKMajor => emit::emit_load_matrix( + ctx, + body, + func, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + dialect_cute::attributes::CuteMatrixRoleAttr::BPairKMajor, + ), + CuteFn::CopyTma2d => emit::emit_copy_tma( + ctx, + body, + func, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::CopyTmaS2g2d => emit::emit_copy_tma_s2g( + ctx, + body, + func, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + CuteFn::CopyG2S => emit::emit_copy_g2s( + ctx, + body, + func, + args, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc.clone(), + ), + }) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs new file mode 100644 index 0000000000..c5c0b8cb98 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/emit.rs @@ -0,0 +1,1473 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! dialect-cute op emission. +//! +//! One function per recognized `cute-rs` API entry point. Per-thread tile +//! calls become one zero-result `cute.copy`; `assume` becomes one SSA +//! fact. Cooperative boundaries become typed semantic operations for the +//! selected backend continuation. Nothing here rebuilds a layout from scalar +//! arithmetic. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::types::translate_type; +use crate::translator::values::ValueMap; +use dialect_cute::gemm_tma_ops::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}; +use dialect_cute::layout::{ComposedLayout, Layout}; +use dialect_cute::ops::{CuteAssumeDivOp, CuteCopyG2SOp, CuteCopyOp}; +use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr}; +use dialect_mir::ops::{MirCastOp, MirExtractFieldOp, MirFieldAddrOp, MirLoadOp, MirPtrOffsetOp}; +use dialect_mir::types::{MirPtrType, MirSliceType, MirStructType, address_space}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::layout::{assume_divisor, tile_elem_and_count}; +use super::smem_mma::emit::aggregate_field; +use super::static_config::{ + decode_copy_g2s_callee, decode_copy_tma_callee, decode_copy_tma_s2g_callee, + decode_load_matrix_callee, +}; + +/// Take the address of field `index` of a by-reference struct argument. +/// +/// The result pointer keeps the carrier's mutability and address space; the +/// field's own type comes from the translated struct, so this helper never +/// guesses an offset. +pub(super) fn struct_field_addr( + ctx: &mut Context, + struct_ptr: Value, + index: usize, + what: &str, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, TypeHandle, Ptr)> { + let looked_up = { + let ty = struct_ptr.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref.downcast_ref::().and_then(|pointer| { + let pointee = pointer.pointee; + let pointee_ref = pointee.deref(ctx); + pointee_ref + .downcast_ref::() + .and_then(|strct| strct.field_types.get(index).copied()) + .map(|field_ty| (field_ty, pointer.is_mutable, pointer.address_space)) + }) + }; + let Some((field_ty, is_mutable, address_space)) = looked_up else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{what} must be a reference to a cute-rs carrier struct with field {index}" + )) + ); + }; + let field_ptr_ty: TypeHandle = MirPtrType::get(ctx, field_ty, is_mutable, address_space).into(); + // Go through the one construction door so `aggregate_ty` is stamped; + // the verifier rejects a `mir.field_addr` without it. + let op = MirFieldAddrOp::build(ctx, struct_ptr, field_ptr_ty, index as u32)?; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block_ptr, prev); + Ok((op.deref(ctx).get_result(0), field_ty, op)) +} + +/// Load one value through a typed pointer. +pub(super) fn load_through( + ctx: &mut Context, + pointer: Value, + pointee: TypeHandle, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![pointee], + vec![pointer], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirLoadOp::new(op).set_volatile(ctx, false); + insert(ctx, op, block_ptr, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Recover the CTA-shared address-space proof at a recognized unsafe API +/// boundary. +/// +/// Rust raw-pointer fields do not carry CUDA address spaces. Consequently, +/// constructing `SmemTile { base, .. }` normalizes an AS3 pointer to the +/// generic AS0 type of its `base: *mut T` field. That normalization is sound +/// for an ordinary Rust aggregate, but `copy_tma_2d` has a stronger contract: +/// its destination must be CTA-local shared memory. The unsafe call is the +/// precise place where the caller asserts that contract, so recover AS3 here +/// instead of teaching the type translator that every generic pointer is +/// shared. +/// +/// Already-shared pointers pass through unchanged. Any future carrier that +/// reaches this boundary with a concrete non-shared address space fails +/// closed rather than being silently reinterpreted. +pub(super) fn recover_cta_shared_pointer( + ctx: &mut Context, + pointer: Value, + expected_pointee: TypeHandle, + what: &str, + block_ptr: Ptr, + prev: Ptr, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let pointer_info = { + let ty = pointer.get_type(ctx); + let ty_ref = ty.deref(ctx); + ty_ref + .downcast_ref::() + .map(|ptr| (ptr.pointee, ptr.is_mutable, ptr.address_space)) + }; + let Some((pointee, is_mutable, source_address_space)) = pointer_info else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} must be a MIR pointer")) + ); + }; + if pointee != expected_pointee || !is_mutable { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{what} must be a mutable pointer to the decoded TMA element type" + )) + ); + } + if source_address_space == address_space::SHARED { + return Ok((pointer, prev)); + } + if source_address_space != address_space::GENERIC { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{what} must be generic or CTA-shared memory, found address space {source_address_space}" + )) + ); + } + + let shared_ty: TypeHandle = + MirPtrType::get(ctx, expected_pointee, true, address_space::SHARED).into(); + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![shared_ty], + vec![pointer], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + insert(ctx, cast, block_ptr, Some(prev)); + Ok((cast.deref(ctx).get_result(0), cast)) +} + +/// Load field `index` of a by-reference carrier struct. When the field is +/// itself a single-field struct (`LeadingDim` is `repr(transparent)` but +/// still translates as a struct), the load reads through that wrapper. +fn struct_field_load( + ctx: &mut Context, + struct_ptr: Value, + index: usize, + what: &str, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let (mut pointer, mut pointee, mut prev_op) = + struct_field_addr(ctx, struct_ptr, index, what, block_ptr, prev, loc)?; + let inner_field = { + let pointee_ref = pointee.deref(ctx); + pointee_ref + .downcast_ref::() + .and_then(|wrapper| (wrapper.field_types.len() == 1).then(|| wrapper.field_types[0])) + }; + if inner_field.is_some() { + let (inner_ptr, inner_ty, inner_op) = + struct_field_addr(ctx, pointer, 0, what, block_ptr, Some(prev_op), loc)?; + pointer = inner_ptr; + pointee = inner_ty; + prev_op = inner_op; + } + let (value, op) = load_through(ctx, pointer, pointee, block_ptr, Some(prev_op), loc); + Ok((value, op)) +} + +/// Extract field `index` from a by-value aggregate (the tile-coordinate +/// tuple), producing a `u64`-typed value. +fn tuple_field( + ctx: &mut Context, + aggregate: Value, + index: usize, + result_ty: TypeHandle, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![aggregate], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block_ptr, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Keep one recognized TMA transfer as three directly connected semantic +/// operations. The view values never enter a Rust local or function ABI: +/// they exist only between this call boundary and the selected backend. +#[allow(clippy::too_many_arguments)] +fn emit_tma_copy_2d_views( + ctx: &mut Context, + descriptor: Value, + smem_base: Value, + smem_capacity: Value, + tile_row: Value, + tile_column: Value, + barrier: Value, + element: TypeHandle, + smem_layout: &ComposedLayout, + smem_alignment_bytes: u64, + block: Ptr, + prev: Option>, + loc: &Location, +) -> Ptr { + let source = CuteTmaGmemViewOp::new(ctx, descriptor, element, smem_layout.clone()); + let source_op = source.get_operation(); + source_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, source_op, block, prev); + let source = source_op.deref(ctx).get_result(0); + + let destination = CuteTmaSmemViewOp::new( + ctx, + smem_base, + smem_capacity, + element, + smem_layout.clone(), + smem_alignment_bytes, + ); + let destination_op = destination.get_operation(); + destination_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, destination_op, block, Some(source_op)); + let destination = destination_op.deref(ctx).get_result(0); + + let copy = CuteTmaCopy2dOp::new(ctx, source, destination, tile_row, tile_column, barrier); + let copy_op = copy.get_operation(); + copy_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, copy_op, block, Some(destination_op)); + copy_op +} + +/// `cute_rs::copy::cooperative::copy_g2s::` +/// +/// The importer's whole job is flattening: read the validated static +/// configuration from the substs, read the runtime facts out of the carrier +/// structs, and emit ONE `cute.copy_g2s` whose nine operands carry every +/// runtime value in a fixed order: +/// +/// ```text +/// &GmemMatrix ──field loads──► base, rows, cols, leading_dim +/// (r, c) ──extracts────► tile_row, tile_column +/// &SmemTile ──field loads──► smem base, capacity +/// tidx ──as is───────► thread index +/// ``` +/// +/// Address math, predication, and target instruction selection happen later +/// in the selected backend, which re-reads the same typed attributes. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_copy_g2s( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 4 { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_g2s expects 4 runtime arguments, got {}", + args.len() + )) + ); + } + let config = match decode_copy_g2s_callee(func) { + Ok(config) => config, + Err(error) => { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_g2s static configuration is invalid: {error}" + )) + ); + } + }; + let elem = translate_type(ctx, &config.element_type)?; + + let (gmem_ref, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let prev = prev.or(prev_op); + let (coord_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + let (smem_ref, prev) = + translate_operand(ctx, body, &args[2], value_map, block_ptr, prev, loc.clone())?; + let (tidx, prev) = + translate_operand(ctx, body, &args[3], value_map, block_ptr, prev, loc.clone())?; + + let (gmem_base, prev) = + struct_field_load(ctx, gmem_ref, 0, "copy_g2s source", block_ptr, prev, &loc)?; + let prev = Some(prev); + let (rows, prev) = + struct_field_load(ctx, gmem_ref, 1, "copy_g2s source", block_ptr, prev, &loc)?; + let prev = Some(prev); + let (cols, prev) = + struct_field_load(ctx, gmem_ref, 2, "copy_g2s source", block_ptr, prev, &loc)?; + let prev = Some(prev); + let (leading_dim, prev) = + struct_field_load(ctx, gmem_ref, 3, "copy_g2s source", block_ptr, prev, &loc)?; + let prev = Some(prev); + + let coordinate_ty = rows.get_type(ctx); + let (tile_row, prev) = tuple_field(ctx, coord_val, 0, coordinate_ty, block_ptr, prev, &loc); + let (tile_col, prev) = tuple_field( + ctx, + coord_val, + 1, + coordinate_ty, + block_ptr, + Some(prev), + &loc, + ); + + let (smem_base, prev) = struct_field_load( + ctx, + smem_ref, + 0, + "copy_g2s destination", + block_ptr, + Some(prev), + &loc, + )?; + let prev = Some(prev); + let (capacity, prev) = struct_field_load( + ctx, + smem_ref, + 1, + "copy_g2s destination", + block_ptr, + prev, + &loc, + )?; + + let copy = CuteCopyG2SOp::new( + ctx, + [ + gmem_base, + rows, + cols, + leading_dim, + tile_row, + tile_col, + smem_base, + capacity, + tidx, + ], + config.atom.bytes, + config.thread_layout.clone(), + config.value_layout.clone(), + config.tile_layout.clone(), + config.smem_layout.clone(), + config.leading_dim_divisor, + elem, + ); + let copy_op = copy.get_operation(); + copy_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, copy_op, block_ptr, Some(prev)); + + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported("copy_g2s call without target not supported".to_string()) + ); + }; + Ok(emit_goto(ctx, *target_idx, copy_op, block_map, loc)) +} + +/// `cute_rs::mma::load_matrix_{a,b}::(src, warp_tile, lane) -> Frag` +/// +/// Flattening only, like `copy_g2s`: the shared base comes out of the +/// carrier, the tuple splits into two coordinates, and the destination +/// local's slot receives the fragment. One zero-result `cute.ldmatrix` +/// carries the layout facts; the selected backend computes per-lane addresses. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load_matrix( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + role: dialect_cute::attributes::CuteMatrixRoleAttr, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "load_matrix expects 3 runtime arguments, got {}", + args.len() + )) + ); + } + let config = match decode_load_matrix_callee(func) { + Ok(config) => config, + Err(error) => { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "load_matrix static configuration is invalid: {error}" + )) + ); + } + }; + let elem: TypeHandle = dialect_mir::types::MirFP16Type::get(ctx).into(); + + let (smem_ref, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let prev = prev.or(prev_op); + let (coord_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + let (lane, prev) = + translate_operand(ctx, body, &args[2], value_map, block_ptr, prev, loc.clone())?; + + let (smem_base, prev) = struct_field_load( + ctx, + smem_ref, + 0, + "load_matrix source", + block_ptr, + prev, + &loc, + )?; + let prev = Some(prev); + + let coordinate_ty: TypeHandle = pliron::builtin::types::IntegerType::get( + ctx, + 64, + pliron::builtin::types::Signedness::Unsigned, + ) + .into(); + let (warp_tile_r, prev) = tuple_field(ctx, coord_val, 0, coordinate_ty, block_ptr, prev, &loc); + let (warp_tile_c, prev) = tuple_field( + ctx, + coord_val, + 1, + coordinate_ty, + block_ptr, + Some(prev), + &loc, + ); + + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "load_matrix destination with projections not supported".to_string() + ) + ); + } + let Some(slot) = value_map.get_slot(destination.local) else { + return input_err!( + loc, + TranslationErr::unsupported( + "load_matrix destination local has no memory slot".to_string() + ) + ); + }; + + let load = dialect_cute::ops::CuteLdmatrixOp::new( + ctx, + [smem_base, warp_tile_r, warp_tile_c, lane, slot], + role, + config.smem_layout.clone(), + elem, + ); + let load_op = load.get_operation(); + load_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, load_op, block_ptr, Some(prev)); + + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported( + "load_matrix call without target not supported".to_string() + ) + ); + }; + Ok(emit_goto(ctx, *target_idx, load_op, block_map, loc)) +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +// Keep these address-space and TMA-view tests next to the helper boundary they +// exercise; the remaining emitters below intentionally follow frontend call order. +#[allow(clippy::items_after_test_module)] +mod tests { + use super::*; + use dialect_cute::types::CuteTmaViewType; + use dialect_mir::ops::MirUndefOp; + use pliron::builtin::types::{IntegerType, Signedness}; + use pliron::common_traits::Verify; + + /// Create a typed pointer block argument for importer boundary tests. + fn pointer_value( + ctx: &mut Context, + block: Ptr, + pointee: TypeHandle, + pointer_address_space: u32, + ) -> (Value, Ptr) { + let pointer_ty: TypeHandle = + MirPtrType::get(ctx, pointee, true, pointer_address_space).into(); + let undef = MirUndefOp::new(ctx, pointer_ty).get_operation(); + insert(ctx, undef, block, None); + (undef.deref(ctx).get_result(0), undef) + } + + /// Read the address space from a pointer-valued test operand. + fn pointer_address_space(ctx: &Context, value: Value) -> u32 { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .unwrap() + .address_space + } + + /// Insert an undefined carrier after the test's current operation. + fn undef_after( + ctx: &mut Context, + block: Ptr, + ty: TypeHandle, + prev: Option>, + ) -> (Value, Ptr) { + let undef = MirUndefOp::new(ctx, ty).get_operation(); + insert(ctx, undef, block, prev); + (undef.deref(ctx).get_result(0), undef) + } + + /// Check that global stores retain the raw mutable pointer authority. + #[test] + fn global_store_destination_keeps_raw_mut_pointer_authority() { + use dialect_mir::{attributes::MirPointerKindAuthorityAttr, types::MirPointerKind}; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let element: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let generic: TypeHandle = crate::translator::facts::mint_ptr_type( + &mut ctx, + element, + address_space::GENERIC, + crate::translator::facts::pointer_origin_of_raw_mutability(mir::Mutability::Mut), + ) + .into(); + let block = BasicBlock::new(&mut ctx, None, vec![generic]); + let value = block.deref(&ctx).arguments().next().unwrap(); + let pointer = generic + .deref(&ctx) + .downcast_ref::() + .unwrap() + .clone(); + let global = global_store_pointer_type(&mut ctx, &pointer); + let op = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![global], + vec![value], + vec![], + 0, + ); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + cast.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::RawAddress); + assert_eq!( + global + .deref(&ctx) + .downcast_ref::() + .unwrap() + .kind, + MirPointerKind::RawMut + ); + assert!(cast.verify(&ctx).is_ok()); + } + + /// Check that TMA accepts only generic or CTA-shared pointer provenance. + #[test] + fn tma_boundary_recovers_only_generic_or_existing_cta_shared_pointers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let elem: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + + let (generic, generic_anchor) = + pointer_value(&mut ctx, block, elem, address_space::GENERIC); + let (recovered, cast) = recover_cta_shared_pointer( + &mut ctx, + generic, + elem, + "test destination", + block, + generic_anchor, + &Location::Unknown, + ) + .unwrap(); + assert_eq!( + pointer_address_space(&ctx, recovered), + address_space::SHARED + ); + assert!(Operation::get_opid(cast, &ctx) == MirCastOp::get_opid_static()); + + let (shared, shared_anchor) = pointer_value(&mut ctx, block, elem, address_space::SHARED); + let (unchanged, unchanged_anchor) = recover_cta_shared_pointer( + &mut ctx, + shared, + elem, + "test destination", + block, + shared_anchor, + &Location::Unknown, + ) + .unwrap(); + assert_eq!(unchanged, shared); + assert_eq!(unchanged_anchor, shared_anchor); + + let (global, global_anchor) = pointer_value(&mut ctx, block, elem, address_space::GLOBAL); + let error = recover_cta_shared_pointer( + &mut ctx, + global, + elem, + "test destination", + block, + global_anchor, + &Location::Unknown, + ) + .unwrap_err() + .to_string(); + assert!(error.contains("generic or CTA-shared memory"), "{error}"); + } + + /// Check that TMA import connects the global source and shared destination views. + #[test] + fn tma_importer_connects_direct_global_and_shared_views() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let smem_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + + let (descriptor, prev) = undef_after(&mut ctx, block, descriptor_ty, None); + let (smem_base, prev) = undef_after(&mut ctx, block, smem_ty, Some(prev)); + let (capacity, prev) = undef_after(&mut ctx, block, u64_ty, Some(prev)); + let (tile_row, prev) = undef_after(&mut ctx, block, u64_ty, Some(prev)); + let (tile_column, prev) = undef_after(&mut ctx, block, u64_ty, Some(prev)); + let (barrier, prev) = undef_after(&mut ctx, block, barrier_ty, Some(prev)); + let layout = ComposedLayout::new( + dialect_cute::layout::Swizzle::new(2, 4, 3), + 0, + "(128,64):(64,1)".parse().unwrap(), + dialect_cute::layout::OffsetUnit::Elements, + ) + .unwrap(); + + let copy_op = emit_tma_copy_2d_views( + &mut ctx, + descriptor, + smem_base, + capacity, + tile_row, + tile_column, + barrier, + u8_ty, + &layout, + 512, + block, + Some(prev), + &Location::Unknown, + ); + let copy = CuteTmaCopy2dOp::wrap(copy_op); + assert!(copy.verify(&ctx).is_ok()); + + let source_op = copy.source(&ctx).defining_op().unwrap(); + let destination_op = copy.destination(&ctx).defining_op().unwrap(); + assert!(Operation::get_opid(source_op, &ctx) == CuteTmaGmemViewOp::get_opid_static()); + assert!(Operation::get_opid(destination_op, &ctx) == CuteTmaSmemViewOp::get_opid_static()); + assert_eq!( + CuteTmaGmemViewOp::wrap(source_op).descriptor(&ctx), + descriptor + ); + let destination = CuteTmaSmemViewOp::wrap(destination_op); + assert_eq!(destination.base(&ctx), smem_base); + assert_eq!(destination.capacity(&ctx), capacity); + + let destination_type = copy.destination(&ctx).get_type(&ctx); + let destination_type = destination_type.deref(&ctx); + let destination_view = destination_type.downcast_ref::().unwrap(); + assert_eq!(destination_view.tensor_view(&ctx).unwrap().alignment.0, 512); + } +} + +/// Insert `op` after `prev` (or at block front when the call is the block's +/// first operation). +fn insert( + ctx: &mut Context, + op: Ptr, + block_ptr: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block_ptr, ctx); + } +} + +/// Extract the data pointer (field 0 of the `{ptr, len}` fat pointer) from a +/// slice value, mirroring `apply_deref_projection`'s slice arm. +fn slice_data_ptr( + ctx: &mut Context, + slice_val: Value, + elem: TypeHandle, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let ptr_ty: TypeHandle = MirPtrType::get_generic(ctx, elem, false).into(); + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![ptr_ty], + vec![slice_val], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(0)); + insert(ctx, op, block_ptr, prev); + (op.deref(ctx).get_result(0), op) +} + +/// `base + idx` (element-scaled, like GEP). +fn ptr_offset( + ctx: &mut Context, + base: Value, + idx: Value, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let res_ty = base.get_type(ctx); + let op = Operation::new( + ctx, + MirPtrOffsetOp::get_concrete_op_info(), + vec![res_ty], + vec![base, idx], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block_ptr, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Reinterpret a pointer's pointee as `new_pointee` (PtrToPtr), keeping its +/// mutability and address space. +fn cast_pointee( + ctx: &mut Context, + ptr: Value, + new_pointee: TypeHandle, + block_ptr: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let (is_mut, addrspace) = { + let ty = ptr.get_type(ctx); + let ty_ref = ty.deref(ctx); + let Some(p) = ty_ref.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported( + "cute-rs tile operand expected to be a pointer".to_string() + ) + ); + }; + (p.is_mutable, p.address_space) + }; + let res_ty: TypeHandle = MirPtrType::get(ctx, new_pointee, is_mut, addrspace).into(); + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![res_ty], + vec![ptr], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(op).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + insert(ctx, op, block_ptr, prev); + Ok((op.deref(ctx).get_result(0), op)) +} + +/// Emit the `cute.copy` plus the mandatory goto epilogue. +#[allow(clippy::too_many_arguments)] +fn finish_with_copy( + ctx: &mut Context, + src: Value, + dst: Value, + n: u64, + elem: TypeHandle, + target: &Option, + block_ptr: Ptr, + prev: Option>, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let copy = CuteCopyOp::new(ctx, src, dst, Layout::contiguous(n as i64), elem); + let copy_op = copy.get_operation(); + copy_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, copy_op, block_ptr, prev); + + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target not supported")) + ); + }; + Ok(emit_goto(ctx, *target_idx, copy_op, block_map, loc)) +} + +/// `cute_rs::copy::vector::load_tile::(src: &[T], idx: usize) -> [T; N]` +/// +/// The returned tile is written directly into the destination local's slot +/// (memory-to-memory `cute.copy`), so there is no SSA result to record. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load_tile( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported("load_tile expects exactly 2 arguments".to_string()) + ); + } + let (elem, n) = tile_elem_and_count(ctx, func, &loc)?; + + let (slice_val, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + { + let ty = slice_val.get_type(ctx); + if ty.deref(ctx).downcast_ref::().is_none() { + return input_err!( + loc, + TranslationErr::unsupported("load_tile source is not a slice value".to_string()) + ); + } + } + let (idx_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + + let (data_ptr, prev) = slice_data_ptr(ctx, slice_val, elem, block_ptr, prev, &loc); + let (src, prev) = ptr_offset(ctx, data_ptr, idx_val, block_ptr, Some(prev), &loc); + + // The destination local's alloca slot IS the tile destination. + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "load_tile destination with projections not supported".to_string() + ) + ); + } + let Some(slot) = value_map.get_slot(destination.local) else { + return input_err!( + loc, + TranslationErr::unsupported( + "load_tile destination local has no memory slot".to_string() + ) + ); + }; + let (dst, prev) = cast_pointee(ctx, slot, elem, block_ptr, Some(prev), &loc)?; + + finish_with_copy( + ctx, + src, + dst, + n, + elem, + target, + block_ptr, + Some(prev), + block_map, + loc, + "load_tile", + ) +} + +/// `cute_rs::copy::vector::assume::(x: usize) -> usize` +/// +/// The result equals the input at runtime, but the SSA edge carries `D`: +/// +/// ```text +/// x ──> cute.assume_div {divisor = D} ──> y ──> later address use +/// M3 reads this fact ──┘ +/// ``` +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_assume_div( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("assume expects exactly 1 argument".to_string()) + ); + } + let divisor = assume_divisor(func, &loc)?; + let (val, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let assume = CuteAssumeDivOp::new(ctx, val, divisor); + let assume_op = assume.get_operation(); + assume_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, assume_op, block_ptr, prev.or(prev_op)); + let assumed = assume_op.deref(ctx).get_result(0); + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "assume destination with projections not supported".to_string() + ) + ); + } + let goto_prev = value_map + .store_local(ctx, destination.local, assumed, block_ptr, Some(assume_op)) + .unwrap_or(assume_op); + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported("assume call without target not supported".to_string()) + ); + }; + Ok(emit_goto(ctx, *target_idx, goto_prev, block_map, loc)) +} + +/// `cute_rs::copy::vector::store_tile::(dst: *mut T, idx: usize, vals: &[T; N])` +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store_tile( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + // The unit return value is never materialized; the copy is the effect. + let _ = destination; + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported("store_tile expects exactly 3 arguments".to_string()) + ); + } + let (elem, n) = tile_elem_and_count(ctx, func, &loc)?; + + let (dst_base, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let (idx_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + let (vals_ref, prev) = + translate_operand(ctx, body, &args[2], value_map, block_ptr, prev, loc.clone())?; + + let (dst, prev) = ptr_offset(ctx, dst_base, idx_val, block_ptr, prev, &loc); + let (src, prev) = cast_pointee(ctx, vals_ref, elem, block_ptr, Some(prev), &loc)?; + + finish_with_copy( + ctx, + src, + dst, + n, + elem, + target, + block_ptr, + Some(prev), + block_map, + loc, + "store_tile", + ) +} + +/// `cute_rs::copy::tma::copy_tma_2d::(desc, (r, c), &mut SmemTile, bar)` +/// +/// The descriptor becomes a global transport view; the selected shared base +/// and exact remaining capacity become a shared transport view. One +/// zero-result semantic copy connects those views to the coordinate and +/// completion barrier. The temporary view values stay in direct SSA and do +/// not change the Rust ABI. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_copy_tma( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 4 { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_tma_2d expects 4 runtime arguments, got {}", + args.len() + )) + ); + } + let config = match decode_copy_tma_callee(func) { + Ok(config) => config, + Err(error) => { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_tma_2d static configuration is invalid: {error}" + )) + ); + } + }; + let elem = translate_type(ctx, &config.element_type)?; + + let (tensor_map, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let prev = prev.or(prev_op); + let (coord_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + let (smem_ref, prev) = + translate_operand(ctx, body, &args[2], value_map, block_ptr, prev, loc.clone())?; + let (barrier, prev) = + translate_operand(ctx, body, &args[3], value_map, block_ptr, prev, loc.clone())?; + + let (smem_base, prev) = struct_field_load( + ctx, + smem_ref, + 0, + "copy_tma_2d destination", + block_ptr, + prev, + &loc, + )?; + let (smem_capacity, prev) = struct_field_load( + ctx, + smem_ref, + 1, + "copy_tma_2d destination capacity", + block_ptr, + Some(prev), + &loc, + )?; + let (smem_base, prev) = recover_cta_shared_pointer( + ctx, + smem_base, + elem, + "copy_tma_2d destination", + block_ptr, + prev, + &loc, + )?; + let prev = Some(prev); + + let coordinate_ty: TypeHandle = pliron::builtin::types::IntegerType::get( + ctx, + 64, + pliron::builtin::types::Signedness::Unsigned, + ) + .into(); + let (tile_r, prev) = tuple_field(ctx, coord_val, 0, coordinate_ty, block_ptr, prev, &loc); + let (tile_c, prev) = tuple_field( + ctx, + coord_val, + 1, + coordinate_ty, + block_ptr, + Some(prev), + &loc, + ); + + let copy_op = emit_tma_copy_2d_views( + ctx, + tensor_map, + smem_base, + smem_capacity, + tile_r, + tile_c, + barrier, + elem, + &config.smem_layout, + config.smem_alignment_bytes, + block_ptr, + Some(prev), + &loc, + ); + + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported( + "copy_tma_2d call without target not supported".to_string() + ) + ); + }; + Ok(emit_goto(ctx, *target_idx, copy_op, block_map, loc)) +} + +/// `cute_rs::copy::tma::copy_tma_s2g_2d::(desc, (r, c), SmemTile)` +/// +/// The recognized boundary recovers the CTA-shared source proof from the +/// carrier, splits the tile coordinate, and creates two short-lived TMA views +/// feeding one semantic store. The store-pipeline acquire/commit/tail calls +/// are recognized separately, so the complete source protocol stays visible. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_copy_tma_s2g( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_tma_s2g_2d expects 3 runtime arguments, got {}", + args.len() + )) + ); + } + let config = match decode_copy_tma_s2g_callee(func) { + Ok(config) => config, + Err(error) => { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "copy_tma_s2g_2d static configuration is invalid: {error}" + )) + ); + } + }; + let elem = translate_type(ctx, &config.element_type)?; + + let (tensor_map, prev) = translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let prev = prev.or(prev_op); + let (coord_val, prev) = + translate_operand(ctx, body, &args[1], value_map, block_ptr, prev, loc.clone())?; + let (smem_tile, prev) = + translate_operand(ctx, body, &args[2], value_map, block_ptr, prev, loc.clone())?; + + let (smem_base, prev) = aggregate_field( + ctx, + smem_tile, + 0, + "copy_tma_s2g_2d source", + block_ptr, + prev, + &loc, + )?; + let (smem_capacity, prev) = aggregate_field( + ctx, + smem_tile, + 1, + "copy_tma_s2g_2d source", + block_ptr, + Some(prev), + &loc, + )?; + let (smem_base, prev) = recover_cta_shared_pointer( + ctx, + smem_base, + elem, + "copy_tma_s2g_2d source", + block_ptr, + prev, + &loc, + )?; + let prev = Some(prev); + + let coordinate_ty: TypeHandle = pliron::builtin::types::IntegerType::get( + ctx, + 64, + pliron::builtin::types::Signedness::Unsigned, + ) + .into(); + let (tile_r, prev) = tuple_field(ctx, coord_val, 0, coordinate_ty, block_ptr, prev, &loc); + let (tile_c, prev) = tuple_field( + ctx, + coord_val, + 1, + coordinate_ty, + block_ptr, + Some(prev), + &loc, + ); + + let destination = + CuteTmaGmemViewOp::new_destination(ctx, tensor_map, elem, config.smem_layout.clone()); + let destination_op = destination.get_operation(); + destination_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, destination_op, block_ptr, Some(prev)); + let destination = destination_op.deref(ctx).get_result(0); + let source = CuteTmaSmemViewOp::new( + ctx, + smem_base, + smem_capacity, + elem, + config.smem_layout.clone(), + config.smem_alignment_bytes, + ); + let source_op = source.get_operation(); + source_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, source_op, block_ptr, Some(destination_op)); + let source = source_op.deref(ctx).get_result(0); + let store = dialect_cute::epilogue_ops::CuteTmaStore2dSemanticOp::new( + ctx, + source, + destination, + tile_r, + tile_c, + ); + let store_op = store.get_operation(); + store_op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, store_op, block_ptr, Some(source_op)); + + let Some(target_idx) = target else { + return input_err!( + loc, + TranslationErr::unsupported( + "copy_tma_s2g_2d call without target not supported".to_string() + ) + ); + }; + Ok(emit_goto(ctx, *target_idx, store_op, block_map, loc)) +} + +/// Preserve a typed cache-aware register-to-global copy until backend selection. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store_global_no_allocate( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + use dialect_cute::copy_ops::CuteStoreGlobalNoAllocateOp; + use dialect_mir::attributes::MirPointerKindAuthorityAttr; + use pliron::common_traits::Verify; + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported("global store expects a destination and register array") + ); + } + let (element, n) = tile_elem_and_count(ctx, func, &loc)?; + let (dst, prev) = translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (src, prev) = translate_operand(ctx, body, &args[1], value_map, block, prev, loc.clone())?; + let (src, source_cast) = cast_pointee(ctx, src, element, block, prev, &loc)?; + let pointer = dst + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned(); + let Some(pointer) = pointer else { + return input_err!( + loc, + TranslationErr::unsupported("global store destination must be a pointer") + ); + }; + if !pointer.is_mutable + || pointer.pointee != element + || !matches!( + pointer.address_space, + address_space::GLOBAL | address_space::GENERIC + ) + { + return input_err!( + loc, + TranslationErr::unsupported( + "global store destination must be a writable global pointer of the element type" + ) + ); + } + let global = global_store_pointer_type(ctx, &pointer); + let (dst, after) = if dst.get_type(ctx) == global { + (dst, source_cast) + } else { + let op = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![global], + vec![dst], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + let cast = MirCastOp::new(op); + cast.set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + if dialect_mir::types::type_contains_concrete_pointer_kind(ctx, global) { + cast.set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::RawAddress); + } + insert(ctx, op, block, Some(source_cast)); + (op.deref(ctx).get_result(0), op) + }; + let store = CuteStoreGlobalNoAllocateOp::new(ctx, dst, src, n); + store + .verify(ctx) + .map_err(|e| pliron::input_error_noloc!(TranslationErr::unsupported(e.to_string())))?; + let op = store.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported("global store requires a continuation") + ); + }; + Ok(emit_goto(ctx, *target, op, block_map, loc)) +} + +/// Translate a global-store destination while preserving raw mutable-pointer authority. +fn global_store_pointer_type(ctx: &mut Context, pointer: &MirPtrType) -> TypeHandle { + crate::translator::facts::mint_ptr_type( + ctx, + pointer.pointee, + address_space::GLOBAL, + crate::translator::facts::pointer_origin_of_ptr_carrier(pointer), + ) + .into() +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/emit.rs new file mode 100644 index 0000000000..4b70bfae0d --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/emit.rs @@ -0,0 +1,493 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the semantic SM120 output-tile protocol. +//! +//! The operations annotate ordinary carriers. The full tile is one pointer, +//! a warp slice is that pointer plus two indices, and both the accumulator and +//! TMA-half structs keep their existing MIR representation. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::{ + CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpilogueSyncPhaseAttr, + CuteTmaStorePipelineAttr, +}; +use dialect_cute::epilogue_ops::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, + CuteEpilogueWarpSliceOp, CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, +}; +use dialect_mir::types::MirFP16Type; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::TypeHandle; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::emit::recover_cta_shared_pointer; +use super::super::epilogue::{ + is_accumulator, is_epilogue_tile, is_epilogue_tile_receiver, is_epilogue_warp_slice_receiver, + one_unsigned_const, tile_type, +}; +use super::super::smem_mma::emit::{ + aggregate_field, build_struct_prefix, destination_rust_type, destination_type, finish, + finish_value, insert, operand_rust_type, pointer_cast, struct_field_type, +}; + +/// Report an unsupported source contract at this CuTe boundary. +fn invalid(what: &str, error: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what}: {error}" + ))) +} + +/// Convert a failed or absent decoded source-type fact into a boundary diagnostic. +fn require_type_fact(result: Result, what: &str) -> TranslationResult<()> { + match result { + Ok(true) => Ok(()), + Ok(false) => Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} has the wrong Rust carrier type") + ))), + Err(error) => Err(invalid(what, error)), + } +} + +/// Find an aggregate field's translated carrier type before extracting it. +fn translated_field_type( + ctx: &mut Context, + body: &mir::Body, + destination: &mir::Place, + field: usize, + what: &str, +) -> TranslationResult<(TypeHandle, TypeHandle)> { + let outer = destination_type(ctx, body, destination, what)?; + let field = struct_field_type(ctx, outer, field, what)?; + Ok((outer, field)) +} + +/// Recover the shared base pointer from an epilogue tile or referenced receiver. +#[allow(clippy::too_many_arguments)] +fn epilogue_base( + ctx: &mut Context, + body: &mir::Body, + operand: &mir::Operand, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, Ptr)> { + require_type_fact( + is_epilogue_tile_receiver(&operand_rust_type(body, operand, what)?), + what, + )?; + let (tile, after) = translate_operand(ctx, body, operand, value_map, block, prev, loc.clone())?; + let (base, after) = aggregate_field(ctx, tile, 0, what, block, after.or(prev), loc)?; + let storage: TypeHandle = MirFP16Type::get(ctx).into(); + recover_cta_shared_pointer(ctx, base, storage, what, block, after, loc) +} + +/// Extract the shared base and warp/lane coordinates of an epilogue slice. +#[allow(clippy::too_many_arguments)] +fn warp_slice_carriers( + ctx: &mut Context, + body: &mir::Body, + operand: &mir::Operand, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, Value, Value, Ptr)> { + require_type_fact( + is_epilogue_warp_slice_receiver(&operand_rust_type(body, operand, what)?), + what, + )?; + let (slice, after) = + translate_operand(ctx, body, operand, value_map, block, prev, loc.clone())?; + let (base, after) = aggregate_field(ctx, slice, 0, what, block, after.or(prev), loc)?; + let (warp, after) = aggregate_field(ctx, slice, 1, what, block, Some(after), loc)?; + let (lane, after) = aggregate_field(ctx, slice, 2, what, block, Some(after), loc)?; + let storage: TypeHandle = MirFP16Type::get(ctx).into(); + let (base, after) = recover_cta_shared_pointer(ctx, base, storage, what, block, after, loc)?; + Ok((base, warp, lane, after)) +} + +/// Overlay the fixed epilogue layout on the caller's shared output storage. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_overlay( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported( + "epilogue_smem_overlay expects one base pointer".to_owned() + ) + ); + } + require_type_fact( + is_epilogue_tile(&destination_rust_type( + body, + destination, + "epilogue overlay", + )?), + "epilogue overlay result", + )?; + let (base, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let storage: TypeHandle = MirFP16Type::get(ctx).into(); + let anchor = after.or(prev).ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "epilogue overlay base produced no insertion anchor".to_owned() + )) + })?; + let (base, anchor) = recover_cta_shared_pointer( + ctx, + base, + storage, + "epilogue overlay base", + block, + anchor, + &loc, + )?; + let tile = tile_type(ctx); + let operation = CuteEpilogueSmemOverlayOp::new(ctx, base, tile); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(anchor)); + + let (outer, base_ty) = translated_field_type(ctx, body, destination, 0, "epilogue overlay")?; + let (base, anchor) = pointer_cast(ctx, operation.base(ctx), base_ty, block, op, &loc)?; + let (result, anchor) = build_struct_prefix(ctx, outer, &[base], block, anchor, &loc)?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "epilogue overlay", + ) +} + +/// Construct the epilogue view owned by one warp and lane. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_warp_slice( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported( + "epilogue_warp_slice expects tile, warp, and lane".to_owned() + ) + ); + } + let (base, after) = epilogue_base( + ctx, + body, + &args[0], + block, + prev, + value_map, + &loc, + "epilogue warp-slice tile", + )?; + let (warp, warp_after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + Some(after), + loc.clone(), + )?; + let anchor = warp_after.unwrap_or(after); + let (lane, lane_after) = translate_operand( + ctx, + body, + &args[2], + value_map, + block, + Some(anchor), + loc.clone(), + )?; + let anchor = lane_after.unwrap_or(anchor); + let tile = tile_type(ctx); + let operation = CuteEpilogueWarpSliceOp::new(ctx, base, warp, lane, tile); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(anchor)); + + let outer = destination_type(ctx, body, destination, "epilogue warp slice")?; + let base_ty = struct_field_type(ctx, outer, 0, "epilogue warp slice")?; + let (base, anchor) = pointer_cast(ctx, operation.base(ctx), base_ty, block, op, &loc)?; + let (result, anchor) = build_struct_prefix( + ctx, + outer, + &[base, operation.warp_id(ctx), operation.lane(ctx)], + block, + anchor, + &loc, + )?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "epilogue warp slice", + ) +} + +/// Store an accumulator fragment through its typed epilogue warp slice. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store_fragment( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "epilogue_store_fragment expects slice and accumulator".to_owned() + ) + ); + } + require_type_fact( + is_accumulator(&operand_rust_type(body, &args[1], "epilogue accumulator")?), + "epilogue accumulator", + )?; + let (base, warp, lane, after) = warp_slice_carriers( + ctx, + body, + &args[0], + block, + prev, + value_map, + &loc, + "epilogue warp slice", + )?; + let (accumulator, accumulator_after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + Some(after), + loc.clone(), + )?; + let after = accumulator_after.unwrap_or(after); + let tile = tile_type(ctx); + let operation = CuteEpilogueStoreFragmentOp::new(ctx, base, warp, lane, accumulator, tile); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + finish(ctx, op, target, block_map, loc, "epilogue store fragment") +} + +/// Emit the reusable or ready-for-TMA handoff with its participant contract. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_sync( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + phase: CuteEpilogueSyncPhaseAttr, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("epilogue sync expects one tile".to_owned()) + ); + } + let (base, after) = epilogue_base( + ctx, + body, + &args[0], + block, + prev, + value_map, + &loc, + "epilogue sync tile", + )?; + let barrier = CuteCountedCtaBarrierAttr::new(2, 0, 8, 9, 32); + let tile = tile_type(ctx); + let operation = CuteEpilogueSyncOp::new(ctx, base, tile, barrier, phase); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + finish(ctx, op, target, block_map, loc, "epilogue sync") +} + +/// Select one contiguous half of the shared epilogue tile for a TMA store. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_half( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("epilogue_half expects one tile".to_owned()) + ); + } + let half = one_unsigned_const(func, "epilogue half") + .map_err(|error| invalid("epilogue half", error))?; + let half = u32::try_from(half).map_err(|_| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "epilogue half index does not fit u32".to_owned() + )) + })?; + let (base, after) = epilogue_base( + ctx, + body, + &args[0], + block, + prev, + value_map, + &loc, + "epilogue half tile", + )?; + let tile = tile_type(ctx); + let operation = CuteEpilogueHalfOp::new(ctx, base, tile, CuteEpilogueHalfAttr(half)); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + + let outer = destination_type(ctx, body, destination, "epilogue half")?; + let base_ty = struct_field_type(ctx, outer, 0, "epilogue half")?; + let (base, anchor) = pointer_cast(ctx, operation.half_base(ctx), base_ty, block, op, &loc)?; + let (result, anchor) = build_struct_prefix( + ctx, + outer, + &[base, operation.capacity(ctx)], + block, + anchor, + &loc, + )?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "epilogue half", + ) +} + +/// Store-pipeline handoff selected by an ergonomic Rust method. +#[derive(Clone, Copy)] +pub(crate) enum StorePipelineEffect { + /// Wait until the oldest shared output slot can be reused. + Acquire, + /// Commit this issuer's pending TMA stores as one group. + Commit, + /// Drain committed stores before releasing their shared source buffers. + Tail, +} + +/// Import a store-pipeline acquire, commit or drain with its static stage count. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store_pipeline_effect( + ctx: &mut Context, + func: &mir::Operand, + args: &[mir::Operand], + effect: StorePipelineEffect, + target: &Option, + block: Ptr, + prev: Option>, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() > 1 { + return input_err!( + loc, + TranslationErr::unsupported( + "TMA store-pipeline operation expects zero or one ZST receiver".to_owned() + ) + ); + } + let stages = one_unsigned_const(func, "TMA store pipeline") + .map_err(|error| invalid("TMA store pipeline", error))?; + if stages != 1 { + return input_err!( + loc, + TranslationErr::unsupported(format!("v0 needs one TMA store stage, got {stages}")) + ); + } + let pipeline = CuteTmaStorePipelineAttr::new(1); + let op = match effect { + StorePipelineEffect::Acquire => CuteTmaStoreAcquireOp::new(ctx, pipeline).get_operation(), + StorePipelineEffect::Commit => CuteTmaStoreCommitOp::new(ctx, pipeline).get_operation(), + StorePipelineEffect::Tail => CuteTmaStoreTailOp::new(ctx, pipeline).get_operation(), + }; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, prev); + finish( + ctx, + op, + target, + block_map, + loc, + "TMA store-pipeline operation", + ) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs new file mode 100644 index 0000000000..d3230058a0 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/epilogue/mod.rs @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural Rust facts for the SM120 output-tile boundary. +//! +//! None of these Rust types becomes a ghost. The importer uses their exact +//! identity to attach one epilogue plan to the same pointer, warp/lane, and +//! accumulator carriers already present in MIR. + +pub(crate) mod emit; + +use dialect_cute::attributes::{CuteEpiloguePlanAttr, CuteTiledMmaPlanAttr}; +use dialect_cute::types::CuteEpilogueTileType; +use dialect_mir::types::MirFP16Type; +use pliron::context::Context; +use pliron::r#type::TypeHandle; +use rustc_public::mir; +use rustc_public::ty::{GenericArgKind, GenericArgs, RigidTy, Ty, TyKind}; + +use super::smem_mma::canonical_data_placement; +use super::tensor::{canonical_adt_path, const_arg, exact_schema, unsigned_const}; + +/// Exact defining path of `cute_rs::epilogue::sm120::Sm120Epilogue128x128`; public reexports do not change it. +const EPILOGUE_TILE: &str = "cute_rs::epilogue::sm120::Sm120Epilogue128x128"; +/// Exact defining path of `cute_rs::epilogue::sm120::Sm120EpilogueWarp128x128`; public reexports do not change it. +const EPILOGUE_WARP_SLICE: &str = "cute_rs::epilogue::sm120::Sm120EpilogueWarp128x128"; +/// Exact defining path of `cute_rs::mma::block_scaled::fragments::Mxf4AccumulatorTile2x8`; public reexports do not change it. +const ACCUMULATOR: &str = "cute_rs::mma::block_scaled::fragments::Mxf4AccumulatorTile2x8"; +/// Exact defining path of `cute_rs::pipeline::PipelineTmaStore`; public reexports do not change it. +const TMA_STORE_PIPELINE: &str = "cute_rs::pipeline::PipelineTmaStore"; + +/// Recognize one exact source ADT rather than a type with a similar printed name. +fn exact_adt(ty: &Ty, expected: &str) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + if canonical_adt_path(&def) != expected { + return Ok(None); + } + Ok(Some(args)) +} + +/// Decode a value directly or through one reference or raw-pointer layer. +fn behind_one_pointer( + ty: &Ty, + decode: impl FnOnce(&Ty) -> Result, +) -> Result { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode(&pointee) + } + _ => decode(ty), + } +} + +/// Check an exact marker identity without accepting unrelated generic arguments. +fn is_exact_marker(ty: &Ty, expected: &str) -> Result { + let Some(args) = exact_adt(ty, expected)? else { + return Ok(false); + }; + exact_schema(&args, &[], expected)?; + Ok(true) +} + +/// Recognize the fixed SM120 shared epilogue tile. +pub(crate) fn is_epilogue_tile(ty: &Ty) -> Result { + is_exact_marker(ty, EPILOGUE_TILE) +} + +/// Recognize the epilogue tile directly or behind one pointer. +pub(crate) fn is_epilogue_tile_receiver(ty: &Ty) -> Result { + behind_one_pointer(ty, is_epilogue_tile) +} + +/// Recognize the fixed warp-owned epilogue view. +pub(crate) fn is_epilogue_warp_slice(ty: &Ty) -> Result { + is_exact_marker(ty, EPILOGUE_WARP_SLICE) +} + +/// Recognize an epilogue warp view directly or behind one pointer. +pub(crate) fn is_epilogue_warp_slice_receiver(ty: &Ty) -> Result { + behind_one_pointer(ty, is_epilogue_warp_slice) +} + +/// Recognize the FP32 accumulator carrier accepted by the epilogue. +pub(crate) fn is_accumulator(ty: &Ty) -> Result { + behind_one_pointer(ty, |ty| is_exact_marker(ty, ACCUMULATOR)) +} + +/// Decode and validate the store pipeline's constant stage count. +pub(crate) fn decode_store_pipeline(ty: &Ty) -> Result, String> { + let decode = |ty: &Ty| -> Result, String> { + let Some(args) = exact_adt(ty, TMA_STORE_PIPELINE)? else { + return Ok(None); + }; + exact_schema(&args, &["Const"], TMA_STORE_PIPELINE)?; + let stages = unsigned_const( + &const_arg(&args, 0, TMA_STORE_PIPELINE)?, + "TMA store-pipeline stage count", + )?; + let stages = u32::try_from(stages) + .map_err(|_| "TMA store-pipeline stage count does not fit u32".to_owned())?; + if stages != 1 { + return Err(format!( + "v0 needs one shared epilogue buffer, got {stages} store-pipeline stages" + )); + } + Ok(Some(stages)) + }; + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode(&pointee) + } + _ => decode(ty), + } +} + +/// Read the one const generic carried by `tma_half` or a store-pipeline helper. +pub(crate) fn one_unsigned_const(func: &mir::Operand, what: &str) -> Result { + let mir::Operand::Constant(constant) = func else { + return Err(format!("{what} must be a direct constant function call")); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err(format!("{what} callee is not a FnDef")); + }; + let constants = args + .0 + .iter() + .filter_map(|arg| match arg { + GenericArgKind::Const(value) => Some(value), + _ => None, + }) + .collect::>(); + if constants.len() != 1 { + return Err(format!( + "{what} needs exactly one const generic, found {}", + constants.len() + )); + } + unsigned_const(constants[0], what) +} + +/// Construct the fixed two-half SM120 epilogue IR type. +pub(crate) fn tile_type(ctx: &mut Context) -> TypeHandle { + let storage: TypeHandle = MirFP16Type::get(ctx).into(); + let mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(canonical_data_placement()); + let plan = CuteEpiloguePlanAttr::sm120_mxf4_128x128(mma); + CuteEpilogueTileType::get(ctx, storage, plan).into() +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that the fixed epilogue type retains two 16 KiB halves. + #[test] + fn fixed_tile_type_keeps_two_16_kib_halves() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let ty = tile_type(&mut ctx); + let ty = ty.deref(&ctx); + let tile = ty.downcast_ref::().unwrap(); + assert_eq!(tile.half_bytes(), Some(16 * 1024)); + assert_eq!(tile.full_bytes(), Some(32 * 1024)); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/identity.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/identity.rs new file mode 100644 index 0000000000..f7103f956e --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/identity.rs @@ -0,0 +1,578 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Canonical defining paths and the closed source-abstraction boundary. + +use super::registry::{CUTE_FNS, CuteFnSpec}; +#[cfg(test)] +use super::registry::{CuteFn, SubstKind}; +use rustc_public::{CrateDef, ty::FnDef}; + +/// Classification of one callee at the closed CuTe compiler boundary. +pub(super) enum CuteIdentity { + /// Callee is not from the cute_rs crate; normal translation proceeds. + NotCuteCrate, + /// A recognized entry point. + Known(#[doc = "Registered compiler-hook specification."] &'static CuteFnSpec), + /// Ordinary zero-cost Rust whose body is intentionally translated. + SourceAbstraction, + /// In the cute_rs crate but not in the table: hard error (fall-through + /// would body-translate and scalarize the layout). + Unknown(#[doc = "Unrecognized defining path for the rejection diagnostic."] String), +} + +/// Definition modules containing ordinary Rust source abstractions rather +/// than importer-owned recognition stubs. +/// +/// This is deliberately a closed list. Registered [`CUTE_FNS`] entries are +/// matched first, including compiler stubs nested below these modules. +/// Every remaining function must have a real, device-safe body, even when an +/// ergonomic method can be recognized separately. An unregistered panic-only +/// stub could otherwise fall through to body translation. Adding a module +/// requires auditing that fallback and updating the exact-list test. +const CUTE_SOURCE_ABSTRACTION_MODULES: &[&str] = &[ + "arch", + "mma::bf16", + "block_scaled", + "mma::block_scaled", + "copy::cooperative", + "epilogue::sm120", + "numeric", + "pipeline", + "reduction", + "scheduler", + "sm100", + "attention::sm100", + "static_layout", + "tensor", + "copy::tiled", +]; + +/// Match a complete module path, including every namespace separator. +/// +/// Auditing `mma::bf16` never admits `mma`, `mma::other`, or `mma::bf16_extra`. +fn path_is_within_module(relative: &str, module: &str) -> bool { + relative + .strip_prefix(module) + .and_then(|suffix| suffix.strip_prefix("::")) + .is_some_and(|item| !item.is_empty()) +} + +/// Accept only the explicitly listed source-helper module prefixes. +fn is_source_abstraction_path(path: &str, crate_name: &str) -> bool { + let Some(rest) = path + .strip_prefix(crate_name) + .and_then(|rest| rest.strip_prefix("::")) + else { + return false; + }; + CUTE_SOURCE_ABSTRACTION_MODULES + .iter() + .any(|module| path_is_within_module(rest, module)) +} + +/// Source-only modules whose function bodies are also callable through Rust's +/// `Fn` traits. The core `tensor` and `pipeline` modules intentionally remain +/// excluded: their ergonomic methods include compiler-owned stubs alongside +/// ordinary source helpers. The source-only `tensor::partition` child is +/// admitted separately. +const CUTE_SOURCE_CALLBACK_MODULES: &[&str] = &[ + "arch", + "mma::bf16", + "numeric", + "tensor::partition", + "reduction", + "static_layout", +]; + +/// Whether a function item may bypass direct CuTe call recognition. +fn is_source_callback_path(path: &str, crate_name: &str) -> bool { + if !matches!(crate_name, "cute_rs" | "cute-rs") || !is_source_abstraction_path(path, crate_name) + { + return false; + } + let Some(relative) = path + .strip_prefix(crate_name) + .and_then(|path| path.strip_prefix("::")) + else { + return false; + }; + // Registered hooks always take precedence, including any future hook + // added under a currently source-only module. Normalize the crate spelling + // in the same way for both accepted crate identities. + if CUTE_FNS + .iter() + .any(|spec| spec.canonical.strip_prefix("cute_rs::") == Some(relative)) + { + return false; + } + CUTE_SOURCE_CALLBACK_MODULES + .iter() + .any(|module| path_is_within_module(relative, module)) +} + +/// Keep compiler-owned CuTe targets on the direct-call recognition path while +/// permitting audited real Rust helpers to serve as tensor-map callbacks. +pub(in crate::translator) fn function_item_requires_direct_dispatch(fn_def: &FnDef) -> bool { + let crate_name = fn_def.krate().name.to_string(); + matches!(crate_name.as_str(), "cute_rs" | "cute-rs") + && !is_source_callback_path(&canonical_path(fn_def, &crate_name), &crate_name) +} + +/// Recover a callee's actual definition path, independent of public aliases. +pub(super) fn canonical_path(fn_def: &FnDef, crate_name: &str) -> String { + canonical_def_path(fn_def.def_id(), crate_name) +} + +/// Recover the defining identity shared by function and ADT recognition. +/// +/// `DefId::name()` uses rustc's diagnostic printer, which can prefer a public +/// reexport even for individual parents. Combining those display names with +/// real parents can manufacture a path that never existed. Read the underlying +/// `DefPath` instead, as the backend's intrinsic classifier does. Its components +/// preserve impl owners and nested modules, including a module named like the +/// crate, and never substitute a compatibility alias. +pub(super) fn canonical_def_path(def: rustc_public::DefId, crate_name: &str) -> String { + rustc_middle::ty::tls::with(|tcx| { + let def = rustc_public::rustc_internal::internal(tcx, def); + format!( + "{crate_name}{}", + tcx.def_path(def).to_string_no_crate_verbose() + ) + }) +} + +#[cfg(test)] +mod definition_tests; + +/// Classify a callee as a registered hook, approved source helper, foreign call or error. +pub(super) fn classify_cute_call(fn_def: &FnDef) -> CuteIdentity { + let crate_name = fn_def.krate().name.to_string(); + if !matches!(crate_name.as_str(), "cute_rs" | "cute-rs") { + return CuteIdentity::NotCuteCrate; + } + let path = canonical_path(fn_def, &crate_name); + match CUTE_FNS.iter().find(|spec| spec.canonical == path) { + Some(spec) => CuteIdentity::Known(spec), + None if is_source_abstraction_path(&path, &crate_name) => CuteIdentity::SourceAbstraction, + None => CuteIdentity::Unknown(path), + } +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod canonical_path_tests { + use super::super::classify::{ + block_scaled_method_name_from_path, pipeline_method_name_from_path, + scheduler_method_name_from_path, smem_mma_method_name_from_path, + tensor_method_name_from_path, + }; + use super::{is_source_abstraction_path, is_source_callback_path}; + + /// Higher-order tensor operations may call real numeric/helper bodies, + /// but neither registered hooks nor ergonomic compiler stubs may escape + /// the direct-call recognition boundary through a function item. + #[test] + fn function_items_allow_only_audited_source_bodies() { + for path in [ + "cute_rs::numeric::silu_bf16", + "cute_rs::numeric::{impl#3}::to_f32", + "cute-rs::numeric::BFloat16::from_f32", + "cute_rs::arch::exp2_f32", + "cute_rs::tensor::partition::{impl#3}::get", + "cute_rs::static_layout::{impl#1}::crd2idx", + "cute_rs::mma::bf16::{impl#1}::get_slice", + "cute_rs::reduction::{impl#0}::observe", + ] { + let crate_name = path.split_once("::").unwrap().0; + assert!(is_source_callback_path(path, crate_name), "rejected {path}"); + } + for path in [ + "cute_rs::copy::vector::load_tile", + "cute_rs::tensor::make_tensor_read", + "cute_rs::tensor::{impl#6}::local_tile", + "cute_rs::tensor::{impl#7}::load", + "cute_rs::pipeline::{impl#4}::advance", + "cute_rs::sm100::tcgen05_alloc", + "cute_rs::numeric_extra::silu_bf16", + "cute_rs::other::numeric::silu_bf16", + "other::numeric::silu_bf16", + ] { + assert!(!is_source_callback_path(path, "cute_rs"), "accepted {path}"); + } + for spec in super::CUTE_FNS { + assert!( + !is_source_callback_path(spec.canonical, "cute_rs"), + "registered compiler hook escaped direct dispatch: {}", + spec.canonical + ); + } + } + + /// Grouped namespaces retain separate audited source and callback boundaries. + /// Sibling MMA/copy/attention families must not inherit admission from a + /// broadly allowed domain parent. + #[test] + fn grouped_domains_keep_source_and_callback_boundaries_separate() { + for crate_name in ["cute_rs", "cute-rs"] { + for module in ["mma::bf16", "tensor::partition"] { + let path = format!("{crate_name}::{module}::{{impl#3}}::load"); + assert!(is_source_abstraction_path(&path, crate_name), "{path}"); + assert!(is_source_callback_path(&path, crate_name), "{path}"); + for relative in [ + format!("{module}_extra::load"), + format!("other::{module}::load"), + ] { + let path = format!("{crate_name}::{relative}"); + assert!(!is_source_callback_path(&path, crate_name), "{path}"); + } + } + for module in [ + "mma::block_scaled", + "copy::cooperative", + "copy::tiled", + "epilogue::sm120", + "attention::sm100", + ] { + let path = format!("{crate_name}::{module}::{{impl#3}}::operation"); + assert!(is_source_abstraction_path(&path, crate_name), "{path}"); + assert!(!is_source_callback_path(&path, crate_name), "{path}"); + let sibling = format!("{crate_name}::{module}_extra::operation"); + assert!( + !is_source_abstraction_path(&sibling, crate_name), + "{sibling}" + ); + } + for relative in [ + "mma::unregistered_instruction", + "mma::other::operation", + "copy::vector::unregistered_copy", + "copy::other::operation", + "attention::metadata::operation", + "attention::other::operation", + "epilogue::other::operation", + "tensor::{impl#0}::load", + ] { + let path = format!("{crate_name}::{relative}"); + assert!(!is_source_callback_path(&path, crate_name), "{path}"); + } + } + } + + /// Check that cache-aware stores match only their exact compiler entry point. + #[test] + fn cache_aware_global_store_has_an_exact_compiler_boundary() { + let spec = super::CUTE_FNS + .iter() + .find(|spec| spec.canonical == "cute_rs::copy::vector::store_global_no_allocate") + .expect("global store marker"); + assert_eq!(spec.func, super::CuteFn::StoreGlobalNoAllocate); + assert_eq!( + spec.schema, + &[super::SubstKind::Type, super::SubstKind::Const] + ); + } + + /// Check that only known source modules are recognized. + #[test] + fn source_abstraction_modules_are_exact_and_closed() { + assert!(is_source_abstraction_path( + "cute_rs::arch::make_warp_uniform", + "cute_rs" + )); + assert_eq!( + super::CUTE_SOURCE_ABSTRACTION_MODULES, + &[ + "arch", + "mma::bf16", + "block_scaled", + "mma::block_scaled", + "copy::cooperative", + "epilogue::sm120", + "numeric", + "pipeline", + "reduction", + "scheduler", + "sm100", + "attention::sm100", + "static_layout", + "tensor", + "copy::tiled", + ] + ); + assert!(is_source_abstraction_path( + "cute_rs::block_scaled::tensor::BlockScaledTensor<...>::thread_row", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute-rs::numeric::Float4E2M1FNx2::to_f32x2", + "cute-rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::tensor::Tensor<...>::layout", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::copy::tiled::ThrCopy<...>::copy", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::mma::block_scaled::tiled::Mxf4TiledMma<...>::gemm", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::epilogue::sm120::Sm120EpilogueWarp128x128::store_atom", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::pipeline::PipelineTmaAsync<...>::consumer_wait", + "cute_rs" + )); + assert!(is_source_abstraction_path( + "cute_rs::scheduler::StaticPersistentTileScheduler<...>::get_current_work", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "cute_rs::copy::vector::load_tile", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "cute_rs::block_scaled_extra::load", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "other::block_scaled::load", + "cute_rs" + )); + } + + /// Tensor partitions and typed BF16 atoms have real Rust bodies; adjacent + /// modules and same-named modules in other crates must remain rejected. + #[test] + fn typed_partition_and_bf16_modules_have_closed_source_boundaries() { + for (module, method) in [ + ("copy::cooperative", "SharedRow<...>::zip_fold"), + ("static_layout", "Layout<...>::crd2idx"), + ("mma::bf16", "Bf16Mma16x8x16::get_slice"), + ("reduction", "argmax_bf16"), + ] { + for crate_name in ["cute_rs", "cute-rs"] { + assert!(is_source_abstraction_path( + &format!("{crate_name}::{module}::{method}"), + crate_name + )); + assert!(!is_source_abstraction_path( + &format!("{crate_name}::{module}_extra::{method}"), + crate_name + )); + assert!(!is_source_abstraction_path( + &format!("{crate_name}::other::{module}::{method}"), + crate_name + )); + assert!(!is_source_abstraction_path( + &format!("other::{module}::{method}"), + crate_name + )); + } + } + // Hardware/compiler stubs continue to require explicit registration. + assert!(!is_source_abstraction_path( + "cute_rs::mma::unregistered_instruction", + "cute_rs" + )); + assert!(!is_source_abstraction_path( + "cute_rs::copy::vector::unregistered_copy", + "cute_rs" + )); + } + + /// Check that tensor methods accept the supported rustc impl-owner spellings. + #[test] + fn tensor_method_path_accepts_rustc_impl_owner_spellings() { + assert_eq!( + tensor_method_name_from_path("cute_rs::tensor::Contiguous1D>::from_slice"), + Some("from_slice") + ); + assert_eq!( + tensor_method_name_from_path("cute_rs::tensor::{impl#6}::from_slice"), + Some("from_slice") + ); + assert_eq!( + tensor_method_name_from_path("cute-rs::tensor::Tile1D>::store_at"), + Some("store_at") + ); + assert_eq!( + tensor_method_name_from_path("cute_rs::tensor::from_slice"), + None, + "a free function is not an inherent method" + ); + assert_eq!( + tensor_method_name_from_path("cute_rs::other::Tile1D>::base"), + None + ); + } + + /// Check that block-scaled method matching ignores only the impl-owner spelling. + #[test] + fn block_scaled_method_path_ignores_only_the_impl_owner_spelling() { + assert_eq!( + block_scaled_method_name_from_path( + "cute_rs::block_scaled::tensor::KMajor>::from_slices" + ), + Some("from_slices") + ); + assert_eq!( + block_scaled_method_name_from_path( + "cute_rs::block_scaled::gemv::{impl#21}::dot_accumulate" + ), + Some("dot_accumulate") + ); + assert_eq!( + block_scaled_method_name_from_path( + "cute_rs::block_scaled::__compiler::block_scaled_dot_k64" + ), + None, + "compiler hooks are recognized only by the exact helper table" + ); + assert_eq!( + block_scaled_method_name_from_path("cute_rs::other::{impl#1}::load"), + None + ); + } + + /// Check that scheduler method matching ignores only the impl-owner spelling. + #[test] + fn scheduler_method_path_ignores_only_the_impl_owner_spelling() { + assert_eq!( + scheduler_method_name_from_path( + "cute_rs::scheduler::StaticPersistentTileScheduler<...>>::has_work" + ), + Some("has_work") + ); + assert_eq!( + scheduler_method_name_from_path("cute_rs::scheduler::{impl#4}::tile_idx"), + Some("tile_idx") + ); + assert_eq!( + scheduler_method_name_from_path("cute_rs::scheduler::__compiler::scheduler_has_work"), + Some("scheduler_has_work"), + "the exact helper table takes precedence over method classification" + ); + } + + /// Check that pipeline method matching ignores only the impl-owner spelling. + #[test] + fn pipeline_method_path_ignores_only_the_impl_owner_spelling() { + assert_eq!( + pipeline_method_name_from_path( + "cute_rs::pipeline::PipelineTmaAsync<...>>::producer_expect_tx" + ), + Some("producer_expect_tx") + ); + assert_eq!( + pipeline_method_name_from_path("cute_rs::pipeline::{impl#4}::advance"), + Some("advance") + ); + assert_eq!( + pipeline_method_name_from_path("cute_rs::pipeline::__compiler::pipeline_consumer_wait"), + Some("pipeline_consumer_wait"), + "the exact helper table takes precedence over method classification" + ); + } + + /// Check that shared MMA matching requires the exact module and method. + #[test] + fn shared_mma_method_path_keeps_module_and_method_exact() { + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::tiled::Mxf4TiledMma<...>>::load_a_128", + "mma::block_scaled::tiled" + ), + Some("load_a_128") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute-rs::copy::tiled::{impl#7}::from_raw_parts", + "copy::tiled" + ), + Some("from_raw_parts") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::other::Mxf4TiledMma<...>>::load_a_128", + "mma::block_scaled::tiled" + ), + None + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::load_a_128", + "mma::block_scaled::tiled" + ), + None, + "a free function is not an inherent method" + ); + } + + /// Check that the accumulator-tile `gemm` is keyed by the fragments module, + /// so the one-cell `Mxf4TiledMma::gemm` in the tiled module is not taken for it. + #[test] + fn accumulator_gemm_is_keyed_by_the_fragments_module() { + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::fragments::{impl#4}::gemm", + "mma::block_scaled::fragments" + ), + Some("gemm") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::fragments::Mxf4AccumulatorTile2x8::gemm", + "mma::block_scaled::fragments" + ), + Some("gemm") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::tiled::Mxf4TiledMma<...>>::gemm", + "mma::block_scaled::fragments" + ), + None + ); + } + + /// Check that epilogue matching requires each exact owner module. + #[test] + fn epilogue_method_paths_keep_each_owner_module_exact() { + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::epilogue::sm120::Sm120Epilogue128x128::sync_ready_for_tma", + "epilogue::sm120" + ), + Some("sync_ready_for_tma") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::mma::block_scaled::epilogue::{impl#9}::store_tile", + "mma::block_scaled::epilogue" + ), + Some("store_tile") + ); + assert_eq!( + pipeline_method_name_from_path( + "cute_rs::pipeline::PipelineTmaStore<1>::producer_commit" + ), + Some("producer_commit") + ); + assert_eq!( + smem_mma_method_name_from_path( + "cute_rs::epilogue::sm120::Sm120EpilogueWarp128x128::store_tile", + "mma::block_scaled::epilogue" + ), + None, + "store_tile is defined in mma::block_scaled::epilogue, not epilogue::sm120" + ); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/identity/definition_tests.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/identity/definition_tests.rs new file mode 100644 index 0000000000..e98ee1874d --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/identity/definition_tests.rs @@ -0,0 +1,164 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Defining identities must survive rustc's alias-aware display names. + +use super::canonical_def_path; + +#[test] +fn defining_paths_ignore_reexport_names_and_keep_same_named_modules() { + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "cuda_oxide_definition_paths_{}_{}", + std::process::id(), + unique + )); + std::fs::create_dir_all(&root).unwrap(); + let library = root.join("library.rs"); + std::fs::write( + &library, + r#" +pub mod mma { + pub mod bf16 { + pub struct Fragment(pub u32); + #[inline(never)] + pub fn identity(value: Fragment) -> Fragment { value } + } +} +pub use mma::bf16 as bf16_mma; +pub use mma::bf16::Fragment as RootFragment; +pub type FragmentAlias = mma::bf16::Fragment; +pub mod cute_rs { + #[inline(never)] + pub fn nested(value: u32) -> u32 { value } +} +"#, + ) + .unwrap(); + let consumer = root.join("consumer.rs"); + std::fs::write( + &consumer, + r#" +pub fn exercise(value: cute_rs::FragmentAlias) -> cute_rs::RootFragment { + let value = cute_rs::bf16_mma::identity(value); + let _ = cute_rs::cute_rs::nested(value.0); + value +} +"#, + ) + .unwrap(); + + let rustc = std::env::var_os("RUSTC").unwrap_or_else(|| "rustc".into()); + let sysroot_output = std::process::Command::new(&rustc) + .args(["--print", "sysroot"]) + .output() + .expect("query rustc sysroot"); + assert!(sysroot_output.status.success()); + let sysroot = String::from_utf8(sysroot_output.stdout).unwrap(); + let library_output = std::process::Command::new(&rustc) + .args([ + "--edition=2024", + "--crate-type=rlib", + "--crate-name=cute_rs", + ]) + .arg(format!("--out-dir={}", root.display())) + .arg(&library) + .output() + .expect("compile external reexport fixture"); + assert!( + library_output.status.success(), + "{}", + String::from_utf8_lossy(&library_output.stderr) + ); + let args = vec![ + "rustc".to_string(), + "--edition=2024".to_string(), + "--crate-type=rlib".to_string(), + "--crate-name=definition_consumer".to_string(), + "--emit=metadata".to_string(), + "-Zmir-opt-level=0".to_string(), + format!( + "--extern=cute_rs={}", + root.join("libcute_rs.rlib").display() + ), + format!("--out-dir={}", root.display()), + format!("--sysroot={}", sysroot.trim()), + consumer.display().to_string(), + ]; + + let identities = std::thread::Builder::new() + .stack_size(64 * 1024 * 1024) + .spawn(move || { + rustc_public::run!(&args, || { + use rustc_public::CrateDef; + use rustc_public::mir::{Operand, TerminatorKind, mono::Instance}; + use rustc_public::ty::{RigidTy, TyKind}; + + let item = rustc_public::all_local_items() + .into_iter() + .find(|item| item.name().ends_with("::exercise")) + .expect("consumer function"); + let instance = Instance::try_from(item).expect("consumer instance"); + let body = instance.body().expect("consumer body"); + let TyKind::RigidTy(RigidTy::Adt(adt, _)) = body.locals()[1].ty.kind() else { + panic!("aliased parameter retains its ADT identity"); + }; + let mut identities = Vec::new(); + let mut record = |def: rustc_public::DefId, expected: &str| { + identities.push(( + def.name().to_string(), + canonical_def_path(def, "cute_rs"), + expected.to_string(), + )); + }; + record(adt.def_id(), "cute_rs::mma::bf16::Fragment"); + for block in &body.blocks { + let TerminatorKind::Call { + func: Operand::Constant(constant), + .. + } = &block.terminator.kind + else { + continue; + }; + let TyKind::RigidTy(RigidTy::FnDef(function, _)) = constant.const_.ty().kind() + else { + continue; + }; + let def = function.def_id(); + match function.name().rsplit("::").next().unwrap() { + "identity" => { + record(def, "cute_rs::mma::bf16::identity"); + record(def.parent().unwrap(), "cute_rs::mma::bf16"); + } + "nested" => { + record(def, "cute_rs::cute_rs::nested"); + record(def.parent().unwrap(), "cute_rs::cute_rs"); + } + other => panic!("unexpected fixture call: {other}"), + } + } + std::ops::ControlFlow::<(), _>::Continue(identities) + }) + }) + .unwrap() + .join() + .unwrap() + .expect("in-process consumer compilation succeeds"); + + std::fs::remove_dir_all(&root).ok(); + assert_eq!(identities.len(), 5, "ADT, both functions and their modules"); + assert!( + identities + .iter() + .any(|(display, _, _)| display.contains("RootFragment") || display.contains("bf16_mma")), + "fixture must exercise rustc's reexport-aware names: {identities:?}" + ); + for (display, canonical, expected) in identities { + assert_eq!(canonical, expected, "display name: {display}"); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/layout.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/layout.rs new file mode 100644 index 0000000000..faa9b96be2 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/layout.rs @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Monomorphized substs -> layout extraction. +//! +//! `cute-rs` encodes static tile parameters as generics (`load_tile::`); +//! by MIR-import time they are concrete in the callee's substs. This module +//! reads them out STRUCTURALLY (const values only exist as allocation bytes +//! in the Debug rendering, so substring matching cannot see them) and hands +//! them to the op emitters. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::types::translate_type; +use pliron::context::Context; +use pliron::input_err; +use pliron::location::Location; +use pliron::r#type::TypeHandle; +use rustc_public::mir; + +/// Read a monomorphized unsigned const without truncating it. +/// +/// Rust's stable compiler API exposes some constants as raw allocation +/// bytes. Convert through `u128` and then check the destination width; `as` +/// would silently turn an invalid large layout constant into a different one. +pub(super) fn const_u64( + c: &rustc_public::ty::TyConst, + what: &str, + loc: &Location, +) -> TranslationResult { + use rustc_public::ty::TyConstKind; + let raw = match c.kind() { + TyConstKind::Value(_, alloc) => match alloc.read_uint() { + Ok(value) => value, + Err(error) => { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{what} const generic could not be read: {error:?}" + )) + ); + } + }, + _ => match c.eval_target_usize() { + Ok(value) => u128::from(value), + Err(error) => { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "{what} const generic could not be evaluated: {error:?}" + )) + ); + } + }, + }; + match u64::try_from(raw) { + Ok(value) => Ok(value), + Err(_) => input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} const generic {raw} does not fit in u64")) + ), + } +} + +/// Extract `(T, N)` from a call to a cute-rs tile function of shape +/// `fn f(...)`: substs position 0 is the element type, +/// position 1 the tile length. Errors loudly on anything unexpected; a +/// silent default would hide a real compiler gap. +pub(crate) fn tile_elem_and_count( + ctx: &mut Context, + func: &mir::Operand, + loc: &Location, +) -> TranslationResult<(TypeHandle, u64)> { + let mir::Operand::Constant(const_op) = func else { + return input_err!( + loc.clone(), + TranslationErr::unsupported( + "cute-rs tile call through a non-constant callee (indirect call?)".to_string() + ) + ); + }; + let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::FnDef(_, substs)) = + const_op.const_.ty().kind() + else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("cute-rs tile callee is not a FnDef".to_string()) + ); + }; + + let Some(rustc_public::ty::GenericArgKind::Type(elem_rust_ty)) = substs.0.first() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported( + "cute-rs tile call missing element type generic at substs[0]".to_string() + ) + ); + }; + let elem = translate_type(ctx, elem_rust_ty)?; + + let Some(rustc_public::ty::GenericArgKind::Const(c)) = substs.0.get(1) else { + return input_err!( + loc.clone(), + TranslationErr::unsupported( + "cute-rs tile call missing tile-length const generic at substs[1]".to_string() + ) + ); + }; + let n = const_u64(c, "cute-rs tile length", loc)?; + if n == 0 || n > i64::MAX as u64 { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "cute-rs tile length must be in 1..={}, got {n}", + i64::MAX + )) + ); + } + + Ok((elem, n)) +} + +/// Extract `D` from `assume::`, checking the promise is meaningful. +pub(crate) fn assume_divisor(func: &mir::Operand, loc: &Location) -> TranslationResult { + let mir::Operand::Constant(const_op) = func else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("assume callee is not constant".to_string()) + ); + }; + let rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::FnDef(_, substs)) = + const_op.const_.ty().kind() + else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("assume callee is not a FnDef".to_string()) + ); + }; + let Some(rustc_public::ty::GenericArgKind::Const(divisor)) = substs.0.first() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported( + "assume is missing divisor const generic at substs[0]".to_string() + ) + ); + }; + let divisor = const_u64(divisor, "assume divisor", loc)?; + if divisor == 0 { + return input_err!( + loc.clone(), + TranslationErr::unsupported("assume divisor must be greater than zero".to_string()) + ); + } + Ok(divisor) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs new file mode 100644 index 0000000000..5d661a1169 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/mod.rs @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Translate `cute-rs` calls into CuTe operations. +//! +//! CuTe-specific call handling lives here. The shared translator calls this +//! module from `translate_call`, as it does for CUDA intrinsics. +//! +//! Calls are matched by their definition path, found by walking the `DefId` +//! parent chain. This keeps matching stable when a function is re-exported or +//! a trait method is printed with a different name: +//! +//! ```text +//! public name: cute_rs::load_tile +//! definition path: cute_rs::copy::vector::load_tile -> CUTE_FNS entry +//! ``` +//! +//! There are three routes through this boundary: +//! +//! - Registered primitives use [`registry::CUTE_FNS`] and validate generic +//! argument kinds before their family decoder reads static configuration. +//! - Ergonomic methods in approved source modules are recognized by exact +//! defining module and receiver/result ADTs in [`classify`]. Their family +//! decoders validate those types without a registry entry. +//! - Other helpers in approved source modules use ordinary body translation. +//! +//! The first two routes emit CuTe operations without translating the Rust body. +//! Unknown `cute_rs` calls outside approved modules are errors: translating a +//! compiler stub's body would lose the operation's meaning. [`identity`] owns +//! the module allowlist; [`dispatch`] selects the decoder and emitter. + +mod classify; +mod dispatch; +mod identity; +mod registry; + +pub(crate) mod block_scaled; +pub(crate) mod emit; +pub(crate) mod epilogue; +pub(crate) mod layout; +pub(crate) mod pipeline; +pub(crate) mod scheduler; +pub(crate) mod sm100; +pub(crate) mod sm100_attention; +pub(crate) mod smem_mma; +pub(crate) mod static_config; +pub(crate) mod tensor; + +pub(crate) use dispatch::try_translate_cute_call; +pub(super) use identity::function_item_requires_direct_dispatch; +pub use registry::is_cute_compiler_hook_path; diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/emit.rs new file mode 100644 index 0000000000..27c421bd76 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/emit.rs @@ -0,0 +1,929 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the semantic TMA load-pipeline protocol. +//! +//! Only the pipeline handle is compiler-only. Cursor state stays as the same +//! two Rust scalars used by the source loop: +//! +//! ```text +//! make -> init +//! state_new -> index -> acquire/expect or wait/release -> advance -> tail +//! ``` + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::types::translate_type; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::{CutePipelineRoleAttr, CutePipelineStateAttr}; +use dialect_cute::pipeline_ops::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, +}; +use dialect_cute::types::CuteTmaLoadPipelineType; +use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr}; +use dialect_mir::ops::{ + MirCastOp, MirExtractFieldOp, MirFieldAddrOp, MirInsertFieldOp, MirLoadOp, MirStoreOp, + MirUndefOp, +}; +use dialect_mir::types::{MirPtrType, MirStructType, address_space}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::types::IntegerType; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::printable::Printable; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::pipeline::{ + PipelineTmaAsyncRust, decode_pipeline_state, decode_pipeline_state_receiver, + decode_tma_load_pipeline, +}; + +/// Storage bytes occupied by one hardware memory barrier. +const BARRIER_BYTES: u64 = 8; + +/// Insert an operation after the current anchor, or at the start of the block. +fn insert( + ctx: &mut Context, + op: Ptr, + block: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block, ctx); + } +} + +/// Branch from the last emitted effect to the call continuation; reject diverging calls. +fn finish( + ctx: &mut Context, + effect: Ptr, + target: &Option, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target is not supported")) + ); + }; + Ok(emit_goto(ctx, *target, effect, block_map, loc)) +} + +/// Store a result in the unprojected MIR destination, then emit the continuation branch. +#[allow(clippy::too_many_arguments)] +fn finish_value( + ctx: &mut Context, + producer: Ptr, + value: Value, + destination: &mir::Place, + target: &Option, + block: Ptr, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "{what} destination with projections is not supported" + )) + ); + } + let after_store = value_map + .store_local(ctx, destination.local, value, block, Some(producer)) + .unwrap_or(producer); + finish(ctx, after_store, target, block_map, loc, what) +} + +/// Translate the MIR destination type into its IR carrier type. +fn destination_type( + ctx: &mut Context, + body: &mir::Body, + destination: &mir::Place, + what: &str, +) -> TranslationResult { + let ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} destination type: {error:?}" + ))) + })?; + translate_type(ctx, &ty).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot translate {what} destination type: {}", + error.disp(ctx) + ))) + }) +} + +/// Build the concrete cursor carrier with its slot, phase and producer/consumer role. +fn build_state( + ctx: &mut Context, + ty: TypeHandle, + slot: Value, + phase: Value, + block: Ptr, + prev: Ptr, + loc: &Location, +) -> (Value, Ptr) { + let undef = MirUndefOp::new(ctx, ty).get_operation(); + undef.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, undef, block, Some(prev)); + let mut value = undef.deref(ctx).get_result(0); + let mut anchor = undef; + for (index, field) in [slot, phase].into_iter().enumerate() { + let op = Operation::new( + ctx, + MirInsertFieldOp::get_concrete_op_info(), + vec![ty], + vec![value, field], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirInsertFieldOp::new(op).set_attr_insert_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block, Some(anchor)); + value = op.deref(ctx).get_result(0); + anchor = op; + } + (value, anchor) +} + +/// Require a recognized load-pipeline type and recover its static barrier plan. +fn pipeline_facts_from_type( + ty: &mir::Operand, + body: &mir::Body, + what: &str, +) -> TranslationResult { + let ty = ty.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} pipeline type: {error:?}" + ))) + })?; + decode_tma_load_pipeline(&ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what} pipeline: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} needs a by-value PipelineTmaAsync receiver" + ))) + }) +} + +/// Recover a cursor operand's stage count and producer/consumer role. +fn state_facts_from_operand( + operand: &mir::Operand, + body: &mir::Body, + what: &str, +) -> TranslationResult { + let ty = operand.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} state type: {error:?}" + ))) + })?; + decode_pipeline_state_receiver(&ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what} state: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} needs a PipelineState receiver" + ))) + }) +} + +/// Require the cursor to match the selected pipeline and expected participant role. +fn require_state( + state: CutePipelineStateAttr, + stages: u64, + role: Option, + what: &str, +) -> TranslationResult<()> { + if state.stages != stages { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!( + "{what} state has {} stages, but its pipeline has {stages}", + state.stages + ) + ))); + } + if let Some(role) = role + && state.role != role + { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} needs {role:?} state, got {:?}", state.role) + ))); + } + Ok(()) +} + +/// Recover a shared barrier pointer at the validated unsafe pipeline boundary. +fn recover_shared_barrier_base( + ctx: &mut Context, + base: Value, + block: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Option>)> { + let pointer = { + let ty = base.get_type(ctx); + ty.deref(ctx).downcast_ref::().cloned() + } + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "TMA load-pipeline base must be a pointer".to_owned() + )) + })?; + if !pointer.is_mutable { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + "TMA load-pipeline base must be mutable".to_owned() + ))); + } + let valid_barrier = { + let pointee = pointer.pointee.deref(ctx); + pointee + .downcast_ref::() + .is_some_and(|barrier| barrier.width() == 64 && barrier.is_unsigned()) + }; + if !valid_barrier { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + "TMA load-pipeline base must point to the canonical u64 Barrier storage".to_owned() + ))); + } + if pointer.address_space == address_space::SHARED { + return Ok((base, prev)); + } + if pointer.address_space != address_space::GENERIC { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!( + "TMA load-pipeline base must be generic or CTA-shared memory, found address space {}", + pointer.address_space + ) + ))); + } + + let shared: TypeHandle = + MirPtrType::get(ctx, pointer.pointee, true, address_space::SHARED).into(); + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![shared], + vec![base], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + insert(ctx, cast, block, prev); + Ok((cast.deref(ctx).get_result(0), Some(cast))) +} + +/// Emit the address of one scalar field in a referenced pipeline-state carrier. +fn field_addr( + ctx: &mut Context, + pointer: Value, + index: usize, + block: Ptr, + prev: Option>, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, TypeHandle, Ptr)> { + let field = { + let ty = pointer.get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} must be a pointer to PipelineState") + ))); + }; + let pointee = pointer.pointee.deref(ctx); + let Some(state) = pointee.downcast_ref::() else { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} must point to PipelineState") + ))); + }; + let Some(field) = state.field_types.get(index).copied() else { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} PipelineState has no field {index}") + ))); + }; + (field, pointer.is_mutable, pointer.address_space) + }; + let field_pointer: TypeHandle = MirPtrType::get(ctx, field.0, field.1, field.2).into(); + // Go through the one construction door so `aggregate_ty` is stamped; + // the verifier rejects a `mir.field_addr` without it. + let op = MirFieldAddrOp::build(ctx, pointer, field_pointer, index as u32)?; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, prev); + Ok((op.deref(ctx).get_result(0), field.0, op)) +} + +/// Emit a scalar load and return both the loaded value and its insertion anchor. +fn load( + ctx: &mut Context, + pointer: Value, + ty: TypeHandle, + block: Ptr, + prev: Ptr, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![ty], + vec![pointer], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirLoadOp::new(op).set_volatile(ctx, false); + insert(ctx, op, block, Some(prev)); + (op.deref(ctx).get_result(0), op) +} + +/// Cursor scalar values and optional addresses needed to update a borrowed cursor. +struct StateValues { + /// Current ring slot as an IR scalar value. + slot: Value, + /// Current barrier parity as an IR scalar value. + phase: Value, + /// Writable slot-field address when the cursor was passed by reference. + slot_pointer: Option, + /// Writable phase-field address when the cursor was passed by reference. + phase_pointer: Option, + /// Last operation emitted while recovering the cursor values. + anchor: Ptr, +} + +/// Read slot and phase from a cursor value or reference, retaining writable field addresses. +fn state_values( + ctx: &mut Context, + value: Value, + block: Ptr, + prev: Option>, + loc: &Location, + what: &str, +) -> TranslationResult { + let ty = value.get_type(ctx); + if ty.deref(ctx).downcast_ref::().is_some() { + let (slot_pointer, slot_ty, slot_addr) = field_addr(ctx, value, 0, block, prev, loc, what)?; + let (slot, slot_load) = load(ctx, slot_pointer, slot_ty, block, slot_addr, loc); + let (phase_pointer, phase_ty, phase_addr) = + field_addr(ctx, value, 1, block, Some(slot_load), loc, what)?; + let (phase, phase_load) = load(ctx, phase_pointer, phase_ty, block, phase_addr, loc); + return Ok(StateValues { + slot, + phase, + slot_pointer: Some(slot_pointer), + phase_pointer: Some(phase_pointer), + anchor: phase_load, + }); + } + + let fields = { + let ty = ty.deref(ctx); + let Some(state) = ty.downcast_ref::() else { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} must be PipelineState or a reference to it") + ))); + }; + if state.field_types.len() < 2 { + return Err(pliron::input_error_noloc!(TranslationErr::unsupported( + format!("{what} PipelineState must contain slot and phase") + ))); + } + [state.field_types[0], state.field_types[1]] + }; + let mut values = Vec::with_capacity(2); + let mut anchor = prev; + for (index, field) in fields.into_iter().enumerate() { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![field], + vec![value], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block, anchor); + values.push(op.deref(ctx).get_result(0)); + anchor = Some(op); + } + Ok(StateValues { + slot: values[0], + phase: values[1], + slot_pointer: None, + phase_pointer: None, + anchor: anchor.expect("two state extracts always produce an anchor"), + }) +} + +/// Construct a load pipeline from its shared barrier storage and static configuration. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_make( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 || !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "TMA load-pipeline construction expects one base and a local destination" + .to_owned() + ) + ); + } + let rust_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read TMA load-pipeline result type: {error:?}" + ))) + })?; + let facts = decode_tma_load_pipeline(&rust_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid TMA load-pipeline result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "TMA load-pipeline construction must return PipelineTmaAsync".to_owned() + )) + })?; + let (base, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (base, after) = recover_shared_barrier_base(ctx, base, block, after.or(prev), &loc)?; + let make = CuteTmaLoadPipelineMakeOp::new( + ctx, + base, + facts.stages, + facts.consumer_warps, + facts.transaction_bytes, + BARRIER_BYTES, + ); + let op = make.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + value_map.set_direct_value(destination.local, make.pipeline(ctx)); + finish(ctx, op, target, block_map, loc, "tma_load_pipeline_make") +} + +/// Emit barrier initialization with the pipeline's arrival and transaction counts. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_init( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "TMA load-pipeline init expects pipeline and thread".to_owned() + ) + ); + } + let _facts = pipeline_facts_from_type(&args[0], body, "pipeline init")?; + let (pipeline, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (thread, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + after.or(prev), + loc.clone(), + )?; + let init = CuteTmaLoadPipelineInitOp::new(ctx, pipeline, thread); + let op = init.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + finish(ctx, op, target, block_map, loc, "tma_load_pipeline_init") +} + +/// Construct a producer or consumer cursor with its initial slot and phase. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_state_new( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if !args.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "pipeline state construction expects no arguments".to_owned() + ) + ); + } + let rust_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read pipeline state result type: {error:?}" + ))) + })?; + let state = decode_pipeline_state(&rust_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid pipeline state result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "pipeline state construction must return PipelineState".to_owned() + )) + })?; + let operation = CutePipelineStateNewOp::new(ctx, state); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, prev); + let ty = destination_type(ctx, body, destination, "pipeline state")?; + let (value, anchor) = build_state( + ctx, + ty, + operation.slot(ctx), + operation.phase(ctx), + block, + op, + &loc, + ); + finish_value( + ctx, + anchor, + value, + destination, + target, + block, + value_map, + block_map, + loc, + "pipeline_state_new", + ) +} + +/// Return the current cursor slot while retaining the semantic state relation. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_state_slot( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("pipeline state slot expects one receiver".to_owned()) + ); + } + let state = state_facts_from_operand(&args[0], body, "pipeline state slot")?; + let (value, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let values = state_values( + ctx, + value, + block, + after.or(prev), + &loc, + "pipeline state slot", + )?; + let operation = CutePipelineStateSlotOp::new(ctx, values.slot, state); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(values.anchor)); + finish_value( + ctx, + op, + operation.index(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "pipeline_state_slot", + ) +} + +/// Advance the cursor slot and toggle parity on wrap, storing referenced state updates. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_state_advance( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("pipeline state advance expects one receiver".to_owned()) + ); + } + let state = state_facts_from_operand(&args[0], body, "pipeline state advance")?; + let (value, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let values = state_values( + ctx, + value, + block, + after.or(prev), + &loc, + "pipeline state advance", + )?; + let Some(slot_pointer) = values.slot_pointer else { + return input_err!( + loc, + TranslationErr::unsupported( + "pipeline state advance receiver must be mutable".to_owned() + ) + ); + }; + let Some(phase_pointer) = values.phase_pointer else { + return input_err!( + loc, + TranslationErr::unsupported( + "pipeline state advance receiver must be mutable".to_owned() + ) + ); + }; + let operation = CutePipelineStateAdvanceOp::new(ctx, values.slot, values.phase, state); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(values.anchor)); + let slot_store = Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![slot_pointer, operation.next_slot(ctx)], + vec![], + 0, + ); + slot_store.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, slot_store, block, Some(op)); + let phase_store = Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![phase_pointer, operation.next_phase(ctx)], + vec![], + 0, + ); + phase_store.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, phase_store, block, Some(slot_store)); + finish( + ctx, + phase_store, + target, + block_map, + loc, + "pipeline_state_advance", + ) +} + +/// Load-pipeline operation and participant role selected for emission. +enum Lifecycle { + /// Wait until an input slot is available for the producer. + ProducerAcquire, + /// Set the expected input-copy byte count for the current slot. + ProducerExpectTx, + /// Wait for the current input slot to become full. + ConsumerWait, + /// Release the current input slot after consumers finish. + ConsumerRelease, + /// Wait for remaining consumers before releasing the input ring. + ProducerTail, +} + +/// Emit one role-checked load-pipeline handoff and its required state operands. +#[allow(clippy::too_many_arguments)] +fn emit_lifecycle( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: Option<&mir::Place>, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + lifecycle: Lifecycle, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "pipeline lifecycle call expects pipeline and state".to_owned() + ) + ); + } + let facts = pipeline_facts_from_type(&args[0], body, "pipeline lifecycle")?; + let state = state_facts_from_operand(&args[1], body, "pipeline lifecycle")?; + let required_role = match lifecycle { + Lifecycle::ConsumerWait | Lifecycle::ConsumerRelease => CutePipelineRoleAttr::Consumer, + Lifecycle::ProducerAcquire | Lifecycle::ProducerExpectTx | Lifecycle::ProducerTail => { + CutePipelineRoleAttr::Producer + } + }; + require_state( + state, + facts.stages, + Some(required_role), + "pipeline lifecycle", + )?; + let (pipeline, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let pipeline_type = pipeline.get_type(ctx); + let matches_pipeline = pipeline_type + .deref(ctx) + .downcast_ref::() + .is_some_and(|ty| { + ty.stages == facts.stages + && ty.consumer_warps == facts.consumer_warps + && ty.transaction_bytes == facts.transaction_bytes + }); + if !matches_pipeline { + return input_err!( + loc, + TranslationErr::unsupported( + "pipeline SSA value does not match its Rust type".to_owned() + ) + ); + } + let (state_value, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + after.or(prev), + loc.clone(), + )?; + let values = state_values( + ctx, + state_value, + block, + after.or(prev), + &loc, + "pipeline lifecycle", + )?; + + let (op, result, what) = match lifecycle { + Lifecycle::ProducerAcquire => { + let operation = + CutePipelineProducerAcquireOp::new(ctx, pipeline, values.slot, values.phase, state); + (operation.get_operation(), None, "pipeline_producer_acquire") + } + Lifecycle::ProducerExpectTx => { + let operation = CutePipelineProducerExpectTxOp::new(ctx, pipeline, values.slot, state); + ( + operation.get_operation(), + Some(operation.completion_barrier(ctx)), + "pipeline_producer_expect_tx", + ) + } + Lifecycle::ConsumerWait => { + let operation = + CutePipelineConsumerWaitOp::new(ctx, pipeline, values.slot, values.phase, state); + (operation.get_operation(), None, "pipeline_consumer_wait") + } + Lifecycle::ConsumerRelease => { + let operation = CutePipelineConsumerReleaseOp::new(ctx, pipeline, values.slot, state); + (operation.get_operation(), None, "pipeline_consumer_release") + } + Lifecycle::ProducerTail => { + let operation = + CutePipelineProducerTailOp::new(ctx, pipeline, values.slot, values.phase, state); + (operation.get_operation(), None, "pipeline_producer_tail") + } + }; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(values.anchor)); + if let Some(result) = result { + let destination = destination.expect("expect-tx is the only value lifecycle operation"); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + what, + ) + } else { + finish(ctx, op, target, block_map, loc, what) + } +} + +/// Define the thin call adapters for the shared load-pipeline lifecycle emitter. +macro_rules! lifecycle_emitter { + ($name:ident, $kind:expr) => { + #[doc = concat!("Import the `", stringify!($kind), "` load-pipeline transition.")] + #[allow(clippy::too_many_arguments)] + pub(crate) fn $name( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + ) -> TranslationResult> { + emit_lifecycle( + ctx, body, args, None, target, block, prev, value_map, block_map, loc, $kind, + ) + } + }; +} + +lifecycle_emitter!(emit_producer_acquire, Lifecycle::ProducerAcquire); +lifecycle_emitter!(emit_consumer_wait, Lifecycle::ConsumerWait); +lifecycle_emitter!(emit_consumer_release, Lifecycle::ConsumerRelease); +lifecycle_emitter!(emit_producer_tail, Lifecycle::ProducerTail); + +/// Import the producer's transaction-byte expectation for the current stage. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_producer_expect_tx( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + emit_lifecycle( + ctx, + body, + args, + Some(destination), + target, + block, + prev, + value_map, + block_map, + loc, + Lifecycle::ProducerExpectTx, + ) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs new file mode 100644 index 0000000000..7af4072e42 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/pipeline/mod.rs @@ -0,0 +1,222 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural Rust types for the TMA load-pipeline boundary. +//! +//! The pipeline handle is compiler-only. Its cursor remains an ordinary Rust +//! pair, so loop state still travels as two `u32` values: +//! +//! ```text +//! PipelineTmaAsync -> !cute.tma_load_pipeline +//! PipelineState -> { slot: u32, phase: u32 } +//! ``` + +pub(crate) mod emit; + +use crate::error::{TranslationErr, TranslationResult}; +use dialect_cute::attributes::CutePipelineStateAttr; +use dialect_cute::types::CuteTmaLoadPipelineType; +use pliron::context::Context; +use pliron::input_error_noloc; +use pliron::r#type::TypeHandle; +use rustc_public::ty::{GenericArgs, RigidTy, Ty, TyKind}; + +use super::tensor::{canonical_adt_path, const_arg, exact_schema, type_arg, unsigned_const}; + +/// Exact defining path of `cute_rs::pipeline::PipelineTmaAsync`; public reexports do not change it. +const TMA_LOAD_PIPELINE: &str = "cute_rs::pipeline::PipelineTmaAsync"; +/// Exact defining path of `cute_rs::pipeline::PipelineState`; public reexports do not change it. +const PIPELINE_STATE: &str = "cute_rs::pipeline::PipelineState"; +/// Exact defining path of `cute_rs::pipeline::Producer`; public reexports do not change it. +const PRODUCER: &str = "cute_rs::pipeline::Producer"; +/// Exact defining path of `cute_rs::pipeline::Consumer`; public reexports do not change it. +const CONSUMER: &str = "cute_rs::pipeline::Consumer"; + +/// Static facts carried by `PipelineTmaAsync`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PipelineTmaAsyncRust { + /// Number of reusable slots in the pipeline ring. + pub stages: u64, + /// Consumer-warp count used to determine release arrivals. + pub consumer_warps: u32, + /// Expected TMA transaction bytes per full stage. + pub transaction_bytes: u32, +} + +/// Recognize one exact source ADT rather than a type with a similar printed name. +fn exact_adt(ty: &Ty, expected: &str) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + if canonical_adt_path(&def) != expected { + return Ok(None); + } + Ok(Some(args)) +} + +/// Decode the compiler-only load-pipeline handle. +pub(crate) fn decode_tma_load_pipeline(ty: &Ty) -> Result, String> { + let Some(args) = exact_adt(ty, TMA_LOAD_PIPELINE)? else { + return Ok(None); + }; + exact_schema(&args, &["Const", "Const", "Const"], TMA_LOAD_PIPELINE)?; + let stages = unsigned_const( + &const_arg(&args, 0, TMA_LOAD_PIPELINE)?, + "pipeline stage count", + )?; + let consumer_warps = unsigned_const( + &const_arg(&args, 1, TMA_LOAD_PIPELINE)?, + "pipeline consumer-warp count", + )?; + let transaction_bytes = unsigned_const( + &const_arg(&args, 2, TMA_LOAD_PIPELINE)?, + "pipeline transaction byte count", + )?; + let consumer_warps = u32::try_from(consumer_warps) + .map_err(|_| "pipeline consumer-warp count does not fit in u32".to_owned())?; + let transaction_bytes = u32::try_from(transaction_bytes) + .map_err(|_| "pipeline transaction byte count does not fit in u32".to_owned())?; + validate_pipeline_facts(stages, consumer_warps, transaction_bytes)?; + Ok(Some(PipelineTmaAsyncRust { + stages, + consumer_warps, + transaction_bytes, + })) +} + +/// Reject unsupported stage, consumer-warp or transaction-byte configurations. +fn validate_pipeline_facts( + stages: u64, + consumer_warps: u32, + transaction_bytes: u32, +) -> Result<(), String> { + if stages == 0 || stages > u64::from(u32::MAX) { + return Err(format!( + "pipeline stage count must be between 1 and {}, got {stages}", + u32::MAX + )); + } + if !(1..=32).contains(&consumer_warps) { + return Err(format!( + "pipeline consumer-warp count must be between 1 and 32, got {consumer_warps}" + )); + } + if transaction_bytes == 0 { + return Err("pipeline transaction byte count must be greater than zero".to_owned()); + } + stages + .checked_mul(16) + .ok_or_else(|| "pipeline barrier storage size overflowed u64".to_owned())?; + Ok(()) +} + +/// Decode the exact role marker; reject unsupported or look-alike role types. +fn decode_role(ty: &Ty) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + let path = canonical_adt_path(&def); + if !matches!(path.as_str(), PRODUCER | CONSUMER) { + return Ok(None); + } + exact_schema(&args, &[], &path)?; + Ok(Some(path == PRODUCER)) +} + +/// Decode the ordinary two-scalar pipeline cursor. +pub(crate) fn decode_pipeline_state(ty: &Ty) -> Result, String> { + let Some(args) = exact_adt(ty, PIPELINE_STATE)? else { + return Ok(None); + }; + exact_schema(&args, &["Type", "Const"], PIPELINE_STATE)?; + let role = type_arg(&args, 0, PIPELINE_STATE)?; + let producer = decode_role(&role)? + .ok_or_else(|| "PipelineState role must be cute-rs Producer or Consumer".to_owned())?; + let stages = unsigned_const( + &const_arg(&args, 1, PIPELINE_STATE)?, + "pipeline-state stage count", + )?; + if stages == 0 || stages > u64::from(u32::MAX) { + return Err(format!( + "pipeline-state stage count must be between 1 and {}, got {stages}", + u32::MAX + )); + } + Ok(Some(if producer { + CutePipelineStateAttr::producer(stages) + } else { + CutePipelineStateAttr::consumer(stages) + })) +} + +/// Decode state behind exactly one reference or raw pointer. +pub(crate) fn decode_pipeline_state_receiver( + ty: &Ty, +) -> Result, String> { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode_pipeline_state(&pointee) + } + _ => decode_pipeline_state(ty), + } +} + +/// Construct the IR pipeline type with every validated static barrier parameter. +fn pipeline_type(ctx: &mut Context, facts: PipelineTmaAsyncRust) -> TypeHandle { + CuteTmaLoadPipelineType::get( + ctx, + facts.stages, + facts.consumer_warps, + facts.transaction_bytes, + ) + .into() +} + +/// Type-translation hook called before generic ADT lowering. +pub(crate) fn try_translate_tma_load_pipeline_type( + ctx: &mut Context, + rust_ty: &Ty, +) -> Option> { + match decode_tma_load_pipeline(rust_ty) { + Ok(Some(facts)) => Some(Ok(pipeline_type(ctx, facts))), + Ok(None) => None, + Err(error) => Some(Err(input_error_noloc!(TranslationErr::unsupported( + format!("invalid cute-rs TMA load-pipeline type: {error}") + )))), + } +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that load-pipeline IR types retain stage count, arrival count, and transaction size. + #[test] + fn load_pipeline_type_keeps_all_static_facts() { + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + let facts = PipelineTmaAsyncRust { + stages: 3, + consumer_warps: 8, + transaction_bytes: 17_408, + }; + let ty = pipeline_type(&mut ctx, facts); + let ty = ty.deref(&ctx); + let pipeline = ty.downcast_ref::().unwrap(); + assert_eq!(pipeline.stages, 3); + assert_eq!(pipeline.consumer_warps, 8); + assert_eq!(pipeline.transaction_bytes, 17_408); + assert_eq!(pipeline.storage_bytes(), Some(48)); + } + + /// Check that unsupported load-pipeline configurations are rejected. + #[test] + fn load_pipeline_facts_fail_closed() { + assert!(validate_pipeline_facts(0, 8, 1).is_err()); + assert!(validate_pipeline_facts(3, 33, 1).is_err()); + assert!(validate_pipeline_facts(3, 8, 0).is_err()); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/registry.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/registry.rs new file mode 100644 index 0000000000..8d8ac92ffb --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/registry.rs @@ -0,0 +1,726 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Exact compiler entry points and their generic-argument schemas. + +use super::{sm100, sm100_attention}; + +/// Recognized cute-rs entry points. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CuteFn { + /// Dispatch a paired SM100 GEMM or pipeline hook. + Sm100(#[doc = "Handler within the paired SM100 operation family."] sm100::Sm100Fn), + /// Dispatch an independent SM100 attention MMA or TMA hook. + Attention( + #[doc = "Handler within the independent attention operation family."] + sm100_attention::AttentionFn, + ), + /// Construct a read-only contiguous tensor over a slice. + MakeTensorRead, + /// Construct a writable tensor over disjoint output storage. + MakeTensorWrite, + /// Split a read-only tensor into fixed-size tiles. + ZippedDivideRead, + /// Split a writable tensor into fixed-size tiles. + ZippedDivideWrite, + /// Select one read-only tensor tile. + SliceRead, + /// Select one writable tensor tile. + SliceWrite, + /// Split a read-only tensor into tiles and select one: CuTe's `local_tile`. + LocalTileRead, + /// Split a writable tensor into tiles and select one: CuTe's `local_tile`. + LocalTileWrite, + /// Test whether every element of the selected tile is in bounds. + TensorTileIsFull, + /// Return the selected tile's absolute starting element index. + TensorTileBase, + /// Load a complete tensor tile into registers. + TensorLoadTile, + /// Store a register tile through a writable tensor view. + TensorStoreTile, + /// Store one element using an absolute storage index. + TensorStoreElementAbs, + /// Pair packed FP4 data with scales and a K-major layout. + BlockScaledMake, + /// Select the row or vector batch owned by one thread. + BlockScaledThreadRow, + /// Select a K64 block from a scaled row. + BlockScaledKTile, + /// Load packed values and scales for the selected K64 block. + BlockScaledLoadK64, + /// Accumulate one scaled K64 dot product in FP32. + BlockScaledDotK64, + /// Initialize persistent work from a linear physical CTA index. + SchedulerNew1d, + /// Test whether the current logical tile remains inside the fixed grid. + SchedulerHasWork, + /// Return the scheduler's current typed work tile. + SchedulerCurrentTile, + /// Decode a work tile's M, N and batch coordinates. + WorkTileCoordinates, + /// Advance persistent work by the physical grid stride. + SchedulerAdvance, + /// Construct a load pipeline with checked shared barrier storage. + TmaLoadPipelineMake, + /// Reject the obsolete untracked raw-pointer pipeline constructor. + TmaLoadPipelineFromRawPartsRejected, + /// Initialize transaction-counted full and reusable empty barriers. + TmaLoadPipelineInit, + /// Create the initial producer or consumer cursor. + PipelineStateNew, + /// Read the current cursor slot. + PipelineStateSlot, + /// Advance a cursor and toggle its parity on wrap. + PipelineStateAdvance, + /// Wait until an input slot is available for the producer. + PipelineProducerAcquire, + /// Set the expected input-copy byte count for the current slot. + PipelineProducerExpectTx, + /// Wait for the current input slot to become full. + PipelineConsumerWait, + /// Release the current input slot after consumers finish. + PipelineConsumerRelease, + /// Wait for remaining consumers before releasing the input ring. + PipelineProducerTail, + /// Overlay a typed data or scale layout on shared storage. + SharedTensorOverlay, + /// Construct the lane-owned tiled-MMA slice. + TiledMmaSlice, + /// Initialize the fixed FP32 accumulator fragment. + MmaFragmentFill, + /// Load A/B scales from the canonical scale atom. + MmaLoadScales, + /// Select scales for one K64 half. + MmaFragmentSliceK, + /// Load the A register fragment for a warp and K64 half. + MmaLoadA, + /// Select the shared B tile for a warp and K64 half. + MmaPartitionB, + /// Accumulate a tiled MXF4 product using paired data and scale operands. + TiledGemm, + /// Overlay the fixed output layout on shared storage. + EpilogueOverlay, + /// Select the warp/lane-owned portion of the epilogue. + EpilogueWarpSlice, + /// Convert and scatter an accumulator fragment into shared FP16 output. + EpilogueStoreFragment, + /// Synchronize participating warps before output-buffer reuse. + EpilogueSyncReusable, + /// Publish shared output writes and synchronize before TMA stores. + EpilogueSyncReadyForTma, + /// Select one contiguous output half for a TMA store. + EpilogueHalf, + /// Wait for the oldest output-store group's shared reads to finish. + TmaStoreAcquire, + /// Commit the current issuer's pending output stores. + TmaStoreCommit, + /// Drain committed output-store groups. + TmaStoreTail, + /// Load one aligned scalar-vector tile into registers. + LoadTile, + /// Store one aligned register tile to memory. + StoreTile, + /// Emit the cache-aware register-to-global store. + StoreGlobalNoAllocate, + /// Record a source-promised divisor for a scalar index. + AssumeDiv, + /// Issue a cooperative global-to-shared tile copy with predicated edges. + CopyG2S, + /// Load the A matrix fragment from its checked shared layout. + LoadMatrixA, + /// Load the B matrix fragment from its checked shared layout. + LoadMatrixB, + /// Load a paired B fragment for adjacent N positions. + LoadMatrixBPair, + /// Load paired B fragments from K-major shared storage. + LoadMatrixBPairKMajor, + /// Issue a typed two-dimensional global-to-shared TMA copy. + CopyTma2d, + /// Issue a typed two-dimensional shared-to-global TMA store. + CopyTmaS2g2d, +} + +/// Expected kind of one generic argument in a recognized function's substs. +#[allow(dead_code)] // `Lifetime` makes future early-bound lifetimes explicit. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SubstKind { + /// One Rust type argument in the hook's generic schema. + Type, + /// One compile-time constant argument in the hook's generic schema. + Const, + /// One lifetime argument retained in the hook's generic schema. + Lifetime, +} + +/// Report whether `path` is the defining path of a recognized cute-rs hook. +/// +/// `path` is the crate-relative form rustc prints for a definition +/// (`::copy::tma::copy_tma_2d`); the registry stores the same path with the +/// `cute_rs` crate prefix. Re-export paths never match, so a caller that +/// resolved an item through the root export `cute_rs::copy_tma_2d` must pass +/// the defining path, not the import path. +pub fn is_cute_compiler_hook_path(path: &str) -> bool { + CUTE_FNS + .iter() + .any(|spec| spec.canonical.strip_prefix("cute_rs") == Some(path)) +} + +/// A recognized function's definition path and expected generic arguments. +/// Re-export paths do not appear here. +pub(crate) struct CuteFnSpec { + /// Exact defining function path used for compiler-hook recognition. + pub canonical: &'static str, + /// Semantic handler selected when the canonical path matches. + pub func: CuteFn, + /// Ordered generic-argument kinds required before decoding hook configuration. + pub schema: &'static [SubstKind], +} + +/// One entry per recognized function, including its generic argument schema. +pub(crate) const CUTE_FNS: &[CuteFnSpec] = &[ + CuteFnSpec { + canonical: "cute_rs::attention::sm100::__compiler::mma_ss", + func: CuteFn::Attention(sm100_attention::AttentionFn::MmaSs), + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + ], + }, + CuteFnSpec { + canonical: "cute_rs::attention::sm100::__compiler::mma_ts", + func: CuteFn::Attention(sm100_attention::AttentionFn::MmaTs), + schema: &[ + SubstKind::Type, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + ], + }, + CuteFnSpec { + canonical: "cute_rs::attention::sm100::__compiler::tma_load", + func: CuteFn::Attention(sm100_attention::AttentionFn::TmaLoad), + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::attention::sm100::__compiler::tma_store", + func: CuteFn::Attention(sm100_attention::AttentionFn::TmaStore), + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tma_store", + func: CuteFn::Sm100(sm100::Sm100Fn::TmaStore), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_commit", + func: CuteFn::Sm100(sm100::Sm100Fn::StoreCommit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_acquire", + func: CuteFn::Sm100(sm100::Sm100Fn::StoreAcquire), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_store_tail", + func: CuteFn::Sm100(sm100::Sm100Fn::StoreTail), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tmem_alloc", + func: CuteFn::Sm100(sm100::Sm100Fn::TmemAlloc), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tmem_dealloc", + func: CuteFn::Sm100(sm100::Sm100Fn::TmemDealloc), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tiled_mma", + func: CuteFn::Sm100(sm100::Sm100Fn::TiledMma), + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Const, + SubstKind::Const, + SubstKind::Const, + ], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_tmem_epilogue", + func: CuteFn::Sm100(sm100::Sm100Fn::TmemEpilogue), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_cluster_tma_load", + func: CuteFn::Sm100(sm100::Sm100Fn::ClusterTmaLoad), + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_init", + func: CuteFn::Sm100(sm100::Sm100Fn::PipelineInit), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_acquire", + func: CuteFn::Sm100(sm100::Sm100Fn::PipelineAcquire), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_expect", + func: CuteFn::Sm100(sm100::Sm100Fn::PipelineExpect), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_wait", + func: CuteFn::Sm100(sm100::Sm100Fn::PipelineWait), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_pipeline_release", + func: CuteFn::Sm100(sm100::Sm100Fn::PipelineRelease), + schema: &[SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_init", + func: CuteFn::Sm100(sm100::Sm100Fn::AccumulatorInit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_acquire", + func: CuteFn::Sm100(sm100::Sm100Fn::AccumulatorAcquire), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_commit", + func: CuteFn::Sm100(sm100::Sm100Fn::AccumulatorCommit), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_wait", + func: CuteFn::Sm100(sm100::Sm100Fn::AccumulatorWait), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::sm100::__compiler::sm100_accumulator_release", + func: CuteFn::Sm100(sm100::Sm100Fn::AccumulatorRelease), + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::make_tensor_read", + func: CuteFn::MakeTensorRead, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::make_tensor_write", + func: CuteFn::MakeTensorWrite, + schema: &[SubstKind::Lifetime, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::zipped_divide_read", + func: CuteFn::ZippedDivideRead, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::zipped_divide_write", + func: CuteFn::ZippedDivideWrite, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::slice_read", + func: CuteFn::SliceRead, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::slice_write", + func: CuteFn::SliceWrite, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::tensor_tile_is_full", + func: CuteFn::TensorTileIsFull, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::tensor_tile_base", + func: CuteFn::TensorTileBase, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::tensor_load_tile", + func: CuteFn::TensorLoadTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::tensor_store_tile", + func: CuteFn::TensorStoreTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::tensor::tensor_store_element_abs", + func: CuteFn::TensorStoreElementAbs, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::block_scaled_make", + func: CuteFn::BlockScaledMake, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::block_scaled_thread_row", + func: CuteFn::BlockScaledThreadRow, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::block_scaled_k_tile", + func: CuteFn::BlockScaledKTile, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::block_scaled_load_k64", + func: CuteFn::BlockScaledLoadK64, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::block_scaled_dot_k64", + func: CuteFn::BlockScaledDotK64, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::nvfp4_make", + func: CuteFn::BlockScaledMake, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::nvfp4_thread_row", + func: CuteFn::BlockScaledThreadRow, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::nvfp4_k_tile", + func: CuteFn::BlockScaledKTile, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::nvfp4_load_k64", + func: CuteFn::BlockScaledLoadK64, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::block_scaled::__compiler::nvfp4_dot_k64", + func: CuteFn::BlockScaledDotK64, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::scheduler::__compiler::scheduler_new_1d", + func: CuteFn::SchedulerNew1d, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::scheduler::__compiler::scheduler_has_work", + func: CuteFn::SchedulerHasWork, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::scheduler::__compiler::scheduler_current_tile", + func: CuteFn::SchedulerCurrentTile, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::scheduler::__compiler::work_tile_coordinates", + func: CuteFn::WorkTileCoordinates, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::scheduler::__compiler::scheduler_advance", + func: CuteFn::SchedulerAdvance, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::tma_load_pipeline_from_raw_base", + func: CuteFn::TmaLoadPipelineMake, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::tma_load_pipeline_init", + func: CuteFn::TmaLoadPipelineInit, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_state_new_producer", + func: CuteFn::PipelineStateNew, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_state_new_consumer", + func: CuteFn::PipelineStateNew, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_state_slot", + func: CuteFn::PipelineStateSlot, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_state_advance", + func: CuteFn::PipelineStateAdvance, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_producer_acquire", + func: CuteFn::PipelineProducerAcquire, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_producer_expect_tx", + func: CuteFn::PipelineProducerExpectTx, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_consumer_wait", + func: CuteFn::PipelineConsumerWait, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_consumer_release", + func: CuteFn::PipelineConsumerRelease, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::pipeline_producer_tail", + func: CuteFn::PipelineProducerTail, + schema: &[SubstKind::Const, SubstKind::Const, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::copy::tiled::__compiler::shared_tensor_overlay", + func: CuteFn::SharedTensorOverlay, + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + ], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::tiled_mma_slice", + func: CuteFn::TiledMmaSlice, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::fragment_fill", + func: CuteFn::MmaFragmentFill, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::mma_load_scales", + func: CuteFn::MmaLoadScales, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::fragment_slice_k", + func: CuteFn::MmaFragmentSliceK, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::mma_load_a", + func: CuteFn::MmaLoadA, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::mma_partition_b", + func: CuteFn::MmaPartitionB, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::tiled_gemm", + func: CuteFn::TiledGemm, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::sm120::__compiler::epilogue_smem_overlay", + func: CuteFn::EpilogueOverlay, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::sm120::__compiler::epilogue_warp_slice", + func: CuteFn::EpilogueWarpSlice, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::mma::block_scaled::__compiler::epilogue_store_fragment", + func: CuteFn::EpilogueStoreFragment, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::sm120::__compiler::epilogue_sync_reusable", + func: CuteFn::EpilogueSyncReusable, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::sm120::__compiler::epilogue_sync_ready_for_tma", + func: CuteFn::EpilogueSyncReadyForTma, + schema: &[], + }, + CuteFnSpec { + canonical: "cute_rs::epilogue::sm120::__compiler::epilogue_half", + func: CuteFn::EpilogueHalf, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::tma_store_producer_acquire", + func: CuteFn::TmaStoreAcquire, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::tma_store_producer_commit", + func: CuteFn::TmaStoreCommit, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::pipeline::__compiler::tma_store_producer_tail", + func: CuteFn::TmaStoreTail, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::copy::vector::load_tile", + func: CuteFn::LoadTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::copy::vector::store_tile", + func: CuteFn::StoreTile, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::copy::vector::store_global_no_allocate", + func: CuteFn::StoreGlobalNoAllocate, + schema: &[SubstKind::Type, SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::copy::vector::assume", + func: CuteFn::AssumeDiv, + schema: &[SubstKind::Const], + }, + CuteFnSpec { + canonical: "cute_rs::mma::load_matrix_a", + func: CuteFn::LoadMatrixA, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::load_matrix_b", + func: CuteFn::LoadMatrixB, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::load_matrix_b_pair", + func: CuteFn::LoadMatrixBPair, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::mma::load_matrix_b_pair_k_major", + func: CuteFn::LoadMatrixBPairKMajor, + schema: &[SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::copy::tma::copy_tma_2d", + func: CuteFn::CopyTma2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::copy::tma::copy_tma_s2g_2d", + func: CuteFn::CopyTmaS2g2d, + schema: &[SubstKind::Type, SubstKind::Type], + }, + CuteFnSpec { + canonical: "cute_rs::copy::cooperative::copy_g2s", + func: CuteFn::CopyG2S, + schema: &[ + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + SubstKind::Type, + // Rust's argument-position `impl LeadingDimMarker` is one + // inferred type. It keeps the user-facing turbofish at six + // entries while making the pitch promise visible here. + SubstKind::Type, + ], + }, +]; + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that matrix-load entry points accept exactly one shared-layout type parameter. + #[test] + fn matrix_load_entry_points_require_one_shared_layout_type() { + for (name, expected) in [ + ("load_matrix_a", CuteFn::LoadMatrixA), + ("load_matrix_b", CuteFn::LoadMatrixB), + ("load_matrix_b_pair", CuteFn::LoadMatrixBPair), + ("load_matrix_b_pair_k_major", CuteFn::LoadMatrixBPairKMajor), + ] { + let path = format!("cute_rs::mma::{name}"); + let spec = CUTE_FNS.iter().find(|spec| spec.canonical == path); + let spec = spec.unwrap_or_else(|| panic!("matrix load is not recognized: {path}")); + assert_eq!(spec.func, expected, "{path}"); + assert_eq!(spec.schema, &[SubstKind::Type], "{path}"); + } + } + + /// Check that the crate-relative hook query accepts defining paths only. + #[test] + fn hook_paths_match_definitions_not_re_exports() { + for path in [ + "::tensor::make_tensor_read", + "::copy::vector::assume", + "::copy::tma::copy_tma_2d", + "::copy::cooperative::copy_g2s", + "::mma::block_scaled::__compiler::tiled_gemm", + "::epilogue::sm120::__compiler::epilogue_half", + "::sm100::__compiler::sm100_tiled_mma", + ] { + assert!(is_cute_compiler_hook_path(path), "{path}"); + } + // Paths that are not definition paths never match: the root module + // aliases cute_rs once had (`tile`, `tma`, `block_scaled_mma`, retired), + // the retired `epilogue::__compiler` re-export and crate-qualified forms. + for path in [ + "::tile::assume", + "::tma::copy_tma_2d", + "::block_scaled_mma::__compiler::tiled_gemm", + "::epilogue::__compiler::epilogue_half", + "cute_rs::copy::tma::copy_tma_2d", + "copy::tma::copy_tma_2d", + ] { + assert!(!is_cute_compiler_hook_path(path), "{path}"); + } + } + + /// Check that registry names are unique and each generic schema matches its handler. + #[test] + fn table_is_well_formed() { + for (i, spec) in CUTE_FNS.iter().enumerate() { + assert!( + spec.canonical.starts_with("cute_rs::"), + "canonical paths are definition paths in the cute_rs crate" + ); + for other in &CUTE_FNS[i + 1..] { + assert_ne!(spec.canonical, other.canonical, "duplicate canonical path"); + } + } + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/emit.rs new file mode 100644 index 0000000000..86c7bb18a0 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/emit.rs @@ -0,0 +1,556 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the five high-level persistent-scheduler operations. +//! +//! The Rust scheduler stays an ordinary `{ current, stride }` value. Only the +//! selected `WorkTileInfo` is a short-lived compiler handle: +//! +//! ```text +//! new_1d -> has_work -> get_current_work -> tile_idx -> advance_to_next_work +//! ``` + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::types::translate_type; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::CuteTileGridAttr; +use dialect_cute::scheduler_ops::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, CuteSchedulerNew1dOp, + CuteWorkTileCoordinatesOp, +}; +use dialect_mir::attributes::FieldIndexAttr; +use dialect_mir::ops::{MirFieldAddrOp, MirInsertFieldOp, MirLoadOp, MirStoreOp, MirUndefOp}; +use dialect_mir::types::{MirPtrType, MirStructType}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::printable::Printable; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::scheduler::{ + decode_scheduler, decode_scheduler_receiver, decode_work_tile, decode_work_tile_receiver, +}; + +/// Insert an operation after the current anchor, or at the start of the block. +fn insert( + ctx: &mut Context, + op: Ptr, + block: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block, ctx); + } +} + +/// Branch from the last emitted effect to the call continuation; reject diverging calls. +fn finish( + ctx: &mut Context, + effect: Ptr, + target: &Option, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target is not supported")) + ); + }; + Ok(emit_goto(ctx, *target, effect, block_map, loc)) +} + +/// Store a result in the unprojected MIR destination, then emit the continuation branch. +#[allow(clippy::too_many_arguments)] +fn finish_value( + ctx: &mut Context, + producer: Ptr, + value: Value, + destination: &mir::Place, + target: &Option, + block: Ptr, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "{what} destination with projections is not supported" + )) + ); + } + let after_store = value_map + .store_local(ctx, destination.local, value, block, Some(producer)) + .unwrap_or(producer); + finish(ctx, after_store, target, block_map, loc, what) +} + +/// Translate the MIR destination type into its IR carrier type. +fn destination_type( + ctx: &mut Context, + body: &mir::Body, + destination: &mir::Place, + _loc: &Location, + what: &str, +) -> TranslationResult { + let ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} destination type: {error:?}" + ))) + })?; + translate_type(ctx, &ty).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot translate {what} destination type: {}", + error.disp(ctx) + ))) + }) +} + +/// Construct a translated scheduler/work-tile carrier from its scalar fields. +fn build_aggregate( + ctx: &mut Context, + aggregate_type: TypeHandle, + fields: &[Value], + block: Ptr, + prev: Ptr, + loc: &Location, +) -> (Value, Ptr) { + let undef = MirUndefOp::new(ctx, aggregate_type).get_operation(); + undef.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, undef, block, Some(prev)); + let mut aggregate = undef.deref(ctx).get_result(0); + let mut anchor = undef; + + for (index, field) in fields.iter().copied().enumerate() { + let op = Operation::new( + ctx, + MirInsertFieldOp::get_concrete_op_info(), + vec![aggregate_type], + vec![aggregate, field], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirInsertFieldOp::new(op).set_attr_insert_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block, Some(anchor)); + aggregate = op.deref(ctx).get_result(0); + anchor = op; + } + (aggregate, anchor) +} + +/// Decode the scheduler or work-tile receiver's compile-time grid. +fn receiver_grid( + body: &mir::Body, + operand: &mir::Operand, + work_tile: bool, + _loc: &Location, + what: &str, +) -> TranslationResult { + let ty = operand.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} receiver type: {error:?}" + ))) + })?; + let decoded = if work_tile { + decode_work_tile_receiver(&ty) + } else { + decode_scheduler_receiver(&ty) + }; + decoded + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what} receiver: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} receiver has the wrong cute-rs type" + ))) + }) +} + +/// Find the mutable persistent-position field of a referenced scheduler. +fn scheduler_field_addr( + ctx: &mut Context, + scheduler: Value, + index: usize, + block: Ptr, + prev: Option>, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, TypeHandle, Ptr)> { + let field = { + let pointer_type = scheduler.get_type(ctx); + let pointer_type = pointer_type.deref(ctx); + let Some(pointer) = pointer_type.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} receiver must be a pointer")) + ); + }; + let pointee = pointer.pointee.deref(ctx); + let Some(scheduler_type) = pointee.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} receiver must point to a scheduler")) + ); + }; + let Some(field) = scheduler_type.field_types.get(index).copied() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} scheduler has no field {index}")) + ); + }; + (field, pointer.is_mutable, pointer.address_space) + }; + let pointer_type: TypeHandle = MirPtrType::get(ctx, field.0, field.1, field.2).into(); + // Go through the one construction door so `aggregate_ty` is stamped; + // the verifier rejects a `mir.field_addr` without it. + let op = MirFieldAddrOp::build(ctx, scheduler, pointer_type, index as u32)?; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, prev); + Ok((op.deref(ctx).get_result(0), field.0, op)) +} + +/// Load a scheduler field while preserving operation insertion order. +fn load_field( + ctx: &mut Context, + pointer: Value, + field_type: TypeHandle, + block: Ptr, + prev: Ptr, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![field_type], + vec![pointer], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirLoadOp::new(op).set_volatile(ctx, false); + insert(ctx, op, block, Some(prev)); + (op.deref(ctx).get_result(0), op) +} + +/// Create a persistent scheduler from the physical CTA index and grid stride. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_new_1d( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if !args.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "scheduler_new_1d expects no runtime arguments".to_string() + ) + ); + } + let rust_type = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read scheduler_new_1d result: {error:?}" + ))) + })?; + let grid = decode_scheduler(&rust_type) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid scheduler_new_1d result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "scheduler_new_1d must return StaticPersistentTileScheduler".to_string() + )) + })?; + + let operation = CuteSchedulerNew1dOp::new(ctx, grid); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, prev); + let result_type = destination_type(ctx, body, destination, &loc, "scheduler_new_1d")?; + let (result, bundle) = build_aggregate( + ctx, + result_type, + &[operation.current(ctx), operation.stride(ctx)], + block, + op, + &loc, + ); + finish_value( + ctx, + bundle, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "scheduler_new_1d", + ) +} + +/// Test whether the scheduler's current logical tile is inside its fixed grid. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_has_work( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("scheduler_has_work expects one receiver".to_string()) + ); + } + let grid = receiver_grid(body, &args[0], false, &loc, "scheduler_has_work")?; + let (scheduler, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (current_pointer, current_type, after) = scheduler_field_addr( + ctx, + scheduler, + 0, + block, + after.or(prev), + &loc, + "scheduler_has_work", + )?; + let (current, after) = load_field(ctx, current_pointer, current_type, block, after, &loc); + let operation = CuteSchedulerHasWorkOp::new(ctx, current, grid); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + finish_value( + ctx, + op, + operation.has_work(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "scheduler_has_work", + ) +} + +/// Construct the typed work tile at the scheduler's current position. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_current_tile( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("scheduler_current_tile expects one receiver".to_string()) + ); + } + let grid = receiver_grid(body, &args[0], false, &loc, "scheduler_current_tile")?; + let result_type = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read scheduler_current_tile result: {error:?}" + ))) + })?; + if decode_work_tile(&result_type).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid scheduler_current_tile result: {error}" + ))) + })? != Some(grid) + { + return input_err!( + loc, + TranslationErr::unsupported( + "scheduler_current_tile result must use the receiver's tile grid".to_string() + ) + ); + } + let (scheduler, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (current_pointer, current_type, after) = scheduler_field_addr( + ctx, + scheduler, + 0, + block, + after.or(prev), + &loc, + "scheduler_current_tile", + )?; + let (current, after) = load_field(ctx, current_pointer, current_type, block, after, &loc); + let operation = CuteSchedulerCurrentOp::new(ctx, current, grid); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + finish_value( + ctx, + op, + operation.work_tile(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "scheduler_current_tile", + ) +} + +/// Project a work tile into its M, N and batch coordinates. +/// +/// The IR operation also yields the linear tile number for the lowering; +/// `tile_idx` returns the `(m, n, l)` triple, so only three results reach +/// the destination tuple. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_coordinates( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("work_tile_coordinates expects one tile".to_string()) + ); + } + let _grid = receiver_grid(body, &args[0], true, &loc, "work_tile_coordinates")?; + let (tile, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let operation = CuteWorkTileCoordinatesOp::new(ctx, tile); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result_type = destination_type(ctx, body, destination, &loc, "work_tile_coordinates")?; + let (result, bundle) = build_aggregate( + ctx, + result_type, + &[ + operation.m_tile(ctx), + operation.n_tile(ctx), + operation.batch(ctx), + ], + block, + op, + &loc, + ); + finish_value( + ctx, + bundle, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "work_tile_coordinates", + ) +} + +/// Advance persistent work by the physical grid stride and update scheduler storage. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_advance( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("scheduler_advance expects one receiver".to_string()) + ); + } + let grid = receiver_grid(body, &args[0], false, &loc, "scheduler_advance")?; + let (scheduler, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (current_pointer, current_type, after) = scheduler_field_addr( + ctx, + scheduler, + 0, + block, + after.or(prev), + &loc, + "scheduler_advance", + )?; + let (current, after) = load_field(ctx, current_pointer, current_type, block, after, &loc); + let (stride_pointer, stride_type, after) = scheduler_field_addr( + ctx, + scheduler, + 1, + block, + Some(after), + &loc, + "scheduler_advance", + )?; + let (stride, after) = load_field(ctx, stride_pointer, stride_type, block, after, &loc); + let operation = CuteSchedulerAdvanceOp::new(ctx, current, stride, grid); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(after)); + + let store = Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![current_pointer, operation.next(ctx)], + vec![], + 0, + ); + store.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, store, block, Some(op)); + finish(ctx, store, target, block_map, loc, "scheduler_advance") +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs new file mode 100644 index 0000000000..667c218209 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/scheduler/mod.rs @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural Rust types for the persistent scheduler boundary. +//! +//! The scheduler itself remains two ordinary `usize` fields. Only the +//! short-lived selected tile becomes a semantic handle: +//! +//! ```text +//! StaticPersistentTileScheduler ordinary Rust state +//! │ get_current_work +//! ▼ +//! WorkTileInfo !cute.work_tile> +//! ``` + +pub(crate) mod emit; + +use crate::error::{TranslationErr, TranslationResult}; +use dialect_cute::attributes::CuteTileGridAttr; +use dialect_cute::types::CuteWorkTileType; +use pliron::context::Context; +use pliron::input_error_noloc; +use pliron::r#type::TypeHandle; +use rustc_public::ty::{GenericArgs, RigidTy, Ty, TyKind}; + +use super::tensor::{canonical_adt_path, const_arg, exact_schema, unsigned_const}; + +/// Exact defining path of `cute_rs::scheduler::StaticPersistentTileScheduler`; public reexports do not change it. +const SCHEDULER: &str = "cute_rs::scheduler::StaticPersistentTileScheduler"; +/// Exact defining path of `cute_rs::scheduler::WorkTileInfo`; public reexports do not change it. +const WORK_TILE: &str = "cute_rs::scheduler::WorkTileInfo"; + +/// Recognize one exact source ADT rather than a type with a similar printed name. +fn exact_adt(ty: &Ty, expected: &str) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + if canonical_adt_path(&def) != expected { + return Ok(None); + } + exact_schema(&args, &["Const", "Const", "Const"], expected)?; + Ok(Some(args)) +} + +/// Decode constant M/N/batch grid dimensions from a scheduler or work-tile ADT. +fn grid_from_args(args: &GenericArgs, owner: &str) -> Result { + let m_tiles = unsigned_const(&const_arg(args, 0, owner)?, "M tile count")?; + let n_tiles = unsigned_const(&const_arg(args, 1, owner)?, "N tile count")?; + let batches = unsigned_const(&const_arg(args, 2, owner)?, "batch count")?; + let grid = CuteTileGridAttr::new(m_tiles, n_tiles, batches); + if m_tiles == 0 || n_tiles == 0 || batches == 0 { + return Err("scheduler tile-grid sizes must be greater than zero".to_string()); + } + if grid.total_tiles().is_none() { + return Err("scheduler tile-grid product must fit in u64".to_string()); + } + Ok(grid) +} + +/// Decode the ordinary two-scalar scheduler ADT. +pub(crate) fn decode_scheduler(ty: &Ty) -> Result, String> { + exact_adt(ty, SCHEDULER)? + .map(|args| grid_from_args(&args, SCHEDULER)) + .transpose() +} + +/// Decode the semantic work-tile ADT. +pub(crate) fn decode_work_tile(ty: &Ty) -> Result, String> { + exact_adt(ty, WORK_TILE)? + .map(|args| grid_from_args(&args, WORK_TILE)) + .transpose() +} + +/// Decode scheduler facts from a value or one borrowed/raw-pointer receiver. +fn decode_behind_one_pointer( + ty: &Ty, + decode: fn(&Ty) -> Result, String>, +) -> Result, String> { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode(&pointee) + } + _ => decode(ty), + } +} + +/// Recover the constant tile grid from a scheduler receiver. +pub(crate) fn decode_scheduler_receiver(ty: &Ty) -> Result, String> { + decode_behind_one_pointer(ty, decode_scheduler) +} + +/// Recover the constant tile grid from a work-tile receiver. +pub(crate) fn decode_work_tile_receiver(ty: &Ty) -> Result, String> { + decode_behind_one_pointer(ty, decode_work_tile) +} + +/// Build a work-tile IR type that retains the complete logical grid. +fn work_tile_type(ctx: &mut Context, grid: CuteTileGridAttr) -> TypeHandle { + CuteWorkTileType::get(ctx, grid).into() +} + +/// Type-translation hook called before generic ADT lowering. +pub(crate) fn try_translate_work_tile_type( + ctx: &mut Context, + rust_ty: &Ty, +) -> Option> { + match decode_work_tile(rust_ty) { + Ok(Some(grid)) => Some(Ok(work_tile_type(ctx, grid))), + Ok(None) => None, + Err(error) => Some(Err(input_error_noloc!(TranslationErr::unsupported( + format!("invalid cute-rs scheduler work tile: {error}") + )))), + } +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that work-tile types retain all three grid dimensions. + #[test] + fn work_tile_type_keeps_the_whole_grid() { + let mut ctx = Context::new(); + dialect_cute::register(&mut ctx); + let ty = work_tile_type(&mut ctx, CuteTileGridAttr::new(16, 16, 1)); + let ty = ty.deref(&ctx); + let tile = ty.downcast_ref::().unwrap(); + assert_eq!(tile.grid, CuteTileGridAttr::new(16, 16, 1)); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs new file mode 100644 index 0000000000..4927284260 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100.rs @@ -0,0 +1,387 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Import SM100 CuTe contracts before any target instruction is selected. +//! +//! Layouts come from structural Rust type arguments. The ordinary scalar +//! carriers keep source control flow intact; only these recognized unsafe +//! boundaries recover the caller's shared-address-space promise. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}; +use dialect_cute::sm100_ops::*; +use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; +use dialect_mir::ops::MirCastOp; +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; +use rustc_public::ty::{GenericArgKind, GenericArgs, RigidTy, TyKind}; + +use super::smem_mma::emit::{finish, insert}; + +/// Paired SM100 GEMM, TMEM and pipeline hooks recognized by the importer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Sm100Fn { + /// Allocate a two-CTA tensor-memory region. + TmemAlloc, + /// Release a completed two-CTA tensor-memory region. + TmemDealloc, + /// Issue the paired A×B0/A×B1 FP16 products with A-collector reuse. + TiledMma, + /// Convert a TMEM accumulator slice into its shared FP16 output tile. + TmemEpilogue, + /// Issue one CTA's portion of a collective two-CTA TMA load. + ClusterTmaLoad, + /// Issue a shared-to-global tile store. + TmaStore, + /// Commit this thread's outstanding output stores. + StoreCommit, + /// Wait until the oldest output-store slot is reusable. + StoreAcquire, + /// Drain committed output stores before freeing shared storage. + StoreTail, + /// Initialize the input ring's full and empty barriers. + PipelineInit, + /// Wait for an empty input slot before overwriting it. + PipelineAcquire, + /// Set the current full barrier's expected transaction bytes. + PipelineExpect, + /// Wait for input copies before tensor-core access. + PipelineWait, + /// Attach input-buffer release to MMA completion. + PipelineRelease, + /// Initialize the single accumulator slot's barriers. + AccumulatorInit, + /// Wait until epilogue readers release the accumulator. + AccumulatorAcquire, + /// Publish the accumulator after the final MMA completes. + AccumulatorCommit, + /// Wait for a ready accumulator before TMEM loads. + AccumulatorWait, + /// Signal one consumer warp's completed accumulator reads. + AccumulatorRelease, +} + +/// Report an unsupported source contract at this CuTe boundary. +fn invalid(message: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid SM100 CuTe boundary: {message}" + ))) +} + +/// Extract monomorphized generic arguments from a direct constant function callee. +fn generics(func: &mir::Operand) -> TranslationResult { + let mir::Operand::Constant(constant) = func else { + return Err(invalid("marker requires a direct constant callee")); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err(invalid("marker callee must be a function definition")); + }; + Ok(args) +} + +/// Evaluate one generic constant as a checked `u32` plan parameter. +fn constant(args: &GenericArgs, index: usize, loc: &Location) -> TranslationResult { + let Some(GenericArgKind::Const(value)) = args.0.get(index) else { + return Err(invalid(format!( + "generic argument {index} must be a constant" + ))); + }; + let value = super::layout::const_u64(value, "SM100 plan", loc)?; + u32::try_from(value).map_err(|_| invalid("SM100 plan constant exceeds u32")) +} + +/// Decode one shared-layout type argument into a verified CuTe layout attribute. +fn layout(args: &GenericArgs, index: usize) -> TranslationResult { + let Some(GenericArgKind::Type(ty)) = args.0.get(index) else { + return Err(invalid(format!( + "generic argument {index} must be a layout type" + ))); + }; + super::static_config::decode_smem_layout(ty) + .map(CuteComposedLayoutAttr) + .map_err(invalid) +} + +/// Build the input-pipeline plan from stage, byte-count and arrival constants. +fn pipeline_plan( + args: &GenericArgs, + loc: &Location, +) -> TranslationResult { + Ok(CuteSm100PipelinePlanAttr { + stages: constant(args, 0, loc)?, + transaction_bytes: constant(args, 1, loc)?, + consumer_arrivals: 1, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + }) +} + +/// Build the single-slot accumulator plan with its consumer-warp arrival count. +fn accumulator_plan( + args: &GenericArgs, + loc: &Location, +) -> TranslationResult { + let consumers = constant(args, 0, loc)? + .checked_mul(2) + .ok_or_else(|| invalid("accumulator consumer arrival count overflows"))?; + Ok(CuteSm100PipelinePlanAttr { + stages: 1, + transaction_bytes: 0, + consumer_arrivals: consumers, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + }) +} + +/// List the operands whose unsafe source contract permits shared-address-space recovery. +fn shared_operands(kind: Sm100Fn) -> &'static [usize] { + use Sm100Fn::*; + match kind { + TmemAlloc => &[0], + TmemDealloc | StoreCommit | StoreAcquire | StoreTail => &[], + TmaStore => &[1], + TiledMma => &[1, 2, 3], + TmemEpilogue => &[1], + ClusterTmaLoad => &[1, 4], + PipelineInit | AccumulatorInit => &[0, 1], + PipelineAcquire | PipelineExpect | PipelineWait | PipelineRelease | AccumulatorAcquire + | AccumulatorCommit | AccumulatorWait | AccumulatorRelease => &[0], + } +} + +/// Return the exact scalar-carrier count required by this compiler hook. +fn operand_count(kind: Sm100Fn) -> usize { + use Sm100Fn::*; + match kind { + TmemAlloc | TmemDealloc | AccumulatorCommit | AccumulatorRelease => 1, + PipelineInit | PipelineExpect | PipelineRelease | AccumulatorInit | AccumulatorAcquire + | AccumulatorWait => 2, + PipelineAcquire | PipelineWait => 3, + TmemEpilogue | TmaStore => 4, + StoreCommit | StoreAcquire | StoreTail => 0, + TiledMma => 5, + ClusterTmaLoad => 6, + } +} + +/// Decode an SM100 hook's static plan and emit its typed operation with scalar carriers. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + kind: Sm100Fn, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != operand_count(kind) { + return Err(invalid(format!( + "{kind:?} expects {} operands, got {}", + operand_count(kind), + args.len() + ))); + } + let static_args = generics(func)?; + let mut operands: Vec = Vec::with_capacity(args.len()); + let mut after = prev; + for arg in args { + let (value, last) = + translate_operand(ctx, body, arg, value_map, block, after, loc.clone())?; + after = last.or(after); + operands.push(value); + } + for &index in shared_operands(kind) { + let pointer = operands[index] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() + .ok_or_else(|| invalid(format!("{kind:?} operand {index} must be a shared pointer")))?; + if !matches!( + pointer.address_space, + address_space::GENERIC | address_space::SHARED + ) { + return Err(invalid(format!( + "{kind:?} operand {index} has non-shared address space {}", + pointer.address_space + ))); + } + let shared: TypeHandle = MirPtrType::get( + ctx, + pointer.pointee, + pointer.is_mutable, + address_space::SHARED, + ) + .into(); + if operands[index].get_type(ctx) != shared { + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![shared], + vec![operands[index]], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + if dialect_mir::types::type_contains_concrete_pointer_kind(ctx, shared) { + MirCastOp::new(cast) + .set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::RawAddress); + } + insert(ctx, cast, block, after); + operands[index] = cast.deref(ctx).get_result(0); + after = Some(cast); + } + } + use Sm100Fn::*; + let operation = match kind { + TmemAlloc | TmemDealloc => { + let plan = CuteSm100TmemPlanAttr { + columns: constant(&static_args, 0, &loc)?, + cta_group: 2, + }; + if kind == TmemAlloc { + CuteSm100TmemAllocOp::new(ctx, operands, plan).get_operation() + } else { + CuteSm100TmemDeallocOp::new(ctx, operands, plan).get_operation() + } + } + TiledMma => { + let plan = CuteSm100MmaPlanAttr { + m: 256, + n0: constant(&static_args, 3, &loc)?, + n1: constant(&static_args, 4, &loc)?, + k: constant(&static_args, 5, &loc)?, + tmem_columns: 512, + cta_group: 2, + collector: CuteSm100CollectorAttr::AFillLastUse, + a_layout: layout(&static_args, 0)?, + b0_layout: layout(&static_args, 1)?, + b1_layout: layout(&static_args, 2)?, + }; + CuteSm100TiledMmaOp::new(ctx, operands, plan).get_operation() + } + TmemEpilogue => CuteSm100TmemEpilogueOp::new( + ctx, + operands, + CuteSm100EpiloguePlanAttr { + layout: layout(&static_args, 0)?, + rows: 128, + columns: 32, + }, + ) + .get_operation(), + ClusterTmaLoad => CuteSm100ClusterTmaLoadOp::new( + ctx, + operands, + CuteSm100ClusterCopyPlanAttr { + layout: layout(&static_args, 0)?, + cta_group: 2, + leader_rank: 0, + }, + ) + .get_operation(), + TmaStore => CuteSm100TmaStoreOp::new( + ctx, + operands, + CuteSm100EpiloguePlanAttr { + layout: layout(&static_args, 0)?, + rows: 128, + columns: 32, + }, + ) + .get_operation(), + StoreCommit => CuteSm100StoreCommitOp::new( + ctx, + operands, + CuteTmaStorePipelineAttr { + stages: constant(&static_args, 0, &loc)?, + }, + ) + .get_operation(), + StoreAcquire => CuteSm100StoreAcquireOp::new( + ctx, + operands, + CuteTmaStorePipelineAttr { + stages: constant(&static_args, 0, &loc)?, + }, + ) + .get_operation(), + StoreTail => CuteSm100StoreTailOp::new( + ctx, + operands, + CuteTmaStorePipelineAttr { + stages: constant(&static_args, 0, &loc)?, + }, + ) + .get_operation(), + PipelineInit => { + CuteSm100PipelineInitOp::new(ctx, operands, pipeline_plan(&static_args, &loc)?) + .get_operation() + } + PipelineAcquire => { + CuteSm100PipelineAcquireOp::new(ctx, operands, pipeline_plan(&static_args, &loc)?) + .get_operation() + } + PipelineExpect => { + CuteSm100PipelineExpectOp::new(ctx, operands, pipeline_plan(&static_args, &loc)?) + .get_operation() + } + PipelineWait => { + CuteSm100PipelineWaitOp::new(ctx, operands, pipeline_plan(&static_args, &loc)?) + .get_operation() + } + PipelineRelease => { + CuteSm100PipelineReleaseOp::new(ctx, operands, pipeline_plan(&static_args, &loc)?) + .get_operation() + } + AccumulatorInit => { + CuteSm100AccumulatorInitOp::new(ctx, operands, accumulator_plan(&static_args, &loc)?) + .get_operation() + } + AccumulatorAcquire => { + CuteSm100AccumulatorAcquireOp::new(ctx, operands, accumulator_plan(&static_args, &loc)?) + .get_operation() + } + AccumulatorCommit => { + CuteSm100AccumulatorCommitOp::new(ctx, operands, accumulator_plan(&static_args, &loc)?) + .get_operation() + } + AccumulatorWait => { + CuteSm100AccumulatorWaitOp::new(ctx, operands, accumulator_plan(&static_args, &loc)?) + .get_operation() + } + AccumulatorRelease => { + CuteSm100AccumulatorReleaseOp::new(ctx, operands, accumulator_plan(&static_args, &loc)?) + .get_operation() + } + }; + operation.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, operation, block, after); + finish( + ctx, + operation, + target, + block_map, + loc, + "SM100 CuTe operation", + ) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs new file mode 100644 index 0000000000..da8681bee2 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/sm100_attention.rs @@ -0,0 +1,229 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Import independent SM100 attention MMA and typed tensor-copy contracts. +//! +//! Layouts come from structural Rust type arguments. The ordinary scalar +//! carriers keep source control flow intact; only these recognized unsafe +//! boundaries recover the caller's shared-address-space promise. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::CuteComposedLayoutAttr; +use dialect_cute::sm100_attention_ops::*; +use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; +use dialect_mir::ops::MirCastOp; +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; +use rustc_public::ty::{GenericArgKind, GenericArgs, RigidTy, TyKind}; + +use super::smem_mma::emit::{finish, insert}; + +/// Independent SM100 attention hooks recognized by the importer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttentionFn { + /// Issue independent shared/shared FP16 MMA into TMEM. + MmaSs, + /// Issue independent TMEM/shared FP16 MMA into TMEM. + MmaTs, + /// Load a rank-three global tile into the attention shared layout. + TmaLoad, + /// Issue a shared-to-global tile store. + TmaStore, +} +/// Report an unsupported source contract at this CuTe boundary. +fn invalid(message: impl core::fmt::Display) -> pliron::result::Error { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid SM100 CuTe boundary: {message}" + ))) +} + +/// Extract monomorphized generic arguments from a direct constant function callee. +fn generics(func: &mir::Operand) -> TranslationResult { + let mir::Operand::Constant(constant) = func else { + return Err(invalid("marker requires a direct constant callee")); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err(invalid("marker callee must be a function definition")); + }; + Ok(args) +} + +/// Evaluate one generic constant as a checked `u32` plan parameter. +fn constant(args: &GenericArgs, index: usize, loc: &Location) -> TranslationResult { + let Some(GenericArgKind::Const(value)) = args.0.get(index) else { + return Err(invalid(format!( + "generic argument {index} must be a constant" + ))); + }; + let value = super::layout::const_u64(value, "SM100 plan", loc)?; + u32::try_from(value).map_err(|_| invalid("SM100 plan constant exceeds u32")) +} + +/// Decode an independent attention layout and its MN-major selector from a type argument. +fn layout(args: &GenericArgs, index: usize) -> TranslationResult<(CuteComposedLayoutAttr, bool)> { + let Some(GenericArgKind::Type(ty)) = args.0.get(index) else { + return Err(invalid("attention layout must be a type argument")); + }; + super::static_config::decode_attention_layout(ty) + .map(|(l, m)| (CuteComposedLayoutAttr(l), m)) + .map_err(invalid) +} +/// List the operands whose unsafe source contract permits shared-address-space recovery. +fn shared_operands(kind: AttentionFn) -> &'static [usize] { + match kind { + AttentionFn::MmaSs => &[1, 2], + AttentionFn::MmaTs => &[2], + AttentionFn::TmaLoad => &[1, 5], + AttentionFn::TmaStore => &[1], + } +} +/// Return the exact scalar-carrier count required by this compiler hook. +fn operand_count(kind: AttentionFn) -> usize { + match kind { + AttentionFn::MmaSs => 6, + AttentionFn::MmaTs => 5, + AttentionFn::TmaLoad => 7, + AttentionFn::TmaStore => 5, + } +} +/// Import independent attention MMA or TMA while preserving source control flow and pointer authority. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit( + ctx: &mut Context, + body: &mir::Body, + func: &mir::Operand, + args: &[mir::Operand], + kind: AttentionFn, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != operand_count(kind) { + return Err(invalid(format!( + "{kind:?} expects {} operands, got {}", + operand_count(kind), + args.len() + ))); + } + let static_args = generics(func)?; + let mut operands: Vec = Vec::with_capacity(args.len()); + let mut after = prev; + for arg in args { + let (value, last) = + translate_operand(ctx, body, arg, value_map, block, after, loc.clone())?; + after = last.or(after); + operands.push(value); + } + for &index in shared_operands(kind) { + let pointer = operands[index] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() + .ok_or_else(|| invalid(format!("{kind:?} operand {index} must be a shared pointer")))?; + if !matches!( + pointer.address_space, + address_space::GENERIC | address_space::SHARED + ) { + return Err(invalid(format!( + "{kind:?} operand {index} has non-shared address space {}", + pointer.address_space + ))); + } + let shared: TypeHandle = MirPtrType::get( + ctx, + pointer.pointee, + pointer.is_mutable, + address_space::SHARED, + ) + .into(); + if operands[index].get_type(ctx) != shared { + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![shared], + vec![operands[index]], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + if dialect_mir::types::type_contains_concrete_pointer_kind(ctx, shared) { + MirCastOp::new(cast) + .set_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::RawAddress); + } + insert(ctx, cast, block, after); + operands[index] = cast.deref(ctx).get_result(0); + after = Some(cast); + } + } + let operation = match kind { + AttentionFn::MmaSs | AttentionFn::MmaTs => { + let ts = matches!(kind, AttentionFn::MmaTs); + let offset = if ts { 1 } else { 2 }; + let m = constant(&static_args, offset, &loc)?; + let n = constant(&static_args, offset + 1, &loc)?; + let k = constant(&static_args, offset + 2, &loc)?; + let cta_group = constant(&static_args, offset + 3, &loc)?; + if !matches!(cta_group, 1 | 2) { + return Err(invalid("CTA group must be 1 or 2")); + } + let (a_layout, a_mn) = if ts { + (shared_layout(m, k, false).map_err(invalid)?, false) + } else { + layout(&static_args, 0)? + }; + let (b_layout, b_mn) = layout(&static_args, if ts { 0 } else { 1 })?; + let plan = CuteSm100MmaAttr { + m, + n, + k, + cta_group, + a_layout, + b_layout, + a_mn, + b_mn, + }; + if ts { + CuteSm100MmaTsOp::new(ctx, operands, plan).get_operation() + } else { + CuteSm100MmaSsOp::new(ctx, operands, plan).get_operation() + } + } + AttentionFn::TmaLoad | AttentionFn::TmaStore => { + let plan = CuteSm100TmaAttr { + layout: layout(&static_args, 0)?.0, + cta_group: constant(&static_args, 1, &loc)?, + }; + if matches!(kind, AttentionFn::TmaLoad) { + CuteSm100TmaLoadOp::new(ctx, operands, plan).get_operation() + } else { + CuteSm100TmaStoreOp::new(ctx, operands, plan).get_operation() + } + } + }; + operation.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, operation, block, after); + finish( + ctx, + operation, + target, + block_map, + loc, + "SM100 attention CuTe operation", + ) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/emit.rs new file mode 100644 index 0000000000..6bb8c191df --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/emit.rs @@ -0,0 +1,1209 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the semantic shared-memory MXFP4 MMA chain. +//! +//! Every result keeps the ordinary Rust carrier it had before recognition: +//! pointers and coordinates stay scalars, fragments stay register aggregates, +//! and the accumulator remains the loop-carried 64-f32 value. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::types::translate_type; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::{ + CuteMmaCarrierKindAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, + CuteTensorLayoutAttr, CuteTiledMmaPlanAttr, +}; +use dialect_cute::smem_mma_ops::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, +}; +use dialect_cute::types::{CuteSmemTensorType, CuteTensorViewType}; +use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr}; +use dialect_mir::ops::{ + MirCastOp, MirExtractFieldOp, MirFloatConstantOp, MirInsertFieldOp, MirStoreOp, MirUndefOp, +}; +use dialect_mir::types::{MirFP16Type, MirPtrType, MirStructType}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::{ + attributes::FPSingleAttr, + types::{FP32Type, IntegerType, Signedness}, +}; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::printable::Printable; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::emit::{load_through, recover_cta_shared_pointer, struct_field_addr}; +use super::super::smem_mma::{ + SharedTensorKind, SharedTensorRust, TiledMmaRust, canonical_data_placement, decode_b_tile, + decode_shared_tensor, decode_shared_tensor_receiver, decode_tiled_mma, + decode_tiled_mma_receiver, is_accumulator, is_scale_k64, is_scale_stage, +}; + +/// Insert an operation after the current anchor, or at the start of the block. +pub(in crate::translator::cute) fn insert( + ctx: &mut Context, + op: Ptr, + block: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block, ctx); + } +} + +/// Branch from the last emitted effect to the call continuation; reject diverging calls. +pub(in crate::translator::cute) fn finish( + ctx: &mut Context, + effect: Ptr, + target: &Option, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target is not supported")) + ); + }; + Ok(emit_goto(ctx, *target, effect, block_map, loc)) +} + +/// Store a result in the unprojected MIR destination, then emit the continuation branch. +#[allow(clippy::too_many_arguments)] +pub(in crate::translator::cute) fn finish_value( + ctx: &mut Context, + producer: Ptr, + value: Value, + destination: &mir::Place, + target: &Option, + block: Ptr, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "{what} destination with projections is not supported" + )) + ); + } + let after = value_map + .store_local(ctx, destination.local, value, block, Some(producer)) + .unwrap_or(producer); + finish(ctx, after, target, block_map, loc, what) +} + +/// Read the Rust destination type before decoding semantic fragment facts. +pub(in crate::translator::cute) fn destination_rust_type( + body: &mir::Body, + destination: &mir::Place, + what: &str, +) -> TranslationResult { + destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} destination type: {error:?}" + ))) + }) +} + +/// Translate the MIR destination type into its IR carrier type. +pub(in crate::translator::cute) fn destination_type( + ctx: &mut Context, + body: &mir::Body, + destination: &mir::Place, + what: &str, +) -> TranslationResult { + let ty = destination_rust_type(body, destination, what)?; + translate_type(ctx, &ty).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot translate {what} destination type: {}", + error.disp(ctx) + ))) + }) +} + +/// Read a MIR operand's Rust type for exact source-API recognition. +pub(in crate::translator::cute) fn operand_rust_type( + body: &mir::Body, + operand: &mir::Operand, + what: &str, +) -> TranslationResult { + operand.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} operand type: {error:?}" + ))) + }) +} + +/// Extract a translated field from a value or referenced aggregate. +pub(in crate::translator::cute) fn aggregate_field( + ctx: &mut Context, + aggregate: Value, + index: usize, + what: &str, + block: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let aggregate_ty = aggregate.get_type(ctx); + if aggregate_ty + .deref(ctx) + .downcast_ref::() + .is_some() + { + let (pointer, field_ty, addr) = + struct_field_addr(ctx, aggregate, index, what, block, prev, loc)?; + return Ok(load_through(ctx, pointer, field_ty, block, Some(addr), loc)); + } + let field_ty = { + let aggregate_ty = aggregate_ty.deref(ctx); + let Some(structure) = aggregate_ty.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} must be a Rust struct carrier")) + ); + }; + let Some(field_ty) = structure.field_types.get(index).copied() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} has no field {index}")) + ); + }; + field_ty + }; + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![field_ty], + vec![aggregate], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block, prev); + Ok((op.deref(ctx).get_result(0), op)) +} + +/// Build a Rust carrier from a semantic prefix, filling remaining fields with undefined values. +pub(in crate::translator::cute) fn build_struct_prefix( + ctx: &mut Context, + ty: TypeHandle, + fields: &[Value], + block: Ptr, + prev: Ptr, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let field_types = { + let ty_ref = ty.deref(ctx); + let Some(structure) = ty_ref.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("compiler result must be a Rust struct".to_string()) + ); + }; + structure.field_types.clone() + }; + if fields.len() > field_types.len() { + return input_err!( + loc.clone(), + TranslationErr::unsupported("compiler result has too few carrier fields".to_string()) + ); + } + let undef = MirUndefOp::new(ctx, ty).get_operation(); + undef.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, undef, block, Some(prev)); + let mut value = undef.deref(ctx).get_result(0); + let mut anchor = undef; + for (index, field) in fields.iter().copied().enumerate() { + if field.get_type(ctx) != field_types[index] { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "compiler result field {index} has the wrong carrier type" + )) + ); + } + let op = Operation::new( + ctx, + MirInsertFieldOp::get_concrete_op_info(), + vec![ty], + vec![value, field], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirInsertFieldOp::new(op).set_attr_insert_index(ctx, FieldIndexAttr(index as u32)); + insert(ctx, op, block, Some(anchor)); + value = op.deref(ctx).get_result(0); + anchor = op; + } + Ok((value, anchor)) +} + +/// Cast a pointer carrier while retaining its concrete pointer-kind authority. +pub(in crate::translator::cute) fn pointer_cast( + ctx: &mut Context, + value: Value, + target: TypeHandle, + block: Ptr, + prev: Ptr, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + if value.get_type(ctx) == target { + return Ok((value, prev)); + } + if value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + || target.deref(ctx).downcast_ref::().is_none() + { + return input_err!( + loc.clone(), + TranslationErr::unsupported("compiler carrier cast must be pointer-to-pointer") + ); + } + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![target], + vec![value], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + // This cast adapts a compiler-produced carrier to the exact Rust type + // declared on the cute-rs stub. When that type names a concrete raw + // pointer kind (e.g. `*mut f16` = RawMut), the kind transition needs an + // explicit authority. `RawAddress` is the same authority the importer + // stamps on `DynamicSharedArray`'s result: the stub hands out a raw + // pointer to compiler-managed storage, exactly a `&raw` boundary. + if dialect_mir::types::type_contains_concrete_pointer_kind(ctx, target) { + MirCastOp::new(cast).set_pointer_kind_authority( + ctx, + dialect_mir::attributes::MirPointerKindAuthorityAttr::RawAddress, + ); + } + insert(ctx, cast, block, Some(prev)); + Ok((cast.deref(ctx).get_result(0), cast)) +} + +/// Build the MXF4 MMA plan from a decoded tiled-MMA layout. +fn plan(facts: &TiledMmaRust) -> CuteTiledMmaPlanAttr { + CuteTiledMmaPlanAttr::mxf4_128x128x128(facts.placement.clone()) +} + +/// Return the canonical SM120 block-scaled MMA plan. +fn fixed_plan() -> CuteTiledMmaPlanAttr { + CuteTiledMmaPlanAttr::mxf4_128x128x128(canonical_data_placement()) +} + +/// Build a shared-view type with packed-data or scale placement and operand role. +fn smem_view_type(ctx: &mut Context, facts: &SharedTensorRust) -> (TypeHandle, TypeHandle) { + let logical: TypeHandle = FP32Type::get(ctx).into(); + let (storage, alignment, format, layout): (TypeHandle, u64, _, _) = match facts.kind { + SharedTensorKind::Data => ( + MirFP16Type::get(ctx).into(), + 16, + CuteTensorFormatAttr::E2M1, + CuteTensorLayoutAttr::KMajor, + ), + SharedTensorKind::Scale => ( + IntegerType::get(ctx, 32, Signedness::Unsigned).into(), + 4, + CuteTensorFormatAttr::UE8M0, + CuteTensorLayoutAttr::BlockScaleKMajor(32), + ), + }; + let tensor: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + alignment, + format, + facts.role, + layout, + ) + .into(); + ( + CuteSmemTensorType::get(ctx, tensor, facts.placement.clone()).into(), + storage, + ) +} + +/// Require a recognized shared tensor operand and recover its layout and role. +fn decode_shared_operand( + body: &mir::Body, + operand: &mir::Operand, + what: &str, +) -> TranslationResult { + let ty = operand_rust_type(body, operand, what)?; + decode_shared_tensor_receiver(&ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what}: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} must be a shared MMA tensor" + ))) + }) +} + +/// Require a recognized tiled-MMA operand and recover its shared placement. +fn decode_mma_operand( + body: &mir::Body, + operand: &mir::Operand, + what: &str, +) -> TranslationResult { + let ty = operand_rust_type(body, operand, what)?; + decode_tiled_mma_receiver(&ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what}: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} must be Mxf4TiledMma" + ))) + }) +} + +/// Extract a shared tensor's storage pointer and capacity from its Rust carrier. +#[allow(clippy::too_many_arguments)] +fn shared_carriers( + ctx: &mut Context, + body: &mir::Body, + operand: &mir::Operand, + facts: &SharedTensorRust, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + loc: &Location, + what: &str, +) -> TranslationResult<(Value, Value, Ptr)> { + let (shared, after) = + translate_operand(ctx, body, operand, value_map, block, prev, loc.clone())?; + let (carrier, after) = aggregate_field(ctx, shared, 0, what, block, after.or(prev), loc)?; + let (base, after) = aggregate_field(ctx, carrier, 0, what, block, Some(after), loc)?; + let (capacity, after) = aggregate_field(ctx, carrier, 1, what, block, Some(after), loc)?; + let (_, storage) = smem_view_type(ctx, facts); + let (base, after) = recover_cta_shared_pointer(ctx, base, storage, what, block, after, loc)?; + Ok((base, capacity, after)) +} + +/// Extract the calling lane from the tiled-MMA slice carrier. +#[allow(clippy::too_many_arguments)] +fn mma_lane( + ctx: &mut Context, + body: &mir::Body, + operand: &mir::Operand, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let (mma, after) = translate_operand(ctx, body, operand, value_map, block, prev, loc.clone())?; + aggregate_field(ctx, mma, 0, "tiled MMA slice", block, after.or(prev), loc) +} + +/// Insert an FP32 zero used to initialize accumulator carriers. +fn emit_zero_f32( + ctx: &mut Context, + block: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let ty: TypeHandle = FP32Type::get(ctx).into(); + let op = Operation::new( + ctx, + MirFloatConstantOp::get_concrete_op_info(), + vec![ty], + vec![], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirFloatConstantOp::new(op).set_attr_float_value(ctx, FPSingleAttr::from(0.0_f32)); + insert(ctx, op, block, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Read a translated structure field type with a checked field index. +pub(in crate::translator::cute) fn struct_field_type( + ctx: &Context, + ty: TypeHandle, + index: usize, + what: &str, +) -> TranslationResult { + ty.deref(ctx) + .downcast_ref::() + .and_then(|structure| structure.field_types.get(index).copied()) + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} has no Rust carrier field {index}" + ))) + }) +} + +/// Construct a semantic shared view over the caller's data or scale storage. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_overlay( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "shared_tensor_overlay expects base and capacity".to_string() + ) + ); + } + let rust_ty = destination_rust_type(body, destination, "shared_tensor_overlay")?; + let facts = decode_shared_tensor(&rust_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid shared_tensor_overlay result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "shared_tensor_overlay must return SharedTensor".to_string() + )) + })?; + let (base, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (capacity, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + after.or(prev), + loc.clone(), + )?; + let (view, storage) = smem_view_type(ctx, &facts); + let anchor = after.or(prev).ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "shared_tensor_overlay arguments produced no insertion anchor".to_string() + )) + })?; + let (base, anchor) = recover_cta_shared_pointer( + ctx, + base, + storage, + "shared_tensor_overlay base", + block, + anchor, + &loc, + )?; + let operation = CuteSmemTensorOverlayOp::new(ctx, base, capacity, view); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(anchor)); + + let outer_ty = destination_type(ctx, body, destination, "shared_tensor_overlay")?; + let carrier_ty = struct_field_type(ctx, outer_ty, 0, "SharedTensor")?; + let base_ty = struct_field_type(ctx, carrier_ty, 0, "SharedTensor carrier")?; + let (base, anchor) = pointer_cast(ctx, operation.base(ctx), base_ty, block, op, &loc)?; + let (carrier, anchor) = build_struct_prefix( + ctx, + carrier_ty, + &[base, operation.capacity(ctx)], + block, + anchor, + &loc, + )?; + let (result, anchor) = build_struct_prefix(ctx, outer_ty, &[carrier], block, anchor, &loc)?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "shared_tensor_overlay", + ) +} + +/// Create the per-lane tiled-MMA carrier with its static shared layout. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_mma_slice( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("tiled_mma_slice expects one lane".to_string()) + ); + } + let rust_ty = destination_rust_type(body, destination, "tiled_mma_slice")?; + let facts = decode_tiled_mma(&rust_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid tiled_mma_slice result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tiled_mma_slice must return Mxf4TiledMma".to_string() + )) + })?; + let (lane, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let operation = CuteTiledMmaSliceOp::new(ctx, lane, plan(&facts)); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result_ty = destination_type(ctx, body, destination, "tiled_mma_slice")?; + let (result, anchor) = build_struct_prefix( + ctx, + result_ty, + &[operation.sliced_lane(ctx)], + block, + op, + &loc, + )?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tiled_mma_slice", + ) +} + +/// Initialize the fixed FP32 accumulator with the supplied scalar, or zero when omitted. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_fragment_fill( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() > 1 { + return input_err!( + loc, + TranslationErr::unsupported("fragment_fill expects zero or one fill value") + ); + } + if !is_accumulator(&destination_rust_type(body, destination, "fragment_fill")?).map_err( + |error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid fragment_fill result: {error}" + ))) + }, + )? { + return input_err!( + loc, + TranslationErr::unsupported("fragment_fill must return Mxf4AccumulatorTile2x8") + ); + } + let (fill, anchor) = if let Some(fill) = args.first() { + let (value, after) = + translate_operand(ctx, body, fill, value_map, block, prev, loc.clone())?; + let anchor = after.or(prev).ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "fragment_fill value produced no insertion anchor" + )) + })?; + (value, anchor) + } else { + emit_zero_f32(ctx, block, prev, &loc) + }; + let result_ty = destination_type(ctx, body, destination, "fragment_fill")?; + let operation = CuteFragmentFillOp::new( + ctx, + fill, + result_ty, + fixed_plan(), + CuteMmaCarrierKindAttr::Accumulator, + ); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(anchor)); + finish_value( + ctx, + op, + operation.fragment(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "fragment_fill", + ) +} + +/// Load the canonical A/B scale atom into typed scale-fragment carriers. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load_scales( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 5 { + return input_err!( + loc, + TranslationErr::unsupported( + "mma_load_scales expects MMA, A/B scales, warp-M, and warp-N" + ) + ); + } + let mma = decode_mma_operand(body, &args[0], "mma_load_scales MMA")?; + let scale_a = decode_shared_operand(body, &args[1], "mma_load_scales A scales")?; + let scale_b = decode_shared_operand(body, &args[2], "mma_load_scales B scales")?; + if scale_a.kind != SharedTensorKind::Scale + || scale_b.kind != SharedTensorKind::Scale + || scale_a.role != dialect_cute::attributes::CuteTensorRoleAttr::Mkl + || scale_b.role != dialect_cute::attributes::CuteTensorRoleAttr::Nkl + || !is_scale_stage(&destination_rust_type( + body, + destination, + "mma_load_scales", + )?) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid mma_load_scales result: {error}" + ))) + })? + { + return input_err!( + loc, + TranslationErr::unsupported( + "mma_load_scales needs M/N scale atoms and Mxf4ScaleTile128 output" + ) + ); + } + + let (lane, after) = mma_lane(ctx, body, &args[0], block, prev, value_map, &loc)?; + let (a_base, a_capacity, after) = shared_carriers( + ctx, + body, + &args[1], + &scale_a, + block, + Some(after), + value_map, + &loc, + "mma_load_scales A scales", + )?; + let (b_base, b_capacity, after) = shared_carriers( + ctx, + body, + &args[2], + &scale_b, + block, + Some(after), + value_map, + &loc, + "mma_load_scales B scales", + )?; + let (warp_m, after) = translate_operand( + ctx, + body, + &args[3], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (warp_n, after) = + translate_operand(ctx, body, &args[4], value_map, block, after, loc.clone())?; + let (a_view, _) = smem_view_type(ctx, &scale_a); + let (b_view, _) = smem_view_type(ctx, &scale_b); + let result_ty = destination_type(ctx, body, destination, "mma_load_scales")?; + let operation = CuteMmaLoadScalesOp::new( + ctx, + lane, + a_base, + a_capacity, + b_base, + b_capacity, + warp_m, + warp_n, + a_view, + b_view, + result_ty, + plan(&mma), + ); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + finish_value( + ctx, + op, + operation.scales(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "mma_load_scales", + ) +} + +/// Select the scale pairs for one K64 half of a loaded scale tile. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_slice_k( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 + || !is_scale_stage(&operand_rust_type( + body, + &args[0], + "fragment_slice_k input", + )?) + .unwrap_or(false) + || !is_scale_k64(&destination_rust_type( + body, + destination, + "fragment_slice_k", + )?) + .unwrap_or(false) + { + return input_err!( + loc, + TranslationErr::unsupported( + "fragment_slice_k expects Mxf4ScaleTile128, K-half, and Mxf4ScalePairs128 output" + ) + ); + } + let (scales, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (half, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + after.or(prev), + loc.clone(), + )?; + let result_ty = destination_type(ctx, body, destination, "fragment_slice_k")?; + let operation = CuteFragmentSliceKOp::new(ctx, scales, half, result_ty, fixed_plan()); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + finish_value( + ctx, + op, + operation.scales(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "fragment_slice_k", + ) +} + +/// Load the A register fragment for the selected warp and K64 half. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load_a( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 4 { + return input_err!( + loc, + TranslationErr::unsupported("mma_load_a expects MMA, shared A, warp-M, and K-half") + ); + } + let mma = decode_mma_operand(body, &args[0], "mma_load_a MMA")?; + let a = decode_shared_operand(body, &args[1], "mma_load_a shared A")?; + if a.kind != SharedTensorKind::Data + || a.role != dialect_cute::attributes::CuteTensorRoleAttr::Mkl + || a.placement != mma.placement + { + return input_err!( + loc, + TranslationErr::unsupported( + "mma_load_a needs an M-role data tile with the MMA shared layout" + ) + ); + } + let (lane, after) = mma_lane(ctx, body, &args[0], block, prev, value_map, &loc)?; + let (base, capacity, after) = shared_carriers( + ctx, + body, + &args[1], + &a, + block, + Some(after), + value_map, + &loc, + "mma_load_a shared A", + )?; + let (warp_m, after) = translate_operand( + ctx, + body, + &args[2], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (half, after) = + translate_operand(ctx, body, &args[3], value_map, block, after, loc.clone())?; + let (view, _) = smem_view_type(ctx, &a); + let result_ty = destination_type(ctx, body, destination, "mma_load_a")?; + let operation = CuteMmaLoadAOp::new( + ctx, + lane, + base, + capacity, + warp_m, + half, + view, + result_ty, + plan(&mma), + ); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + finish_value( + ctx, + op, + operation.fragment(ctx), + destination, + target, + block, + value_map, + block_map, + loc, + "mma_load_a", + ) +} + +/// Construct the shared B view for a warp N tile and K64 half. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_partition_b( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 4 { + return input_err!( + loc, + TranslationErr::unsupported( + "mma_partition_b expects MMA, shared B, warp-N, and K-half" + ) + ); + } + let mma = decode_mma_operand(body, &args[0], "mma_partition_b MMA")?; + let b = decode_shared_operand(body, &args[1], "mma_partition_b shared B")?; + let result_facts = decode_b_tile(&destination_rust_type( + body, + destination, + "mma_partition_b", + )?) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid mma_partition_b result: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "mma_partition_b must return Mxf4BTileK64" + )) + })?; + if b.kind != SharedTensorKind::Data + || b.role != dialect_cute::attributes::CuteTensorRoleAttr::Nkl + || b.placement != mma.placement + || result_facts != mma + { + return input_err!( + loc, + TranslationErr::unsupported( + "mma_partition_b needs one N-role tile and matching MMA/B result layouts" + ) + ); + } + let (_lane, after) = mma_lane(ctx, body, &args[0], block, prev, value_map, &loc)?; + let (base, capacity, after) = shared_carriers( + ctx, + body, + &args[1], + &b, + block, + Some(after), + value_map, + &loc, + "mma_partition_b shared B", + )?; + let (warp_n, after) = translate_operand( + ctx, + body, + &args[2], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (half, after) = + translate_operand(ctx, body, &args[3], value_map, block, after, loc.clone())?; + let (view, _) = smem_view_type(ctx, &b); + let operation = CuteMmaPartitionBOp::new(ctx, base, capacity, warp_n, half, view, plan(&mma)); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + + let result_ty = destination_type(ctx, body, destination, "mma_partition_b")?; + let base_ty = struct_field_type(ctx, result_ty, 0, "Mxf4BTileK64")?; + let (base, anchor) = pointer_cast(ctx, operation.base(ctx), base_ty, block, op, &loc)?; + let (result, anchor) = build_struct_prefix( + ctx, + result_ty, + &[ + base, + operation.capacity(ctx), + operation.warp_n(ctx), + operation.k_half(ctx), + ], + block, + anchor, + &loc, + )?; + finish_value( + ctx, + anchor, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "mma_partition_b", + ) +} + +/// Emit the MXF4 MMA chain joining accumulator, A/B operands and matching scale fragments. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_tiled_gemm( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + _destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 5 { + return input_err!( + loc, + TranslationErr::unsupported( + "tiled_gemm expects accumulator, MMA, A, lazy B, and scales" + ) + ); + } + let acc_ty = operand_rust_type(body, &args[0], "tiled_gemm accumulator")?; + let mma = decode_mma_operand(body, &args[1], "tiled_gemm MMA")?; + let b_facts = decode_b_tile(&operand_rust_type(body, &args[3], "tiled_gemm lazy B")?) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid tiled_gemm lazy B: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported("tiled_gemm needs Mxf4BTileK64")) + })?; + if !is_accumulator(&acc_ty).unwrap_or(false) + || !is_scale_k64(&operand_rust_type(body, &args[4], "tiled_gemm scales")?).unwrap_or(false) + || b_facts != mma + { + return input_err!( + loc, + TranslationErr::unsupported( + "tiled_gemm accumulator, scales, or lazy-B layout is invalid" + ) + ); + } + + let (acc_pointer, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let acc_pointee = { + let ty = acc_pointer.get_type(ctx); + let ty = ty.deref(ctx); + let Some(pointer) = ty.downcast_ref::() else { + return input_err!( + loc, + TranslationErr::unsupported( + "tiled_gemm accumulator receiver must be a mutable reference" + ) + ); + }; + if !pointer.is_mutable { + return input_err!( + loc, + TranslationErr::unsupported("tiled_gemm accumulator receiver must be mutable") + ); + } + pointer.pointee + }; + let (accumulator, after) = + load_through(ctx, acc_pointer, acc_pointee, block, after.or(prev), &loc); + let (lane, after) = mma_lane(ctx, body, &args[1], block, Some(after), value_map, &loc)?; + let (a, after) = translate_operand( + ctx, + body, + &args[2], + value_map, + block, + Some(after), + loc.clone(), + )?; + let (b, after) = translate_operand(ctx, body, &args[3], value_map, block, after, loc.clone())?; + let (b_base, after) = aggregate_field(ctx, b, 0, "tiled_gemm lazy B", block, after, &loc)?; + let (b_capacity, after) = + aggregate_field(ctx, b, 1, "tiled_gemm lazy B", block, Some(after), &loc)?; + let (b_warp_n, after) = + aggregate_field(ctx, b, 2, "tiled_gemm lazy B", block, Some(after), &loc)?; + let (b_half, after) = + aggregate_field(ctx, b, 3, "tiled_gemm lazy B", block, Some(after), &loc)?; + let b_shared = SharedTensorRust { + kind: SharedTensorKind::Data, + role: dialect_cute::attributes::CuteTensorRoleAttr::Nkl, + placement: mma.placement.clone(), + }; + let (b_view, b_storage) = smem_view_type(ctx, &b_shared); + let (b_base, after) = recover_cta_shared_pointer( + ctx, + b_base, + b_storage, + "tiled_gemm lazy B base", + block, + after, + &loc, + )?; + let (scales, after) = translate_operand( + ctx, + body, + &args[4], + value_map, + block, + Some(after), + loc.clone(), + )?; + let operation = CuteTiledGemmOp::new( + ctx, + lane, + a, + b_base, + b_capacity, + b_warp_n, + b_half, + scales, + accumulator, + b_view, + plan(&mma), + ); + let op = operation.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after); + let store = Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![acc_pointer, operation.accumulator(ctx)], + vec![], + 0, + ); + store.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, store, block, Some(op)); + finish(ctx, store, target, block_map, loc, "tiled_gemm") +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs new file mode 100644 index 0000000000..4f9cbbeb56 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/smem_mma/mod.rs @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural Rust facts for the shared-memory MXFP4 MMA boundary. +//! +//! The runtime carriers stay ordinary Rust values. This module only decodes +//! the exact types that give those values tensor and tiled-MMA meaning. + +pub(crate) mod emit; + +use dialect_cute::attributes::CuteTensorRoleAttr; +use dialect_cute::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; +use rustc_public::ty::{FloatTy, GenericArgs, RigidTy, Ty, TyKind, UintTy}; + +use super::static_config::decode_smem_layout; +use super::tensor::{canonical_adt_path, exact_schema, type_arg}; + +/// Exact defining path of `cute_rs::copy::tiled::SharedTensor`; public reexports do not change it. +const SHARED_TENSOR: &str = "cute_rs::copy::tiled::SharedTensor"; +/// Exact defining path of `cute_rs::numeric::Mxf4E2M1`; public reexports do not change it. +const MXF4_E2M1: &str = "cute_rs::numeric::Mxf4E2M1"; +/// Exact defining path of `cute_rs::numeric::UE8M0x4`; public reexports do not change it. +const UE8M0_X4: &str = "cute_rs::numeric::UE8M0x4"; +/// Exact defining path of `cute_rs::block_scaled::layouts::Sm120ScaleAtom`; public reexports do not change it. +const SM120_SCALE_ATOM: &str = "cute_rs::block_scaled::layouts::Sm120ScaleAtom"; +/// Exact defining path of `cute_rs::block_scaled::Mkl`; public reexports do not change it. +const MKL: &str = "cute_rs::block_scaled::Mkl"; +/// Exact defining path of `cute_rs::block_scaled::Nkl`; public reexports do not change it. +const NKL: &str = "cute_rs::block_scaled::Nkl"; +/// Exact defining path of `cute_rs::mma::block_scaled::tiled::Mxf4TiledMma`; public reexports do not change it. +const TILED_MMA: &str = "cute_rs::mma::block_scaled::tiled::Mxf4TiledMma"; +/// Exact defining path of `cute_rs::mma::block_scaled::fragments::Mxf4AccumulatorTile2x8`; public reexports do not change it. +const ACCUMULATOR: &str = "cute_rs::mma::block_scaled::fragments::Mxf4AccumulatorTile2x8"; +/// Exact defining path of `cute_rs::mma::block_scaled::fragments::Mxf4ScaleTile128`; public reexports do not change it. +const SCALE_STAGE: &str = "cute_rs::mma::block_scaled::fragments::Mxf4ScaleTile128"; +/// Exact defining path of `cute_rs::mma::block_scaled::fragments::Mxf4ScalePairs128`; public reexports do not change it. +const SCALE_K64: &str = "cute_rs::mma::block_scaled::fragments::Mxf4ScalePairs128"; +/// Exact defining path of `cute_rs::mma::block_scaled::fragments::Mxf4BTileK64`; public reexports do not change it. +const B_TILE_K64: &str = "cute_rs::mma::block_scaled::fragments::Mxf4BTileK64"; + +/// Whether shared storage holds packed matrix data or scale words. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SharedTensorKind { + /// Packed FP4 matrix values carried in shared FP16 storage. + Data, + /// Packed UE8M0 scale words in the canonical shared atom. + Scale, +} + +/// Static kind, role and physical placement decoded from a shared tensor type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SharedTensorRust { + /// Distinguishes packed-data storage from packed scale words. + pub kind: SharedTensorKind, + /// Matrix-A or matrix-B role carried by the exact source marker. + pub role: CuteTensorRoleAttr, + /// Physical shared-memory placement, including composition and offset units. + pub placement: ComposedLayout, +} + +/// Shared placement retained by a tiled-MMA or B-tile source carrier. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TiledMmaRust { + /// Physical shared-memory placement, including composition and offset units. + pub placement: ComposedLayout, +} + +/// Return an ADT's exact defining path and generic arguments, or no match. +fn adt(ty: &Ty) -> Option<(String, GenericArgs)> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return None; + }; + Some((canonical_adt_path(&def), args)) +} + +/// Match the exact marker definition and require its expected empty generic schema. +fn exact_marker(ty: &Ty, expected: &str) -> Result { + let Some((path, args)) = adt(ty) else { + return Ok(false); + }; + if path != expected { + return Ok(false); + } + exact_schema(&args, &[], expected)?; + Ok(true) +} + +/// Decode the matrix-A or matrix-B role from its exact marker identity. +fn role(ty: &Ty) -> Result { + if exact_marker(ty, MKL)? { + Ok(CuteTensorRoleAttr::Mkl) + } else if exact_marker(ty, NKL)? { + Ok(CuteTensorRoleAttr::Nkl) + } else { + Err(format!("shared tensor role must be `{MKL}` or `{NKL}`")) + } +} + +/// Recognize the compiler's primitive FP16 type. +fn is_f16(ty: &Ty) -> bool { + matches!(ty.kind(), TyKind::RigidTy(RigidTy::Float(FloatTy::F16))) +} + +/// Recognize the compiler's primitive unsigned 32-bit type. +fn is_u32(ty: &Ty) -> bool { + matches!(ty.kind(), TyKind::RigidTy(RigidTy::Uint(UintTy::U32))) +} + +/// Return the fixed 32-by-4 word layout of one SM120 scale atom. +fn scale_placement() -> ComposedLayout { + let outer = "(32,4):(4,1)" + .parse() + .expect("fixed SM120 scale placement is valid"); + ComposedLayout::from_layout(outer, OffsetUnit::Elements) +} + +/// Return the fixed swizzled 128-by-32 packed-data placement. +pub(crate) fn canonical_data_placement() -> ComposedLayout { + let outer: Layout = "(128,32):(32,1)" + .parse() + .expect("fixed SM120 packed-data placement is valid"); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, outer, OffsetUnit::Elements) + .expect("fixed SM120 packed-data swizzle is valid") +} + +/// Decode packed data or scale tensor markers, preserving layout and matrix role. +pub(crate) fn decode_shared_tensor(ty: &Ty) -> Result, String> { + let Some((path, args)) = adt(ty) else { + return Ok(None); + }; + if path != SHARED_TENSOR { + return Ok(None); + } + exact_schema(&args, &["Type", "Type", "Type", "Type"], SHARED_TENSOR)?; + let element = type_arg(&args, 0, SHARED_TENSOR)?; + let storage = type_arg(&args, 1, SHARED_TENSOR)?; + let layout = type_arg(&args, 2, SHARED_TENSOR)?; + let role = role(&type_arg(&args, 3, SHARED_TENSOR)?)?; + + if exact_marker(&element, MXF4_E2M1)? && is_f16(&storage) { + return Ok(Some(SharedTensorRust { + kind: SharedTensorKind::Data, + role, + placement: decode_smem_layout(&layout)?, + })); + } + if exact_marker(&element, UE8M0_X4)? + && is_u32(&storage) + && exact_marker(&layout, SM120_SCALE_ATOM)? + { + return Ok(Some(SharedTensorRust { + kind: SharedTensorKind::Scale, + role, + placement: scale_placement(), + })); + } + Err("shared MMA tensor must be packed E2M1/f16 data or UE8M0x4/u32 scales".to_string()) +} + +/// Decode a value directly or through one reference or raw-pointer layer. +fn behind_one_pointer( + ty: &Ty, + decode: impl FnOnce(&Ty) -> Result, String>, +) -> Result, String> { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode(&pointee) + } + _ => decode(ty), + } +} + +/// Decode a shared tensor value or one borrowed/raw-pointer receiver. +pub(crate) fn decode_shared_tensor_receiver(ty: &Ty) -> Result, String> { + behind_one_pointer(ty, decode_shared_tensor) +} + +/// Recognize the tiled-MMA definition and decode its shared placement. +pub(crate) fn decode_tiled_mma(ty: &Ty) -> Result, String> { + let Some((path, args)) = adt(ty) else { + return Ok(None); + }; + if path != TILED_MMA { + return Ok(None); + } + exact_schema(&args, &["Type"], TILED_MMA)?; + Ok(Some(TiledMmaRust { + placement: decode_smem_layout(&type_arg(&args, 0, TILED_MMA)?)?, + })) +} + +/// Decode a tiled-MMA value or one borrowed/raw-pointer receiver. +pub(crate) fn decode_tiled_mma_receiver(ty: &Ty) -> Result, String> { + behind_one_pointer(ty, decode_tiled_mma) +} + +/// Recognize one exact source ADT rather than a type with a similar printed name. +fn exact_adt(ty: &Ty, expected: &str) -> Result { + exact_marker(ty, expected) +} + +/// Recognize the fixed accumulator carrier directly or behind one pointer. +pub(crate) fn is_accumulator(ty: &Ty) -> Result { + behind_one_pointer(ty, |ty| exact_adt(ty, ACCUMULATOR).map(Some)) + .map(|value| value.unwrap_or(false)) +} + +/// Recognize the complete 128-row scale-tile carrier. +pub(crate) fn is_scale_stage(ty: &Ty) -> Result { + exact_adt(ty, SCALE_STAGE) +} + +/// Recognize the scale-pair carrier selected for one K64 half. +pub(crate) fn is_scale_k64(ty: &Ty) -> Result { + exact_adt(ty, SCALE_K64) +} + +/// Decode the B-tile lifetime/layout schema and recover its shared placement. +pub(crate) fn decode_b_tile(ty: &Ty) -> Result, String> { + let Some((path, args)) = adt(ty) else { + return Ok(None); + }; + if path != B_TILE_K64 { + return Ok(None); + } + exact_schema(&args, &["Lifetime", "Type"], B_TILE_K64)?; + Ok(Some(TiledMmaRust { + placement: decode_smem_layout(&type_arg(&args, 1, B_TILE_K64)?)?, + })) +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check that scale placement uses exactly one 128-word atom. + #[test] + fn fixed_scale_placement_has_one_128_word_atom() { + assert_eq!(scale_placement().outer().checked_size(), Some(128)); + } + + /// Check that data placement fixes the canonical swizzled 128-by-32 tile. + #[test] + fn canonical_data_placement_locks_the_128_by_32_swizzled_tile() { + let placement = canonical_data_placement(); + assert_eq!(placement.outer().checked_size(), Some(128 * 32)); + assert_eq!(placement.inner(), Swizzle::new(2, 3, 3)); + assert_eq!(placement.offset(), 0); + assert_eq!(placement.unit(), OffsetUnit::Elements); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs new file mode 100644 index 0000000000..6f4331382c --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/static_config.rs @@ -0,0 +1,710 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural decoder for nested `cute-rs` type markers. +//! +//! The decoder follows ADT identities and generic-argument kinds: +//! +//! ```text +//! Layout<(C<32>, C<8>), (C<8>, C<1>)> +//! │ └─── shape ──┘ └─ stride ──┘ +//! └─ exact DefId: cute_rs::markers::Layout +//! ``` +//! +//! Shape and stride modes are Rust tuples of `C` leaves. A tuple type has +//! no defining path of its own, so the identity check holds at the leaves. +//! +//! It never searches a debug-printed type name. A near-miss from another +//! crate, a swapped `Type`/`Const`, or an unevaluated constant is a loud +//! boundary error. + +use dialect_cute::layout::{ + ComposedLayout, CooperativeCopyPlan, IntTuple, Layout, OffsetUnit, Swizzle, + validate_cooperative_copy_plan, validate_ldmatrix_source, +}; +use rustc_public::CrateDef; +use rustc_public::mir; +use rustc_public::ty::{ + AdtDef, FloatTy, GenericArgKind, GenericArgs, IntTy, RigidTy, Ty, TyConst, TyConstKind, TyKind, + UintTy, +}; + +/// Exact defining path of `cute_rs::markers::C`; public reexports do not change it. +const C: &str = "cute_rs::markers::C"; +/// Exact defining path of `cute_rs::markers::Layout`; public reexports do not change it. +const LAYOUT: &str = "cute_rs::markers::Layout"; +/// Exact defining path of `cute_rs::markers::RowMajor`; public reexports do not change it. +const ROW_MAJOR: &str = "cute_rs::markers::RowMajor"; +/// Exact defining path of `cute_rs::markers::ColMajor`; public reexports do not change it. +const COL_MAJOR: &str = "cute_rs::markers::ColMajor"; +/// Exact defining path of `cute_rs::markers::CopyG2S`; public reexports do not change it. +const COPY_G2S: &str = "cute_rs::markers::CopyG2S"; +/// Exact defining path of `cute_rs::markers::LeadingDim`; public reexports do not change it. +const LEADING_DIM: &str = "cute_rs::markers::LeadingDim"; +/// Exact defining path of `cute_rs::markers::Swizzle`; public reexports do not change it. +const SWIZZLE: &str = "cute_rs::markers::Swizzle"; +/// Exact defining path of `cute_rs::markers::ComposedLayout`; public reexports do not change it. +const COMPOSED: &str = "cute_rs::markers::ComposedLayout"; + +/// Validated cooperative-copy instruction width. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CopyAtom { + /// Bytes transferred by one copy atom. + pub bytes: u32, +} + +/// Decoded copy types and layouts together with the checked cooperative transfer plan. +#[derive(Debug, Clone)] +pub(crate) struct CooperativeConfig { + /// Validated asynchronous-copy instruction selection. + pub atom: CopyAtom, + /// Map from participating thread coordinates to linear thread IDs. + pub thread_layout: Layout, + /// Map describing the per-thread values in one cooperative transfer. + pub value_layout: Layout, + /// Logical global tile shape and coordinate mapping. + pub tile_layout: Layout, + /// Physical destination layout, including its shared-memory swizzle. + pub smem_layout: ComposedLayout, + /// Divisor of the runtime row pitch, measured in elements. + pub leading_dim_divisor: u64, + /// The exact checked maps M3 must lower; never recompute them later. + #[allow(dead_code)] + pub plan: CooperativeCopyPlan, + // M3 emission needs the decoded Rust type; the pre-M3 sentinel only prints it. + /// Decoded Rust scalar type retained for IR carrier translation. + #[allow(dead_code)] + pub element_type: Ty, + // Keep the validated byte width beside the type so emission cannot re-guess it. + /// Validated number of storage bytes per scalar element. + #[allow(dead_code)] + pub element_bytes: i64, +} + +/// Return the supported scalar element's byte width and diagnostic name. +fn cooperative_element(ty: &Ty) -> Result<(i64, &'static str), String> { + let info = match ty.kind() { + TyKind::RigidTy(RigidTy::Int(kind)) => match kind { + IntTy::I8 => (1, "i8"), + IntTy::I16 => (2, "i16"), + IntTy::I32 => (4, "i32"), + IntTy::I64 => (8, "i64"), + IntTy::I128 => (16, "i128"), + IntTy::Isize => (8, "isize"), + }, + TyKind::RigidTy(RigidTy::Uint(kind)) => match kind { + UintTy::U8 => (1, "u8"), + UintTy::U16 => (2, "u16"), + UintTy::U32 => (4, "u32"), + UintTy::U64 => (8, "u64"), + UintTy::U128 => (16, "u128"), + UintTy::Usize => (8, "usize"), + }, + TyKind::RigidTy(RigidTy::Float(kind)) => match kind { + FloatTy::F16 => (2, "f16"), + FloatTy::F32 => (4, "f32"), + FloatTy::F64 => (8, "f64"), + FloatTy::F128 => { + return Err("copy_g2s does not support f128 elements".to_string()); + } + }, + // Rust `char` already lowers exactly to an unsigned 32-bit dialect integer. + TyKind::RigidTy(RigidTy::Char) => (4, "char"), + _ => { + return Err(format!( + "copy_g2s element must be a supported scalar, found {ty:?}" + )); + } + }; + Ok(info) +} + +/// Recover an ADT's actual defining path, independent of public aliases. +fn canonical_adt_path(def: &AdtDef) -> String { + super::identity::canonical_def_path(def.def_id(), &def.krate().name.to_string()) +} + +/// Require an exact ADT owner and return its generic arguments. +fn adt_owned(ty: &Ty, expected: &str) -> Result { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Err(format!("expected `{expected}`, found non-ADT type {ty:?}")); + }; + let found = canonical_adt_path(&def); + if found != expected { + return Err(format!("expected `{expected}`, found `{found}`")); + } + Ok(args) +} + +/// Check the ordered Type/Const/Lifetime kinds before decoding generic arguments. +fn expect_schema(args: &GenericArgs, expected: &[&str], owner: &str) -> Result<(), String> { + let found: Vec<_> = args + .0 + .iter() + .map(|arg| match arg { + GenericArgKind::Type(_) => "Type", + GenericArgKind::Const(_) => "Const", + GenericArgKind::Lifetime(_) => "Lifetime", + }) + .collect(); + if found == expected { + Ok(()) + } else { + Err(format!( + "`{owner}` expects generics [{}], found [{}]", + expected.join(", "), + found.join(", ") + )) + } +} + +/// Evaluate a source constant and reject values that do not fit in `u64`. +fn unsigned_const(value: &TyConst, what: &str) -> Result { + let raw = match value.kind() { + TyConstKind::Value(_, allocation) => allocation + .read_uint() + .map_err(|error| format!("cannot read {what}: {error:?}"))?, + _ => u128::from( + value + .eval_target_usize() + .map_err(|error| format!("cannot evaluate {what}: {error:?}"))?, + ), + }; + u64::try_from(raw).map_err(|_| format!("{what} value {raw} does not fit in u64")) +} + +/// Evaluate a signed source constant and reject values outside `i64`. +fn signed_const(value: &TyConst, what: &str) -> Result { + let raw = match value.kind() { + TyConstKind::Value(_, allocation) => allocation + .read_int() + .map_err(|error| format!("cannot read {what}: {error:?}"))?, + _ => i128::from( + value + .eval_target_usize() + .map_err(|error| format!("cannot evaluate {what}: {error:?}"))?, + ), + }; + i64::try_from(raw).map_err(|_| format!("{what} value {raw} does not fit in i64")) +} + +/// Read a type argument at the required generic position, rejecting other argument kinds. +fn type_arg(args: &GenericArgs, index: usize, owner: &str) -> Result { + match args.0.get(index) { + Some(GenericArgKind::Type(value)) => Ok(*value), + _ => Err(format!("`{owner}` generic {index} must be a type")), + } +} + +/// Read a constant argument at the required generic position, rejecting other argument kinds. +fn const_arg(args: &GenericArgs, index: usize, owner: &str) -> Result { + match args.0.get(index) { + Some(GenericArgKind::Const(value)) => Ok(value.clone()), + _ => Err(format!("`{owner}` generic {index} must be a const")), + } +} + +/// Recursively decode integer leaves and tuple modes into a layout tuple. +/// +/// A mode is the leaf marker `C` or a Rust tuple `(A, B, ...)` of further +/// modes, so `(C<2>, (C<3>, C<4>))` decodes to `(2,(3,4))`. A tuple type has +/// no defining path, so the near-miss guard applies at the leaves: every leaf +/// must be the exact `cute_rs::markers::C` ADT. +fn decode_tuple(ty: &Ty) -> Result { + if let TyKind::RigidTy(RigidTy::Tuple(items)) = ty.kind() { + if items.is_empty() { + return Err("layout tuple must have at least one mode".to_string()); + } + return items + .iter() + .map(decode_tuple) + .collect::, _>>() + .map(IntTuple::Tuple); + } + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Err(format!( + "layout tuple mode must be a tuple or `{C}`, found {ty:?}" + )); + }; + match canonical_adt_path(&def).as_str() { + C => { + expect_schema(&args, &["Const"], C)?; + Ok(IntTuple::Leaf(signed_const( + &const_arg(&args, 0, C)?, + "C value", + )?)) + } + found => Err(format!("expected a tuple or `{C}`, found `{found}`")), + } +} + +/// Compute the shape product with checked arithmetic and positive extents. +fn checked_shape_size(shape: &IntTuple) -> Result { + match shape { + IntTuple::Leaf(value) if *value > 0 => Ok(*value), + IntTuple::Leaf(value) => Err(format!("layout shape leaves must be positive, got {value}")), + IntTuple::Tuple(items) => items.iter().try_fold(1_i64, |size, item| { + size.checked_mul(checked_shape_size(item)?) + .ok_or_else(|| "layout shape product overflows i64".to_string()) + }), + } +} + +/// Decode a static layout marker, including nested shape and stride tuples. +pub(crate) fn decode_layout(ty: &Ty) -> Result { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Err(format!("layout marker must be an ADT, found {ty:?}")); + }; + let path = canonical_adt_path(&def); + let layout = match path.as_str() { + LAYOUT => { + expect_schema(&args, &["Type", "Type"], LAYOUT)?; + let shape = decode_tuple(&type_arg(&args, 0, LAYOUT)?)?; + let stride = decode_tuple(&type_arg(&args, 1, LAYOUT)?)?; + if !shape.is_congruent(&stride) { + return Err("layout shape and stride have different nesting".to_string()); + } + Layout::new(shape, stride) + } + ROW_MAJOR | COL_MAJOR => { + expect_schema(&args, &["Const", "Const"], &path)?; + let rows = signed_const(&const_arg(&args, 0, &path)?, "row count")?; + let cols = signed_const(&const_arg(&args, 1, &path)?, "column count")?; + if rows <= 0 || cols <= 0 { + return Err(format!("`{path}` dimensions must be positive")); + } + let shape = IntTuple::Tuple(vec![IntTuple::Leaf(rows), IntTuple::Leaf(cols)]); + let stride = if path == ROW_MAJOR { + IntTuple::Tuple(vec![IntTuple::Leaf(cols), IntTuple::Leaf(1)]) + } else { + IntTuple::Tuple(vec![IntTuple::Leaf(1), IntTuple::Leaf(rows)]) + }; + Layout::new(shape, stride) + } + _ => return Err(format!("expected a cute-rs layout marker, found `{path}`")), + }; + checked_shape_size(&layout.shape)?; + Ok(layout) +} + +/// Decode and validate a `CopyG2S` copy atom's byte width. +fn decode_atom(ty: &Ty) -> Result { + let args = adt_owned(ty, COPY_G2S)?; + expect_schema(&args, &["Const"], COPY_G2S)?; + let bytes = u32::try_from(unsigned_const( + &const_arg(&args, 0, COPY_G2S)?, + "cp.async width", + )?) + .map_err(|_| "cp.async width does not fit in u32".to_string())?; + if !matches!(bytes, 4 | 8 | 16) { + return Err(format!( + "cp.async width must be 4, 8, or 16 bytes, got {bytes}" + )); + } + Ok(CopyAtom { bytes }) +} + +/// Decode the guaranteed runtime row-pitch divisor from its static marker. +fn decode_leading_dim_divisor(ty: &Ty) -> Result { + let args = adt_owned(ty, LEADING_DIM)?; + expect_schema(&args, &["Const"], LEADING_DIM)?; + let divisor = unsigned_const( + &const_arg(&args, 0, LEADING_DIM)?, + "leading-dimension divisor", + )?; + if divisor == 0 { + return Err("leading-dimension divisor must be greater than zero".to_string()); + } + Ok(divisor) +} + +/// Decode XOR-swizzle bits, base and shift from constant marker arguments. +fn decode_swizzle(ty: &Ty) -> Result { + let args = adt_owned(ty, SWIZZLE)?; + expect_schema(&args, &["Const", "Const", "Const"], SWIZZLE)?; + let bits = u32::try_from(unsigned_const(&const_arg(&args, 0, SWIZZLE)?, "swizzle B")?) + .map_err(|_| "swizzle B does not fit in u32".to_string())?; + let base = u32::try_from(unsigned_const(&const_arg(&args, 1, SWIZZLE)?, "swizzle M")?) + .map_err(|_| "swizzle M does not fit in u32".to_string())?; + let shift = i32::try_from(signed_const(&const_arg(&args, 2, SWIZZLE)?, "swizzle S")?) + .map_err(|_| "swizzle S does not fit in i32".to_string())?; + Swizzle::try_new(bits, base, shift).ok_or_else(|| { + format!("invalid CuTe swizzle S<{bits},{base},{shift}>: require |S| >= B and M+B+|S| <= 32") + }) +} + +/// Decode the sealed attention layout identity and its major mode. +pub(crate) fn decode_attention_layout(ty: &Ty) -> Result<(ComposedLayout, bool), String> { + /// Exact defining path of `cute_rs::attention::sm100::shared::SharedLayout`; public reexports do not change it. + const PATH: &str = "cute_rs::attention::sm100::shared::SharedLayout"; + let args = adt_owned(ty, PATH)?; + expect_schema(&args, &["Const", "Const", "Const"], PATH)?; + let rows = u32::try_from(unsigned_const( + &const_arg(&args, 0, PATH)?, + "attention rows", + )?) + .map_err(|_| "attention rows exceed u32")?; + let cols = u32::try_from(unsigned_const( + &const_arg(&args, 1, PATH)?, + "attention columns", + )?) + .map_err(|_| "attention columns exceed u32")?; + let major = unsigned_const(&const_arg(&args, 2, PATH)?, "attention major mode")?; + if major > 1 { + return Err("attention major mode must be boolean".into()); + } + let mn = major == 1; + Ok(( + dialect_cute::sm100_attention_ops::shared_layout(rows, cols, mn)?.0, + mn, + )) +} + +/// Decode the supported plain or composed shared-memory layout markers. +pub(crate) fn decode_smem_layout(ty: &Ty) -> Result { + if let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() + && canonical_adt_path(&def) == COMPOSED + { + // `ComposedLayout`: position 0 is the inner swizzle and + // position 2 the outer layout, in the order of CuTeDSL's + // `make_composed_layout(inner, offset, outer)`. + expect_schema(&args, &["Type", "Const", "Type"], COMPOSED)?; + let swizzle = decode_swizzle(&type_arg(&args, 0, COMPOSED)?)?; + let offset = signed_const(&const_arg(&args, 1, COMPOSED)?, "composed offset")?; + let layout = decode_layout(&type_arg(&args, 2, COMPOSED)?)?; + // Rust marker layouts count elements. Keeping the unit explicit here + // lets the dialect verifier recast the entire expression—including + // the swizzle bit fields—to bytes once `T` is known. + return ComposedLayout::new(swizzle, offset, layout, OffsetUnit::Elements) + .map_err(|error| error.to_string()); + } + Ok(ComposedLayout::from_layout( + decode_layout(ty)?, + OffsetUnit::Elements, + )) +} + +/// Decode the six explicit parameters and inferred pitch marker at the +/// frozen `copy_g2s` boundary. +pub(crate) fn decode_cooperative_config(args: &GenericArgs) -> Result { + expect_schema( + args, + &["Type", "Type", "Type", "Type", "Type", "Type", "Type"], + "cute_rs::copy::cooperative::copy_g2s", + )?; + let atom = decode_atom(&type_arg(args, 0, "copy_g2s")?)?; + let thread_layout = decode_layout(&type_arg(args, 1, "copy_g2s")?)?; + let value_layout = decode_layout(&type_arg(args, 2, "copy_g2s")?)?; + let tile_layout = decode_layout(&type_arg(args, 3, "copy_g2s")?)?; + let smem_layout = decode_smem_layout(&type_arg(args, 4, "copy_g2s")?)?; + let element_type = type_arg(args, 5, "copy_g2s")?; + let (element_bytes, _spelling) = cooperative_element(&element_type)?; + let leading_dim_divisor = decode_leading_dim_divisor(&type_arg(args, 6, "copy_g2s")?)?; + let leading_dim_alignment = leading_dim_divisor + .checked_mul(element_bytes as u64) + .ok_or_else(|| "leading-dimension byte alignment overflows u64".to_string())?; + if leading_dim_alignment % u64::from(atom.bytes) != 0 { + return Err(format!( + "leading-dimension promise is too weak: {leading_dim_divisor} elements x \ + {element_bytes} bytes is not a multiple of the {}-byte atom", + atom.bytes + )); + } + + let plan = validate_cooperative_copy_plan( + atom.bytes, + &thread_layout, + &value_layout, + &tile_layout, + &smem_layout, + element_bytes, + ) + .map_err(|error| format!("invalid cooperative copy plan: {error}"))?; + + Ok(CooperativeConfig { + atom, + thread_layout, + value_layout, + tile_layout, + smem_layout, + leading_dim_divisor, + plan, + element_type, + element_bytes, + }) +} + +/// Reach the static decoder from a direct call operand. +pub(crate) fn decode_copy_g2s_callee(func: &mir::Operand) -> Result { + let mir::Operand::Constant(constant) = func else { + return Err("copy_g2s must be a direct constant function call".to_string()); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err("copy_g2s callee is not a FnDef".to_string()); + }; + decode_cooperative_config(&args) +} + +/// Checked static facts for one warp-cooperative `ldmatrix` fragment load. +#[derive(Debug, Clone)] +pub(crate) struct MatrixLoadConfig { + /// Shared tile map, in elements; backend lowering recasts it to bytes. + pub smem_layout: ComposedLayout, +} + +/// Decode `load_matrix_{a,b}::`'s one generic and prove the shared +/// tile can feed `ldmatrix` (two modes, 8-multiple extents, 16-byte-aligned +/// contiguous row segments through any swizzle). +pub(crate) fn decode_load_matrix_callee(func: &mir::Operand) -> Result { + let mir::Operand::Constant(constant) = func else { + return Err("load_matrix must be a direct constant function call".to_string()); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err("load_matrix callee is not a FnDef".to_string()); + }; + expect_schema(&args, &["Type"], "cute_rs::mma::load_matrix")?; + let smem_layout = decode_smem_layout(&type_arg(&args, 0, "load_matrix")?)?; + + let modes = smem_layout.outer().modes(); + if modes.len() != 2 { + return Err("load_matrix shared layout needs row and column modes".to_string()); + } + let (Some(rows), Some(columns)) = (modes[0].checked_size(), modes[1].checked_size()) else { + return Err("load_matrix shared layout extents are invalid".to_string()); + }; + let byte_layout = smem_layout + .to_byte_offsets(2) + .map_err(|error| format!("load_matrix shared layout has no byte form: {error}"))?; + validate_ldmatrix_source(&byte_layout, rows, columns) + .map_err(|error| format!("load_matrix source is not ldmatrix-loadable: {error}"))?; + + Ok(MatrixLoadConfig { smem_layout }) +} + +/// Checked static facts for one hardware (TMA) tile copy. +#[derive(Debug, Clone)] +pub(crate) struct TmaCopyConfig { + /// Physical destination layout, including its shared-memory swizzle. + pub smem_layout: ComposedLayout, + /// Decoded Rust scalar type retained for IR carrier translation. + pub element_type: Ty, + /// Minimum shared-base alignment that preserves this layout's swizzle + /// phase. The recognized unsafe boundary is the exact place where the + /// source program promises it. + pub smem_alignment_bytes: u64, +} + +/// Compute the base alignment needed to preserve the selected TMA swizzle phase. +fn tma_copy_alignment_bytes( + smem_layout: &ComposedLayout, + element_bytes: i64, +) -> Result { + let alignment = cute_layout::tma_phase_alignment_bytes(smem_layout, element_bytes) + .map_err(|error| format!("TMA layout has no valid shared-base alignment: {error}"))?; + u64::try_from(alignment).map_err(|_| "TMA shared-base alignment does not fit u64".to_string()) +} + +/// Decode `copy_tma_2d::` and prove `SmemL` is TMA-encodable at +/// kernel compile time (the host encoder re-checks at descriptor creation). +pub(crate) fn decode_copy_tma_callee(func: &mir::Operand) -> Result { + let mir::Operand::Constant(constant) = func else { + return Err("copy_tma_2d must be a direct constant function call".to_string()); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err("copy_tma_2d callee is not a FnDef".to_string()); + }; + expect_schema(&args, &["Type", "Type"], "cute_rs::copy::tma::copy_tma_2d")?; + let element_type = type_arg(&args, 0, "copy_tma_2d")?; + let (element_bytes, _) = cooperative_element(&element_type)?; + let smem_layout = decode_smem_layout(&type_arg(&args, 1, "copy_tma_2d")?)?; + let smem_alignment_bytes = tma_copy_alignment_bytes(&smem_layout, element_bytes) + .map_err(|error| format!("copy_tma_2d layout is not TMA-encodable: {error}"))?; + Ok(TmaCopyConfig { + smem_layout, + element_type, + smem_alignment_bytes, + }) +} + +/// Decode `copy_tma_s2g_2d::` and prove the shared source layout +/// is TMA-encodable before constructing the store operation. +pub(crate) fn decode_copy_tma_s2g_callee(func: &mir::Operand) -> Result { + let mir::Operand::Constant(constant) = func else { + return Err("copy_tma_s2g_2d must be a direct constant function call".to_string()); + }; + let TyKind::RigidTy(RigidTy::FnDef(_, args)) = constant.const_.ty().kind() else { + return Err("copy_tma_s2g_2d callee is not a FnDef".to_string()); + }; + expect_schema( + &args, + &["Type", "Type"], + "cute_rs::copy::tma::copy_tma_s2g_2d", + )?; + let element_type = type_arg(&args, 0, "copy_tma_s2g_2d")?; + let (element_bytes, _) = cooperative_element(&element_type)?; + let smem_layout = decode_smem_layout(&type_arg(&args, 1, "copy_tma_s2g_2d")?)?; + let smem_alignment_bytes = tma_copy_alignment_bytes(&smem_layout, element_bytes) + .map_err(|error| format!("copy_tma_s2g_2d layout is not TMA-encodable: {error}"))?; + Ok(TmaCopyConfig { + smem_layout, + element_type, + smem_alignment_bytes, + }) +} + +/// Regression tests for exact source identity and semantic type preservation. +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + + /// Compile `source` in-process as a crate named `cute_rs`, so a `markers` + /// module inside it sits at the defining paths the decoder pins, then + /// decode the first parameter type of every function item and return the + /// results by short function name. + /// + /// `Ty` values exist only inside a compiler session, so the decoded + /// layouts are plain data compared outside it. The rustup shim resolves + /// the pinned toolchain this test binary was built with, so the in-process + /// driver and the sysroot agree. + fn decode_fixture_parameters( + source: &str, + tag: &str, + ) -> BTreeMap> { + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system time before unix epoch") + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "cuda_oxide_static_config_{tag}_{}_{unique}", + std::process::id() + )); + std::fs::create_dir_all(&root).unwrap(); + let fixture = root.join("fixture.rs"); + std::fs::write(&fixture, source).unwrap(); + + let rustc = std::env::var_os("RUSTC").unwrap_or_else(|| "rustc".into()); + let sysroot_output = std::process::Command::new(rustc) + .args(["--print", "sysroot"]) + .output() + .expect("query rustc sysroot"); + assert!(sysroot_output.status.success(), "rustc --print sysroot"); + let sysroot = String::from_utf8(sysroot_output.stdout) + .expect("sysroot path is UTF-8") + .trim() + .to_string(); + + let args = vec![ + "rustc".to_string(), + "--edition=2024".to_string(), + "--crate-type=rlib".to_string(), + "--crate-name=cute_rs".to_string(), + "--emit=metadata".to_string(), + format!("--out-dir={}", root.display()), + format!("--sysroot={sysroot}"), + fixture.display().to_string(), + ]; + + // rustc needs more stack than the default test-thread allowance. + let decoded = std::thread::Builder::new() + .stack_size(64 * 1024 * 1024) + .spawn(move || { + rustc_public::run!(&args, || { + let mut decoded = BTreeMap::new(); + for item in rustc_public::all_local_items() { + if item.kind() != rustc_public::ItemKind::Fn { + continue; + } + let Some(body) = item.body() else { + continue; + }; + let Some(parameter) = body.locals().get(1) else { + continue; + }; + let name = item.name(); + let short = name.rsplit("::").next().unwrap_or(&name).to_string(); + decoded.insert(short, decode_layout(¶meter.ty)); + } + std::ops::ControlFlow::<(), _>::Continue(decoded) + }) + }) + .unwrap() + .join() + .unwrap() + .expect("in-process fixture compilation succeeds"); + + std::fs::remove_dir_all(&root).ok(); + decoded + } + + /// Rust tuples carry nested shapes and strides: a tuple of `C` leaves + /// or of further tuples decodes to the nested layout it spells, an empty + /// tuple is rejected, and a leaf that is not the exact + /// `cute_rs::markers::C` ADT fails the identity check. + #[test] + fn tuple_shapes_decode_to_nested_layouts() { + let decoded = decode_fixture_parameters( + r#" +pub mod markers { + use core::marker::PhantomData; + pub struct C; + pub struct Layout { _layout: PhantomData<(Shape, Stride)> } +} +pub mod other { + pub struct C; +} +use markers::{C, Layout}; +pub fn leaf(_: Layout, C<17>>) {} +pub fn pitched(_: Layout<(C<2>, C<3>), (C<17>, C<4>)>) {} +pub fn nested(_: Layout<(C<4>, ((C<2>, C<2>), C<2>)), (C<2>, ((C<1>, C<128>), C<8>))>) {} +pub fn triple(_: Layout<(C<2>, C<3>, C<4>), (C<12>, C<4>, C<1>)>) {} +pub fn empty_mode(_: Layout<(), ()>) {} +pub fn primitive_leaf(_: Layout<(u32, C<1>), (C<1>, C<1>)>) {} +pub fn foreign_leaf(_: Layout<(other::C<2>, C<3>), (C<3>, C<1>)>) {} +"#, + "tuple_shapes", + ); + let expect = |name: &str, layout: &str| { + assert_eq!( + decoded[name], + Ok(layout.parse::().unwrap()), + "{name}" + ); + }; + expect("leaf", "2:17"); + expect("pitched", "(2,3):(17,4)"); + expect("nested", "(4,((2,2),2)):(2,((1,128),8))"); + expect("triple", "(2,3,4):(12,4,1)"); + let error = |name: &str| decoded[name].as_ref().expect_err(name).clone(); + assert!(error("empty_mode").contains("at least one mode")); + assert!(error("primitive_leaf").contains("must be a tuple or `cute_rs::markers::C`")); + assert_eq!( + error("foreign_leaf"), + "expected a tuple or `cute_rs::markers::C`, found `cute_rs::other::C`" + ); + } + + /// Check that TMA copy alignment preserves the selected swizzle phase. + #[test] + fn tma_copy_alignment_keeps_the_selected_swizzle_phase() { + let byte_layout = ComposedLayout::new( + Swizzle::new(2, 4, 3), + 0, + "(128,64):(64,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(); + let scale_layout = ComposedLayout::new( + Swizzle::IDENTITY, + 0, + "(1,256):(256,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(); + + assert_eq!(tma_copy_alignment_bytes(&byte_layout, 1), Ok(512)); + assert_eq!(tma_copy_alignment_bytes(&scale_layout, 2), Ok(16)); + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/emit.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/emit.rs new file mode 100644 index 0000000000..9e777315c9 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/emit.rs @@ -0,0 +1,827 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Emit the high-level elementwise tensor chain. +//! +//! Each recognized Rust call becomes one operation. Pointer arithmetic stays +//! out of the importer; the selected backend lowers it from the preserved view. + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::rvalue::translate_operand; +use crate::translator::terminator::helpers::emit_goto; +use crate::translator::types::translate_type; +use crate::translator::values::ValueMap; +use dialect_cute::attributes::CuteTensorAccessAttr; +use dialect_cute::tensor_ops::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, +}; +use dialect_cute::types::CuteTensorViewType; +use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr}; +use dialect_mir::ops::{MirCastOp, MirExtractFieldOp, MirLoadOp}; +use dialect_mir::types::{MirDisjointSliceType, MirPtrType, MirSliceType}; +use pliron::basic_block::BasicBlock; +use pliron::context::{Context, Ptr}; +use pliron::input_err; +use pliron::location::Location; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use rustc_public::mir; + +use super::super::tensor::{decode_register_tile, decode_tensor_view}; + +/// Insert an operation after the current anchor, or at the start of the block. +fn insert( + ctx: &mut Context, + op: Ptr, + block: Ptr, + prev: Option>, +) { + if let Some(prev) = prev { + op.insert_after(ctx, prev); + } else { + op.insert_at_front(block, ctx); + } +} + +/// Branch from the last emitted effect to the call continuation; reject diverging calls. +fn finish( + ctx: &mut Context, + effect: Ptr, + target: &Option, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + let Some(target) = target else { + return input_err!( + loc, + TranslationErr::unsupported(format!("{what} call without target is not supported")) + ); + }; + Ok(emit_goto(ctx, *target, effect, block_map, loc)) +} + +/// Store a result in the unprojected MIR destination, then emit the continuation branch. +#[allow(clippy::too_many_arguments)] +fn finish_value( + ctx: &mut Context, + producer: Ptr, + value: Value, + destination: &mir::Place, + target: &Option, + block: Ptr, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + what: &str, +) -> TranslationResult> { + if !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported(format!( + "{what} destination with projections is not supported" + )) + ); + } + let after_store = value_map + .store_local(ctx, destination.local, value, block, Some(producer)) + .unwrap_or(producer); + finish(ctx, after_store, target, block_map, loc, what) +} + +/// Emit one aggregate-field extraction and return its value and insertion anchor. +fn extract_field( + ctx: &mut Context, + aggregate: Value, + index: u32, + result_ty: TypeHandle, + block: Ptr, + prev: Option>, + loc: &Location, +) -> (Value, Ptr) { + let op = Operation::new( + ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![result_ty], + vec![aggregate], + vec![], + 0, + ); + op.deref_mut(ctx).set_loc(loc.clone()); + MirExtractFieldOp::new(op).set_attr_index(ctx, FieldIndexAttr(index)); + insert(ctx, op, block, prev); + (op.deref(ctx).get_result(0), op) +} + +/// Load the translated carrier behind a tensor reference or raw pointer. +fn load_pointer_pointee( + ctx: &mut Context, + pointer: Value, + expected: fn(&dyn pliron::r#type::Type) -> bool, + what: &str, + block: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, TypeHandle, Ptr)> { + let pointee = { + let pointer_ty = pointer.get_type(ctx); + let pointer_ty = pointer_ty.deref(ctx); + let Some(pointer) = pointer_ty.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} must be a MIR pointer")) + ); + }; + let pointee_ty = pointer.pointee.deref(ctx); + if !expected(&*pointee_ty) { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!("{what} has the wrong pointee type")) + ); + } + pointer.pointee + }; + let load = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![pointee], + vec![pointer], + vec![], + 0, + ); + load.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, load, block, prev); + Ok((load.deref(ctx).get_result(0), pointee, load)) +} + +/// Recover a tensor carrier from either a direct value or one pointer layer. +fn load_tensor_if_referenced( + ctx: &mut Context, + value: Value, + block: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Option>)> { + if value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some() + { + return Ok((value, prev)); + } + let (loaded, _, op) = load_pointer_pointee( + ctx, + value, + |ty| ty.downcast_ref::().is_some(), + "tensor receiver", + block, + prev, + loc, + )?; + Ok((loaded, Some(op))) +} + +/// Cast raw storage to its logical scalar element type at the typed tensor boundary. +fn cast_carrier_to_element( + ctx: &mut Context, + pointer: Value, + element: TypeHandle, + mutable: bool, + block: Ptr, + prev: Option>, + loc: &Location, +) -> TranslationResult<(Value, Ptr)> { + let (address_space, source_mutable) = { + let pointer_ty = pointer.get_type(ctx); + let pointer_ty = pointer_ty.deref(ctx); + let Some(pointer) = pointer_ty.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("register tile carrier must be a pointer".to_string()) + ); + }; + (pointer.address_space, pointer.is_mutable) + }; + if mutable && !source_mutable { + return input_err!( + loc.clone(), + TranslationErr::unsupported("load destination carrier must be mutable".to_string()) + ); + } + let result_ty: TypeHandle = MirPtrType::get(ctx, element, source_mutable, address_space).into(); + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![result_ty], + vec![pointer], + vec![], + 0, + ); + cast.deref_mut(ctx).set_loc(loc.clone()); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + insert(ctx, cast, block, prev); + Ok((cast.deref(ctx).get_result(0), cast)) +} + +/// Combine source and destination type facts for a tensor-view transfer. +fn view_transfer_facts( + ctx: &Context, + tensor: Value, + loc: &Location, +) -> TranslationResult<(TypeHandle, u64, u64)> { + let view = tensor.get_type(ctx); + let view = view.deref(ctx); + let Some(view) = view.downcast_ref::() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("tensor transfer operand is not a tensor view".to_string()) + ); + }; + let Some(tile) = view.selected_tile_size() else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("tensor transfer needs a selected tile".to_string()) + ); + }; + let Some(storage_bytes) = view.storage_bytes(ctx) else { + return input_err!( + loc.clone(), + TranslationErr::unsupported("tensor storage width is not supported".to_string()) + ); + }; + let alignment = tile.checked_mul(storage_bytes).ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tensor transfer byte width overflows u64".to_string() + )) + })?; + Ok((view.storage, tile, alignment)) +} + +/// Import a slice-backed read-only or writable contiguous tensor view. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_make( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + expected_access: CuteTensorAccessAttr, +) -> TranslationResult> { + let destination_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read tensor_make destination type: {error:?}" + ))) + })?; + let decoded = decode_tensor_view(&destination_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid tensor_make destination: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tensor_make destination is not a v0 tensor view".to_string() + )) + })?; + if decoded.access != expected_access { + return input_err!( + loc, + TranslationErr::unsupported("tensor_make access does not match its API".to_string()) + ); + } + let logical = translate_type(ctx, &decoded.logical)?; + let storage = translate_type(ctx, &decoded.storage)?; + let alignment = decoded + .storage + .layout() + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read tensor storage alignment: {error:?}" + ))) + })? + .shape() + .abi_align; + + let (data, len, after_inputs) = if expected_access == CuteTensorAccessAttr::ReadOnly { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("read tensor_make expects one slice".to_string()) + ); + } + let (slice, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + if slice + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + return input_err!( + loc, + TranslationErr::unsupported("read tensor_make input is not a slice".to_string()) + ); + } + let data_ty: TypeHandle = MirPtrType::get_generic(ctx, storage, false).into(); + let (data, data_op) = extract_field(ctx, slice, 0, data_ty, block, after, &loc); + let len_ty: TypeHandle = crate::translator::types::get_usize_type(ctx).into(); + let (len, len_op) = extract_field(ctx, slice, 1, len_ty, block, Some(data_op), &loc); + (data, len, Some(len_op)) + } else if args.len() == 2 { + // Stable free boundary: `(raw pointer, length)`. + let (data, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (len, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + (data, len, after) + } else if args.len() == 1 { + // Ergonomic method boundary: `&mut DisjointSlice`. + let (receiver, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (slice, pointee, load) = load_pointer_pointee( + ctx, + receiver, + |ty| ty.downcast_ref::().is_some(), + "TensorMut::from_disjoint_slice receiver", + block, + after, + &loc, + )?; + let element = pointee + .deref(ctx) + .downcast_ref::() + .expect("checked above") + .element_type(); + if element != storage { + return input_err!( + loc, + TranslationErr::unsupported( + "DisjointSlice element does not match tensor storage".to_string() + ) + ); + } + let data_ty: TypeHandle = MirPtrType::get_generic(ctx, storage, true).into(); + let (data, data_op) = extract_field(ctx, slice, 0, data_ty, block, Some(load), &loc); + let len_ty: TypeHandle = crate::translator::types::get_usize_type(ctx).into(); + let (len, len_op) = extract_field(ctx, slice, 1, len_ty, block, Some(data_op), &loc); + (data, len, Some(len_op)) + } else { + return input_err!( + loc, + TranslationErr::unsupported( + "writable tensor_make expects a DisjointSlice or raw pointer and length" + .to_string() + ) + ); + }; + + let make = CuteTensorMakeOp::new(ctx, data, len, logical, storage, expected_access, alignment); + let op = make.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after_inputs.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tensor_make", + ) +} + +/// Read the tile width from a divided or selected tensor-view destination. +/// +/// `what` names the Rust call in diagnostics (`zipped_divide` or `local_tile`). +fn destination_tile_size( + body: &mir::Body, + destination: &mir::Place, + what: &str, +) -> TranslationResult { + let destination_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read {what} destination type: {error:?}" + ))) + })?; + let result = decode_tensor_view(&destination_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid {what} destination: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} destination is not a v0 tensor view" + ))) + })?; + result.layout.tile_size().ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "{what} result does not carry a tile width" + ))) + }) +} + +/// Insert one `cute.tensor_zipped_divide` after `anchor`; return the op and its view. +fn insert_zipped_divide( + ctx: &mut Context, + tensor: Value, + tile_size: u64, + block: Ptr, + anchor: Option>, + loc: &Location, +) -> (Ptr, Value) { + let divide = CuteTensorZippedDivideOp::new(ctx, tensor, tile_size); + let op = divide.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, anchor); + let result = op.deref(ctx).get_result(0); + (op, result) +} + +/// Insert one `cute.tensor_slice` after `anchor`; return the op and its view. +fn insert_slice( + ctx: &mut Context, + tensor: Value, + tile_index: Value, + block: Ptr, + anchor: Option>, + loc: &Location, +) -> (Ptr, Value) { + let slice = CuteTensorSliceOp::new(ctx, tensor, tile_index); + let op = slice.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, anchor); + let result = op.deref(ctx).get_result(0); + (op, result) +} + +/// Split a contiguous tensor into tile-local and tile-index coordinates. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_zipped_divide( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("tensor zipped_divide expects one view".to_string()) + ); + } + let tile_size = destination_tile_size(body, destination, "zipped_divide")?; + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (op, result) = insert_zipped_divide(ctx, tensor, tile_size, block, after.or(prev), &loc); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tensor_zipped_divide", + ) +} + +/// Select one tile while retaining the tensor's original storage extent. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_slice( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported("tensor slice expects a view and tile index".to_string()) + ); + } + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (tile_index, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + let (op, result) = insert_slice(ctx, tensor, tile_index, block, after.or(prev), &loc); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tensor_slice", + ) +} + +/// Split a contiguous tensor into tiles and select one in a single call. +/// +/// This is CuTe's `local_tile` for a one-dimensional tiler and a single tile +/// coordinate. It inserts the same `cute.tensor_zipped_divide` followed by +/// `cute.tensor_slice` that the two-step `zipped_divide` then `slice` chain +/// produces, through the same builders, so the dialect verifier and the +/// CUTLASS compiler see one program for both forms. The tile width comes +/// from the `Tile1D` destination type; the tile index is the second operand. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_local_tile( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported( + "tensor local_tile expects a view and tile index".to_string() + ) + ); + } + let tile_size = destination_tile_size(body, destination, "local_tile")?; + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (divide, divided) = + insert_zipped_divide(ctx, tensor, tile_size, block, after.or(prev), &loc); + let (tile_index, after) = translate_operand( + ctx, + body, + &args[1], + value_map, + block, + Some(divide), + loc.clone(), + )?; + let (op, result) = insert_slice( + ctx, + divided, + tile_index, + block, + after.or(Some(divide)), + &loc, + ); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + "tensor_local_tile", + ) +} + +/// Import the full-tile or base-index query without reading tensor memory. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_query( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + is_full: bool, +) -> TranslationResult> { + if args.len() != 1 { + return input_err!( + loc, + TranslationErr::unsupported("tensor query expects one selected tile".to_string()) + ); + } + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let op = if is_full { + CuteTensorIsFullOp::new(ctx, tensor).get_operation() + } else { + CuteTensorBaseOp::new(ctx, tensor).get_operation() + }; + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + let result = op.deref(ctx).get_result(0); + finish_value( + ctx, + op, + result, + destination, + target, + block, + value_map, + block_map, + loc, + if is_full { + "tensor_is_full" + } else { + "tensor_base" + }, + ) +} + +/// Load a complete per-thread tile into a register carrier. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_load( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 1 || !destination.projection.is_empty() { + return input_err!( + loc, + TranslationErr::unsupported( + "tensor load expects one view and a local destination".to_string() + ) + ); + } + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (storage, tile_size, alignment) = view_transfer_facts(ctx, tensor, &loc)?; + let destination_ty = destination.ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read tensor load destination type: {error:?}" + ))) + })?; + let register = decode_register_tile(&destination_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid tensor load destination: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tensor load destination is not RegisterTile".to_string() + )) + })?; + if register.size != tile_size || translate_type(ctx, ®ister.element)? != storage { + return input_err!( + loc, + TranslationErr::unsupported( + "tensor load tile shape does not match RegisterTile".to_string() + ) + ); + } + let Some(slot) = value_map.get_slot(destination.local) else { + return input_err!( + loc, + TranslationErr::unsupported("tensor load destination has no local slot".to_string()) + ); + }; + let (destination, cast) = + cast_carrier_to_element(ctx, slot, storage, true, block, after, &loc)?; + let load = CuteTensorLoadIntoOp::new(ctx, tensor, destination, alignment); + let op = load.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(cast)); + finish(ctx, op, target, block_map, loc, "tensor_load_into") +} + +/// Store a register tile through its writable tensor view. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 2 { + return input_err!( + loc, + TranslationErr::unsupported("tensor store expects a view and RegisterTile".to_string()) + ); + } + let (tensor, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (storage, tile_size, alignment) = view_transfer_facts(ctx, tensor, &loc)?; + let value_ty = args[1].ty(body.locals()).map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read tensor store source type: {error:?}" + ))) + })?; + let register = decode_register_tile(&value_ty) + .map_err(|error| { + pliron::input_error_noloc!(TranslationErr::unsupported(format!( + "invalid tensor store source: {error}" + ))) + })? + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tensor store source is not RegisterTile".to_string() + )) + })?; + if register.size != tile_size || translate_type(ctx, ®ister.element)? != storage { + return input_err!( + loc, + TranslationErr::unsupported( + "tensor store tile shape does not match RegisterTile".to_string() + ) + ); + } + let source = match &args[1] { + mir::Operand::Copy(place) | mir::Operand::Move(place) if place.projection.is_empty() => { + value_map.get_slot(place.local) + } + _ => None, + } + .ok_or_else(|| { + pliron::input_error_noloc!(TranslationErr::unsupported( + "tensor store RegisterTile must live in a local slot".to_string() + )) + })?; + let (source, cast) = cast_carrier_to_element(ctx, source, storage, false, block, after, &loc)?; + let store = CuteTensorStoreFromOp::new(ctx, source, tensor, alignment); + let op = store.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, Some(cast)); + finish(ctx, op, target, block_map, loc, "tensor_store_from") +} + +/// Store one scalar element by absolute storage index for a tensor tail. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_store_element_abs( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block: Ptr, + prev: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + if args.len() != 3 { + return input_err!( + loc, + TranslationErr::unsupported( + "tensor tail store expects a view, absolute index, and value".to_string() + ) + ); + } + let (receiver, after) = + translate_operand(ctx, body, &args[0], value_map, block, prev, loc.clone())?; + let (tensor, after) = load_tensor_if_referenced(ctx, receiver, block, after, &loc)?; + let (index, after) = + translate_operand(ctx, body, &args[1], value_map, block, after, loc.clone())?; + let (value, after) = + translate_operand(ctx, body, &args[2], value_map, block, after, loc.clone())?; + let store = CuteTensorStoreElementAbsOp::new(ctx, tensor, index, value); + let op = store.get_operation(); + op.deref_mut(ctx).set_loc(loc.clone()); + insert(ctx, op, block, after.or(prev)); + finish(ctx, op, target, block_map, loc, "tensor_store_element_abs") +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs new file mode 100644 index 0000000000..fae1bf039e --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/cute/tensor/mod.rs @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural decoding for the first high-level tensor flow. +//! +//! The Rust types keep the story visible: +//! +//! ```text +//! Tensor<_, Contiguous1D> +//! -> Tensor<_, Zipped1D> +//! -> Tensor<_, Tile1D> +//! ``` +//! +//! This module turns only those exact `cute-rs` ADTs into the ghost +//! `!cute.tensor_view` type. Other `Tensor` layouts continue through the +//! ordinary Rust-struct path until their own high-level operation set exists. + +pub(crate) mod emit; + +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::types::translate_type; +use dialect_cute::attributes::{ + CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorLayoutAttr, +}; +use dialect_cute::types::CuteTensorViewType; +use pliron::context::Context; +use pliron::input_error_noloc; +use pliron::r#type::TypeHandle; +use rustc_public::CrateDef; +use rustc_public::ty::{ + AdtDef, GenericArgKind, GenericArgs, RigidTy, Ty, TyConst, TyConstKind, TyKind, +}; + +/// Exact defining path of `cute_rs::tensor::Tensor`; public reexports do not change it. +const TENSOR: &str = "cute_rs::tensor::Tensor"; +/// Exact defining path of `cute_rs::tensor::TensorMut`; public reexports do not change it. +const TENSOR_MUT: &str = "cute_rs::tensor::TensorMut"; +/// Exact defining path of `cute_rs::tensor::Contiguous1D`; public reexports do not change it. +const CONTIGUOUS_1D: &str = "cute_rs::tensor::Contiguous1D"; +/// Exact defining path of `cute_rs::tensor::Zipped1D`; public reexports do not change it. +const ZIPPED_1D: &str = "cute_rs::tensor::Zipped1D"; +/// Exact defining path of `cute_rs::tensor::Tile1D`; public reexports do not change it. +const TILE_1D: &str = "cute_rs::tensor::Tile1D"; +/// Exact defining path of `cute_rs::tensor::RegisterTile`; public reexports do not change it. +const REGISTER_TILE: &str = "cute_rs::tensor::RegisterTile"; + +/// Rust-side facts needed by call classification and operation emission. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct TensorViewRust { + /// Logical tensor element format, which may differ from its packed storage type. + pub logical: Ty, + /// Rust scalar type holding the tensor's physical bits. + pub storage: Ty, + /// Read-only or writable authority encoded by the source tensor type. + pub access: CuteTensorAccessAttr, + /// Decoded contiguous, divided or selected-tile coordinate mapping. + pub layout: CuteTensorLayoutAttr, +} + +/// Rust-side facts for `RegisterTile`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct RegisterTileRust { + /// Logical scalar type held by each register-tile position. + pub element: Ty, + /// Number of scalar elements in the fixed register tile. + pub size: u64, +} + +/// Recover an ADT's actual defining path, independent of public aliases. +pub(super) fn canonical_adt_path(def: &AdtDef) -> String { + super::identity::canonical_def_path(def.def_id(), &def.krate().name.to_string()) +} + +/// Require the expected ordered generic-argument kinds for a source tensor type. +pub(super) fn exact_schema( + args: &GenericArgs, + expected: &[&str], + owner: &str, +) -> Result<(), String> { + let found: Vec<_> = args + .0 + .iter() + .map(|arg| match arg { + GenericArgKind::Type(_) => "Type", + GenericArgKind::Const(_) => "Const", + GenericArgKind::Lifetime(_) => "Lifetime", + }) + .collect(); + if found == expected { + Ok(()) + } else { + Err(format!( + "`{owner}` expects generics [{}], found [{}]", + expected.join(", "), + found.join(", ") + )) + } +} + +/// Read a type argument at the required generic position, rejecting other argument kinds. +pub(super) fn type_arg(args: &GenericArgs, index: usize, owner: &str) -> Result { + match args.0.get(index) { + Some(GenericArgKind::Type(value)) => Ok(*value), + _ => Err(format!("`{owner}` generic {index} must be a type")), + } +} + +/// Read a constant argument at the required generic position, rejecting other argument kinds. +pub(super) fn const_arg(args: &GenericArgs, index: usize, owner: &str) -> Result { + match args.0.get(index) { + Some(GenericArgKind::Const(value)) => Ok(value.clone()), + _ => Err(format!("`{owner}` generic {index} must be a const")), + } +} + +/// Evaluate a source constant and reject values that do not fit in `u64`. +pub(super) fn unsigned_const(value: &TyConst, what: &str) -> Result { + let raw = match value.kind() { + TyConstKind::Value(_, allocation) => allocation + .read_uint() + .map_err(|error| format!("cannot read {what}: {error:?}"))?, + _ => u128::from( + value + .eval_target_usize() + .map_err(|error| format!("cannot evaluate {what}: {error:?}"))?, + ), + }; + u64::try_from(raw).map_err(|_| format!("{what} value {raw} does not fit in u64")) +} + +/// Decode the contiguous, divided or selected-tile layout of a source tensor. +fn decode_layout(ty: &Ty) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + let path = canonical_adt_path(&def); + match path.as_str() { + CONTIGUOUS_1D => { + exact_schema(&args, &[], CONTIGUOUS_1D)?; + Ok(Some(CuteTensorLayoutAttr::Contiguous1D)) + } + ZIPPED_1D | TILE_1D => { + exact_schema(&args, &["Const"], &path)?; + let tile = unsigned_const(&const_arg(&args, 0, &path)?, "tensor tile width")?; + if tile == 0 { + return Err("tensor tile width must be greater than zero".to_string()); + } + Ok(Some(if path == ZIPPED_1D { + CuteTensorLayoutAttr::Zipped1D(tile) + } else { + CuteTensorLayoutAttr::Tile1D(tile) + })) + } + _ => Ok(None), + } +} + +/// Decode an exact v0 `Tensor` or `TensorMut` ADT. +/// +/// `Ok(None)` means either “not a tensor” or “a tensor using a layout not in +/// v0”; both intentionally keep the existing generic ADT translation. +pub(crate) fn decode_tensor_view(ty: &Ty) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + let path = canonical_adt_path(&def); + let access = match path.as_str() { + TENSOR => CuteTensorAccessAttr::ReadOnly, + TENSOR_MUT => CuteTensorAccessAttr::ReadWrite, + _ => return Ok(None), + }; + exact_schema(&args, &["Lifetime", "Type", "Type", "Type"], &path)?; + let logical = type_arg(&args, 1, &path)?; + let Some(layout) = decode_layout(&type_arg(&args, 2, &path)?)? else { + return Ok(None); + }; + let storage = type_arg(&args, 3, &path)?; + Ok(Some(TensorViewRust { + logical, + storage, + access, + layout, + })) +} + +/// Decode a tensor view behind exactly one Rust reference or raw pointer. +pub(crate) fn decode_tensor_view_receiver(ty: &Ty) -> Result, String> { + match ty.kind() { + TyKind::RigidTy(RigidTy::Ref(_, pointee, _) | RigidTy::RawPtr(pointee, _)) => { + decode_tensor_view(&pointee) + } + _ => decode_tensor_view(ty), + } +} + +/// Decode an exact `RegisterTile`. +pub(crate) fn decode_register_tile(ty: &Ty) -> Result, String> { + let TyKind::RigidTy(RigidTy::Adt(def, args)) = ty.kind() else { + return Ok(None); + }; + if canonical_adt_path(&def) != REGISTER_TILE { + return Ok(None); + } + exact_schema(&args, &["Type", "Const"], REGISTER_TILE)?; + let element = type_arg(&args, 0, REGISTER_TILE)?; + let size = unsigned_const(&const_arg(&args, 1, REGISTER_TILE)?, "register tile width")?; + if size == 0 { + return Err("register tile width must be greater than zero".to_string()); + } + Ok(Some(RegisterTileRust { element, size })) +} + +/// Type-translation hook called before generic ADT lowering. +pub(crate) fn try_translate_tensor_view_type( + ctx: &mut Context, + rust_ty: &Ty, +) -> Option> { + let decoded = match decode_tensor_view(rust_ty) { + Ok(Some(decoded)) => decoded, + Ok(None) => return None, + Err(error) => { + return Some(Err(input_error_noloc!(TranslationErr::unsupported( + format!("invalid cute-rs tensor view type: {error}") + )))); + } + }; + + Some((|| { + let logical = translate_type(ctx, &decoded.logical)?; + let storage = translate_type(ctx, &decoded.storage)?; + let alignment = decoded + .storage + .layout() + .map_err(|error| { + input_error_noloc!(TranslationErr::unsupported(format!( + "cannot read cute-rs tensor storage alignment: {error:?}" + ))) + })? + .shape() + .abi_align; + Ok(CuteTensorViewType::get( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Gmem, + decoded.access, + alignment, + decoded.layout, + ) + .into()) + })()) +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/facts.rs b/cuda-oxide/crates/mir-importer/src/translator/facts.rs index c9542c38ff..f2c6bae88e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/facts.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/facts.rs @@ -417,15 +417,6 @@ mod pointer_origin { } } - /// ABI RULE: `cuda_device::SharedPtr` stores a `*mut SharedSlot`, so - /// the shared pointer it lowers to is always `RawMut`. - pub(crate) fn abi_shared_ptr_field() -> PointerOrigin { - PointerOrigin { - kind: MirPointerKind::RawMut, - mutable: true, - } - } - /// ABI RULE: `cuda_device::DisjointSlice<'_, T>` stores its data pointer /// as `*mut T`, so its data/element pointers are always `RawMut`. pub(crate) fn abi_disjoint_slice_data_ptr() -> PointerOrigin { @@ -505,6 +496,7 @@ pub(crate) struct ReferenceParamValidity { /// `DisjointSlice` produce no fact. Unsupported reference pass modes and source /// argument mismatches are errors rather than guesses about the ABI mapping. pub(crate) fn reference_param_validity( + instance: &rustc_public::mir::mono::Instance, ty: &rustc_public::ty::Ty, fn_abi: &rustc_public::abi::FnAbi, source_index: usize, @@ -542,15 +534,45 @@ pub(crate) fn reference_param_validity( return Ok(None); } - let attrs = match (&argument.mode, is_slice) { - (PassMode::Direct(attrs), false) | (PassMode::Pair(attrs, _), true) => attrs, + match (&argument.mode, is_slice) { + (PassMode::Direct(_), false) | (PassMode::Pair(_, _), true) => {} _ => { return Err(format!( "kernel reference argument {source_index} has an unsupported rustc FnAbi pass mode" )); } - }; - let flags = attrs.regular(); + } + + let (no_alias, no_free, pointee_size_bytes) = rustc_middle::ty::tls::with(|tcx| { + use rustc_middle::ty::{List, TypingEnv}; + use rustc_target::callconv::{ArgAttribute, PassMode as InternalPassMode}; + + let internal_instance = rustc_public::rustc_internal::internal(tcx, *instance); + let query_input = + TypingEnv::fully_monomorphized().as_query_input((internal_instance, List::empty())); + let internal_fn_abi = tcx + .fn_abi_of_instance(query_input) + .map_err(|err| format!("could not query internal FnAbi: {err:?}"))?; + let internal_arg = internal_fn_abi.args.get(source_index).ok_or_else(|| { + format!("kernel source argument index {source_index} is absent from internal FnAbi") + })?; + let attrs = match (&internal_arg.mode, is_slice) { + (InternalPassMode::Direct(attrs), false) | (InternalPassMode::Pair(attrs, _), true) => { + attrs + } + _ => { + return Err(format!( + "kernel reference argument {source_index} has an unsupported rustc FnAbi pass mode" + )); + } + }; + Ok(( + attrs.regular.contains(ArgAttribute::NoAlias), + attrs.regular.contains(ArgAttribute::NoFree), + attrs.pointee_size.bytes(), + )) + })?; + // Do not consume rustc's `read_only`: NVPTX (including libNVVM) can lower // volatile reads through a readonly + noalias kernel parameter to cached // non-volatile loads. Let the optimizer infer read-only access from the @@ -560,14 +582,10 @@ pub(crate) fn reference_param_validity( // dereferenceability as lasting for the whole call. rustc's NoFree is the // proof needed by all supported consumers; producer LLVM version alone is // insufficient. Zero covers both ZST pointees and unknown slice lengths. - let dereferenceable_bytes = if flags.no_free { - attrs.pointee_size().bytes() as u64 - } else { - 0 - }; + let dereferenceable_bytes = if no_free { pointee_size_bytes } else { 0 }; Ok(Some(ReferenceParamValidity { pointee_alignment: shape.abi_align, - no_alias: flags.no_alias, + no_alias, dereferenceable_bytes, })) } @@ -1055,7 +1073,7 @@ pub fn reference_validity( .iter() .enumerate() .map(|(index, ty)| { - reference_param_validity(ty, &fn_abi, index) + reference_param_validity(&instance, ty, &fn_abi, index) .expect("source parameter matches ABI") }) .collect::>(); @@ -1063,19 +1081,27 @@ pub fn reference_validity( // Reusing a neighboring argument after dropping an ignored // ZST or flattening a slice must not mint reference facts. let mut errors = vec![ - reference_param_validity(&source_types[1], &fn_abi, 0).unwrap_err(), - reference_param_validity(&source_types[1], &fn_abi, fn_abi.args.len()) + reference_param_validity(&instance, &source_types[1], &fn_abi, 0) .unwrap_err(), + reference_param_validity( + &instance, + &source_types[1], + &fn_abi, + fn_abi.args.len(), + ) + .unwrap_err(), ]; let mut wrong_mode = fn_abi.clone(); wrong_mode.args[1].mode = fn_abi.args[15].mode.clone(); assert!(matches!(wrong_mode.args[1].mode, PassMode::Indirect { .. })); errors.push( - reference_param_validity(&source_types[1], &wrong_mode, 1).unwrap_err(), + reference_param_validity(&instance, &source_types[1], &wrong_mode, 1) + .unwrap_err(), ); wrong_mode.args[3].mode = fn_abi.args[1].mode.clone(); errors.push( - reference_param_validity(&source_types[3], &wrong_mode, 3).unwrap_err(), + reference_param_validity(&instance, &source_types[3], &wrong_mode, 3) + .unwrap_err(), ); std::ops::ControlFlow::<(), _>::Continue((facts, errors)) }) diff --git a/cuda-oxide/crates/mir-importer/src/translator/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/mod.rs index 784ca58758..8d11a85ed0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/mod.rs @@ -50,6 +50,7 @@ pub mod block; pub mod body; +pub(crate) mod cute; pub(crate) mod facts; pub(crate) mod layout; pub(crate) mod location; @@ -108,6 +109,9 @@ pub fn register_dialects(ctx: &mut Context) { // dialect-nvvm is required for thread / block / warp intrinsics. dialect_nvvm::register(ctx); + // dialect-cute keeps tensor and tile meaning visible to both backends. + dialect_cute::register(ctx); + // The builtin dialect (ModuleOp etc.) is auto-registered by pliron 0.14. } diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs index 5685038de6..8131b47bca 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_alloc.rs @@ -506,10 +506,10 @@ fn decode_relocation_addend( /// Materialize a thin pointer field from an aggregate constant's allocation. /// /// Aggregate **const** values with thin pointers to device statics are -/// materialized via [`MirGlobalAllocOp`] per field (addend taken from the +/// materialized via [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) per field (addend taken from the /// relocation's stored bytes). This does **not** mean device-static /// *initializers* that themselves contain pointer relocations are supported — -/// those remain rejected by [`allocation_initializer_data`]. +/// those remain rejected by [`allocation_initializer_data`](super::statics::allocation_initializer_data). /// /// When the field has no provenance entry, falls back to the existing /// inttoptr-of-bytes path (null / exposed provenance). diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs index b275dac7a5..dd7f699a36 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_bytes.rs @@ -37,7 +37,7 @@ use std::num::NonZeroUsize; /// [..]` indexed by runtime value) to a `MirConstructArrayOp`. Element stride /// and aggregate field offsets come from rustc layout. Thin pointer fields /// inside elements that relocate to device statics are materialized via -/// [`MirGlobalAllocOp`] per field. +/// [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) per field. pub(super) fn translate_array_value_constant( ctx: &mut Context, constant: &mir::ConstOperand, @@ -466,7 +466,7 @@ fn translate_array_value_constant_inner( /// constant is a promoted `&Struct`) /// 2. Decode each field at its rustc layout offset from that allocation /// 3. Thin pointer fields that relocate to device statics become -/// [`MirGlobalAllocOp`] results; other fields use the byte decoder +/// [`MirGlobalAllocOp`](dialect_mir::ops::MirGlobalAllocOp) results; other fields use the byte decoder /// 4. Create MirConstructStructOp with those operands /// /// Each field is read at the byte offset rustc's layout records for it, so diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs index 7ce8468d31..777de9ca3d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs @@ -368,6 +368,20 @@ pub fn translate_rvalue( )?; let result_type = types::translate_type(ctx, ty)?; + let result_type = if matches!( + kind, + mir::CastKind::PtrToPtr + | mir::CastKind::PointerCoercion(mir::PointerCoercion::MutToConstPointer) + ) && let Some(src_as) = + super::super::values::pointer_addr_space(ctx, operand_val.get_type(ctx)) + && src_as != dialect_mir::types::address_space::GENERIC + && super::super::values::pointer_addr_space(ctx, result_type) + == Some(dialect_mir::types::address_space::GENERIC) + { + super::super::values::align_pointer_addr_space(ctx, result_type, src_as) + } else { + result_type + }; // A cast whose source and target lower to the same dialect type // changes nothing. `isize` and `i64` are that pair on this target: diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/fn_ptr.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/fn_ptr.rs index c58e9cbe87..69f8061d66 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/fn_ptr.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/fn_ptr.rs @@ -28,6 +28,10 @@ use std::num::NonZeroUsize; /// name, never 0 so it cannot look like a null pointer) and casts it /// int -> ptr. See the comment at the `Rvalue::Cast` arm for why a token /// stands in for a code address on the device. +fn is_cute_rs_crate(crate_name: &str) -> bool { + matches!(crate_name, "cute_rs" | "cute-rs") +} + pub(super) fn translate_reify_fn_pointer( ctx: &mut Context, body: &mir::Body, @@ -56,6 +60,16 @@ pub(super) fn translate_reify_fn_pointer( )) ); }; + let crate_name = fn_def.krate().name; + if is_cute_rs_crate(crate_name.as_str()) { + let item_name = fn_def.name(); + return input_err!( + loc, + TranslationErr::unsupported(format!( + "CuTe function `{item_name}` must be called directly; converting it to a function pointer would erase its layout configuration" + )) + ); + } let raw_intrinsic = crate::translator::terminator::intrinsics::generated::require_supported_raw_intrinsic( fn_def, &loc, diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs index b5c1097339..89d9a9406e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_addr.rs @@ -64,7 +64,6 @@ pub(crate) fn translate_place_address( prev_op: Option>, loc: Location, ) -> TranslationResult>)>> { - let place = &*super::place_read::strip_shared_ptr_field_projections(body, place); let Some(slot) = value_map.get_slot(place.local) else { return Ok(None); }; @@ -572,11 +571,11 @@ pub(super) fn emit_array_subslice_value( /// Walks the projection chain and emits the correct pliron ops for each /// element: /// -/// - `Field(idx, _)` → [`MirFieldAddrOp`] +/// - `Field(idx, _)` → [`MirFieldAddrOp`](dialect_mir::ops::MirFieldAddrOp) /// - `ConstantIndex {offset, from_end, ..}` → an element address. Forward /// indexes materialize `offset` directly; from-end indexes on slices extract /// the fat-pointer length and materialize `len - offset` at runtime. -/// - `Index(local)` → `load_local(local)` + [`MirArrayElementAddrOp`] +/// - `Index(local)` → `load_local(local)` + [`MirArrayElementAddrOp`](dialect_mir::ops::MirArrayElementAddrOp) /// (array pointee) or `load_local(local)` + [`MirPtrOffsetOp`] (slice data pointer) /// - `Deref` → `MirLoadOp` of the pointer (the loaded pointer IS /// the pointee's address); subsequent projections apply to the pointee. diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_read.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_read.rs index dc3f96e9d1..193b39c8ba 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_read.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_read.rs @@ -80,7 +80,11 @@ pub fn translate_place( prev_op: Option>, loc: Location, ) -> TranslationResult<(Value, Option>)> { - let place = &*strip_shared_ptr_field_projections(body, place); + if place.projection.is_empty() + && let Some(value) = value_map.get_direct_value(place.local) + { + return Ok((value, prev_op)); + } match classify_place_read_strategy(ctx, place, value_map)? { PlaceReadStrategy::Address => { if let Some((value, last_op)) = translate_place_load_from_address( @@ -110,59 +114,6 @@ pub fn translate_place( } } -/// Drop the projections of a place onto the field of a `cuda_device::SharedPtr`. -/// -/// `SharedPtr { ptr: *mut SharedSlot }` is a scalar-lowered newtype: its -/// value is its one field, a CTA-shared pointer (see `types::translate_type`). -/// Projecting that field is therefore the identity. Without this, the place -/// walkers would read the pointer value of a `SharedPtr` nested in another -/// place, such as `(region.base).ptr`, as the address of an aggregate. -/// -/// `ThreadIndex`, the other scalar-lowered newtype, needs no strip. Its scalar -/// is an integer, and `dialect-mir` and `mir-lower` already treat -/// `extract_field` on an integer scalar as the identity. The iterative place -/// walker in `place_iter.rs` takes any pointer-typed current value as an -/// address to load from, so only a newtype whose scalar is itself a pointer is -/// misread. Follow-up: have that walker track whether its current value is an -/// address or a value, then drop this strip. -pub(crate) fn strip_shared_ptr_field_projections<'a>( - body: &mir::Body, - place: &'a mir::Place, -) -> std::borrow::Cow<'a, mir::Place> { - use rustc_public::ty::{RigidTy, TyKind}; - - if !place - .projection - .iter() - .any(|elem| matches!(elem, mir::ProjectionElem::Field(0, _))) - { - return std::borrow::Cow::Borrowed(place); - } - let mut place_ty = body.locals()[place.local].ty; - let mut kept = Vec::with_capacity(place.projection.len()); - for elem in &place.projection { - let is_shared_ptr = matches!( - place_ty.kind(), - TyKind::RigidTy(RigidTy::Adt(adt_def, _)) if facts::is_cuda_device_adt(&adt_def, "SharedPtr") - ); - let Ok(next_ty) = elem.ty(place_ty) else { - return std::borrow::Cow::Borrowed(place); - }; - if !(is_shared_ptr && matches!(elem, mir::ProjectionElem::Field(0, _))) { - kept.push(elem.clone()); - } - place_ty = next_ty; - } - if kept.len() == place.projection.len() { - std::borrow::Cow::Borrowed(place) - } else { - std::borrow::Cow::Owned(mir::Place { - local: place.local, - projection: kept, - }) - } -} - // ============================================================================ // Place Read Strategy and Address Path // ============================================================================ diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs index 301415d0a4..78012fbfe0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs @@ -369,7 +369,7 @@ pub fn emit_function_call( // arity is invalid and must remain visible to MirCallOp verification. if arg_values.len() == expected_arguments.len() { for (argument, expected) in arg_values.iter_mut().zip(expected_arguments) { - let (normalized, normalized_last_op) = normalize_foreign_pointer_argument( + let (normalized, normalized_last_op) = normalize_call_pointer_argument( ctx, *argument, expected, @@ -383,6 +383,45 @@ pub fn emit_function_call( } } + // Ordinary Rust helper signatures also expose generic pointers. Resolve + // the monomorphized operand type (not an unresolved trait signature), and + // adapt only address space. All other type/provenance mismatches still + // reach the independently typed callee's verifier unchanged. + if external_callee_type.is_none() { + for (argument, operand) in arg_values.iter_mut().zip(args) { + if argument + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_none() + { + continue; + } + let rust_type = match operand.ty(body.locals()) { + Ok(ty) => ty, + Err(error) => { + return input_err!( + loc.clone(), + TranslationErr::unsupported(format!( + "failed to resolve Rust call argument type: {error:?}" + )) + ); + } + }; + let declared = crate::translator::types::translate_type(ctx, &rust_type)?; + let (normalized, tail) = normalize_call_pointer_argument( + ctx, + *argument, + declared, + block_ptr, + last_op, + loc.clone(), + ); + *argument = normalized; + last_op = tail; + } + } + let prepared_destination = if target.is_some() { let (prepared, prepared_last_op) = prepare_destination_write( ctx, @@ -458,9 +497,9 @@ pub fn emit_function_call( )) } -/// Adapt a foreign-call pointer argument only when the address space is the +/// Adapt a call pointer argument only when the address space is the /// sole difference from its declared ABI parameter type. -fn normalize_foreign_pointer_argument( +fn normalize_call_pointer_argument( ctx: &mut Context, value: Value, declared_type: TypeHandle, @@ -872,7 +911,7 @@ mod tests { let block = BasicBlock::new(&mut ctx, None, vec![shared_raw_mut]); let argument = block.deref(&ctx).get_argument(0); - let (normalized, cast_op) = normalize_foreign_pointer_argument( + let (normalized, cast_op) = normalize_call_pointer_argument( &mut ctx, argument, abi_raw_mut, @@ -913,7 +952,7 @@ mod tests { let block = BasicBlock::new(&mut ctx, None, vec![shared_raw_mut]); let argument = block.deref(&ctx).get_argument(0); - let (unchanged, cast_op) = normalize_foreign_pointer_argument( + let (unchanged, cast_op) = normalize_call_pointer_argument( &mut ctx, argument, abi_unique_ref, @@ -941,7 +980,7 @@ mod tests { let block = BasicBlock::new(&mut ctx, None, vec![abi_raw_mut]); let argument = block.deref(&ctx).get_argument(0); - let (unchanged, cast_op) = normalize_foreign_pointer_argument( + let (unchanged, cast_op) = normalize_call_pointer_argument( &mut ctx, argument, abi_raw_mut, @@ -974,7 +1013,7 @@ mod tests { let block = BasicBlock::new(&mut ctx, None, vec![shared_raw_mut]); let argument = block.deref(&ctx).get_argument(0); - let (unchanged, cast_op) = normalize_foreign_pointer_argument( + let (unchanged, cast_op) = normalize_call_pointer_argument( &mut ctx, argument, wrong_pointee, @@ -1004,7 +1043,7 @@ mod tests { let erased_block = BasicBlock::new(&mut ctx, None, vec![mutable_erased]); let erased_argument = erased_block.deref(&ctx).get_argument(0); - let (unchanged, cast_op) = normalize_foreign_pointer_argument( + let (unchanged, cast_op) = normalize_call_pointer_argument( &mut ctx, erased_argument, immutable_erased, diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs index 230d9328f1..a24f6f9395 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs index 8564244bef..55baf4c4e7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `clc` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs index 0003ca73b3..dbbee9abaa 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs index 9b98425297..1ec9e8cfab 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs index 2f094cbde2..c968753ea0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs index d7121a8933..34f9e703f8 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs index 8a0f114705..af3cdbcd30 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs index 5e57988c8f..84c3287f38 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `elect` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs index ae5162f1bf..e61e80d9ce 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs index 0ca7643ffb..258bb7f2ce 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs index a75e5f01c3..8dd7273051 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/helpers.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Shared glue for the generated intrinsic dispatch shards. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs index 81160e1b04..b5fedc0e92 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs index 89203b59b7..4afd34b8f0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs index abbfe8b90c..5f6fc97b20 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs index e7b455627c..cfec72c760 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `mbarrier_extended` intrinsics. @@ -13,10 +13,12 @@ use super::helpers::require_arity; use crate::error::{TranslationErr, TranslationResult}; use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{ - FenceMbarrierInitReleaseClusterOp, FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceAcquireSharedClusterClusterOp, FenceMbarrierInitReleaseClusterOp, + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, FenceProxyAsyncGenericReleaseSharedCtaClusterOp, FenceProxyAsyncSharedCtaOp, - MbarrierArriveClusterOp, MbarrierArriveExpectTxClusterOp, MbarrierArriveExpectTxSharedOp, - MbarrierTryWaitParityClusterOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, + MbarrierArriveClusterOp, MbarrierArriveClusterRelaxedOp, MbarrierArriveExpectTxClusterOp, + MbarrierArriveExpectTxSharedOp, MbarrierTryWaitParityClusterOp, + MbarrierTryWaitParityClusterRelaxedOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, NanosleepOp, }; use pliron::basic_block::BasicBlock; @@ -41,6 +43,31 @@ pub(super) fn try_dispatch( loc: Location, ) -> TranslationResult>> { match name { + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" + | "cuda_device::barrier::fence_acquire_shared_cluster_cluster" => { + require_arity(name, args.len(), 0, &loc)?; + let last_op = prev_op; + let extended = FenceAcquireSharedClusterClusterOp::build(ctx); + extended.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, extended, "v1:i1042"); + helpers::insert_op(ctx, extended, block_ptr, last_op); + if let Some(target_idx) = target { + Ok(Some(helpers::emit_goto( + ctx, + *target_idx, + extended, + block_map, + loc, + ))) + } else { + input_err!( + loc, + TranslationErr::unsupported( + "fence_acquire_shared_cluster_cluster call without target block".to_owned() + ) + ) + } + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_device::barrier::fence_mbarrier_init_release_cluster" => { require_arity(name, args.len(), 0, &loc)?; @@ -164,6 +191,39 @@ pub(super) fn try_dispatch( ) } } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" + | "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" => { + require_arity(name, args.len(), 1, &loc)?; + let (address, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let extended = MbarrierArriveClusterRelaxedOp::build(ctx, address); + extended.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, extended, "v1:i1040"); + helpers::insert_op(ctx, extended, block_ptr, last_op); + if let Some(target_idx) = target { + Ok(Some(helpers::emit_goto( + ctx, + *target_idx, + extended, + block_map, + loc, + ))) + } else { + input_err!( + loc, + TranslationErr::unsupported( + "mbarrier_arrive_cluster_relaxed call without target block".to_owned() + ) + ) + } + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_device::barrier::mbarrier_arrive_expect_tx" => { require_arity(name, args.len(), 3, &loc)?; @@ -404,6 +464,54 @@ pub(super) fn try_dispatch( "mbarrier_try_wait_parity_cluster call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" + | "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" => { + require_arity(name, args.len(), 2, &loc)?; + let (barrier, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let (parity, last_op) = rvalue::translate_operand( + ctx, + body, + &args[1], + value_map, + block_ptr, + last_op, + loc.clone(), + )?; + let extended = MbarrierTryWaitParityClusterRelaxedOp::build(ctx, barrier, parity); + extended.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, extended, "v1:i1041"); + let (prepared_destination, prepared_last_op) = helpers::prepare_destination_write( + ctx, + body, + destination, + value_map, + block_ptr, + last_op, + loc.clone(), + )?; + helpers::insert_op(ctx, extended, block_ptr, prepared_last_op); + let result = extended.deref(ctx).get_result(0); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + extended, + value_map, + block_map, + loc, + "mbarrier_try_wait_parity_cluster_relaxed call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0316" | "cuda_device::barrier::nanosleep" => { require_arity(name, args.len(), 1, &loc)?; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs index e3355982d3..07bc667300 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility path dispatch for CUDA intrinsics. @@ -619,6 +619,8 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::float::ex2_approx_f32" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820" | "cuda_device::float::ex2_approx_ftz_f32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" + | "cuda_device::barrier::fence_acquire_shared_cluster_cluster" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_device::barrier::fence_mbarrier_init_release_cluster" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" @@ -848,6 +850,8 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::barrier::mbarrier_arrive" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308" | "cuda_device::barrier::mbarrier_arrive_cluster" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" + | "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_device::barrier::mbarrier_arrive_expect_tx" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" @@ -866,6 +870,8 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::barrier::mbarrier_try_wait_parity" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311" | "cuda_device::barrier::mbarrier_try_wait_parity_cluster" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" + | "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_device::bf16::min_bf16" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067" @@ -1176,6 +1182,8 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::wmma::mma_m16n8k32_s32_u8_satfinite" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0520" | "cuda_device::wmma::mma_m16n8k4_f32_tf32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" + | "cuda_device::wmma::mma_m16n8k64_mxf4_f32" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" | "cuda_device::wmma::mma_m16n8k64_s32_s4" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0153" @@ -2472,6 +2480,7 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1016" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0819" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0314" @@ -2586,6 +2595,7 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0576" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0098" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1017" @@ -2595,6 +2605,7 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0309" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0310" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830" @@ -2750,6 +2761,7 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0124" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0123" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0520" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0153" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0146" @@ -3766,6 +3778,8 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::float::ex2_approx_f32" => Some("v1:i0819"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820" | "cuda_device::float::ex2_approx_ftz_f32" => Some("v1:i0820"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" + | "cuda_device::barrier::fence_acquire_shared_cluster_cluster" => Some("v1:i1042"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_device::barrier::fence_mbarrier_init_release_cluster" => Some("v1:i0313"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" @@ -4002,6 +4016,8 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::barrier::mbarrier_arrive" => Some("v1:i0098"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308" | "cuda_device::barrier::mbarrier_arrive_cluster" => Some("v1:i0308"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" + | "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" => Some("v1:i1040"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_device::barrier::mbarrier_arrive_expect_tx" => Some("v1:i0306"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" @@ -4020,6 +4036,8 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::barrier::mbarrier_try_wait_parity" => Some("v1:i0310"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311" | "cuda_device::barrier::mbarrier_try_wait_parity_cluster" => Some("v1:i0311"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" + | "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" => Some("v1:i1041"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_device::bf16::min_bf16" => { Some("v1:i0838") } @@ -4332,6 +4350,8 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::wmma::mma_m16n8k32_s32_u8_satfinite" => Some("v1:i0123"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0520" | "cuda_device::wmma::mma_m16n8k4_f32_tf32" => Some("v1:i0520"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" + | "cuda_device::wmma::mma_m16n8k64_mxf4_f32" => Some("v1:i1039"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" | "cuda_device::wmma::mma_m16n8k64_s32_s4" => Some("v1:i0145"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0153" diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs index e88cca166a..6e935cd437 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs index 2dbde4941c..74bd7cb7e3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs index 17ace3efe1..e0a7401bb3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs index 72cca157a1..cb88a6307f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs index d0ca47ff84..4eaadfa8e9 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `prmt` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs index a4266d3dcb..58f1cbd014 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `redux` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs index 96fbc9a321..556940ad95 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `register_mma` intrinsics. @@ -7298,6 +7298,65 @@ pub(super) fn try_dispatch( "mma_m16n8k4_f32_tf32 call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" + | "cuda_device::wmma::mma_m16n8k64_mxf4_f32" => { + require_arity(name, args.len(), 9, &loc)?; + let (operands, last_op, result_ty, result_count) = import_generated_mma_operands( + ctx, + body, + args, + block_ptr, + prev_op, + value_map, + loc.clone(), + GeneratedMmaImportAdapter::C4F32A4U32B2U32Scales2U32Selectors4U16ToD4F32, + )?; + let mma = Operation::new( + ctx, + RegisterMmaOp::get_concrete_op_info(), + vec![result_ty; result_count], + operands, + vec![], + 0, + ); + mma.deref_mut(ctx).set_loc(loc.clone()); + let mma = RegisterMmaOp::new(mma); + mma.set_attr_nvvm_register_mma_shape(ctx, RegisterMmaShapeAttr::M16n8k64); + mma.set_attr_nvvm_register_mma_operation(ctx, RegisterMmaOperationAttr::Multiply); + mma.set_attr_nvvm_register_mma_kind(ctx, RegisterMmaKindAttr::Mxf4); + mma.set_attr_nvvm_register_mma_accumulator(ctx, RegisterMmaAccumulatorAttr::F32); + mma.set_attr_nvvm_register_mma_a_element(ctx, RegisterMmaElementAttr::E2m1); + mma.set_attr_nvvm_register_mma_b_element(ctx, RegisterMmaElementAttr::E2m1); + mma.set_attr_nvvm_register_mma_a_layout(ctx, RegisterMmaLayoutAttr::Row); + mma.set_attr_nvvm_register_mma_b_layout(ctx, RegisterMmaLayoutAttr::Col); + mma.set_attr_nvvm_register_mma_overflow(ctx, RegisterMmaOverflowAttr::NotApplicable); + let mma = mma.get_operation(); + helpers::set_generated_intrinsic_marker(ctx, mma, "v1:i1039"); + let (prepared_destination, prepared_last_op) = helpers::prepare_destination_write( + ctx, + body, + destination, + value_map, + block_ptr, + Some(last_op), + loc.clone(), + )?; + helpers::insert_op(ctx, mma, block_ptr, prepared_last_op); + let (result, last_op) = + bundle_generated_mma_results(ctx, mma, result_ty, result_count, loc.clone()); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + last_op, + value_map, + block_map, + loc, + "mma_m16n8k64_mxf4_f32 call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" | "cuda_device::wmma::mma_m16n8k64_s32_s4" => { require_arity(name, args.len(), 3, &loc)?; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs index 5169f0a38e..ddee176556 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs index fa13d0392c..218b38c3ae 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs index e3dcd08795..a9dfbc33c3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs index 0caab1ecbc..734e1c2ac7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs index 424884c665..10439f8b22 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sreg` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs index 91eaeb2c98..7269a3bb1e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs index d512c13821..0ef28b175c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sync` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs index 33d2e9000d..67f719821c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_cp.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_cp` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs index c0f96923fa..c51a81c377 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_ld.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_ld` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs index c7bcf88c9f..1dc99e58ea 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs index d335198d0b..5d6d4cdf06 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_other.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_other` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs index 9a24687307..af5044a941 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tcgen05_st.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tcgen05_st` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs index 76af4ebd4b..f8c8ddc1c6 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs index fadcb28f60..b164c399cb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `clc` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs index 37efbd61da..d65ec6b9ca 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs index f2539775fc..47fd9d2510 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs index fdb56c3acd..10fb3aaa3d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs index 1d3e283004..488826b2b1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs index 6044b37f76..ed6dffac3d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs index 6154ad507d..1e73fcb049 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `elect` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs index 32eb53bd92..320c7ff08c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs index ebaedbc13c..0b5c6773a3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs index 4ed38845d3..8eba2f1917 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs index 33ee956869..d7db141e4c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs index 6d036a77a1..9c81695619 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_extended.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_extended.rs index 18c76eafc5..d0a4755f75 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_extended.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mbarrier_extended.rs @@ -4,13 +4,49 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `mbarrier_extended` intrinsics. use super::super::*; +#[test] +fn fence_acquire_shared_cluster_cluster_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042"), + Some("v1:i1042") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::barrier::fence_acquire_shared_cluster_cluster" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::barrier::fence_acquire_shared_cluster_cluster" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::barrier::fence_acquire_shared_cluster_cluster" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::barrier::fence_acquire_shared_cluster_cluster"), + Some("v1:i1042") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1042" + )); +} #[test] fn fence_mbarrier_init_release_cluster_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( @@ -198,6 +234,42 @@ fn mbarrier_arrive_cluster_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn mbarrier_arrive_cluster_relaxed_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040"), + Some("v1:i1040") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::barrier::mbarrier_arrive_cluster_relaxed" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::barrier::mbarrier_arrive_cluster_relaxed"), + Some("v1:i1040") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1040" + )); +} +#[test] fn mbarrier_arrive_expect_tx_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" @@ -378,6 +450,44 @@ fn mbarrier_try_wait_parity_cluster_uses_only_canonical_or_compatibility_defpath )); } #[test] +fn mbarrier_try_wait_parity_cluster_relaxed_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041"), + Some("v1:i1041") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster_relaxed" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" + )); + assert_eq!( + generated_intrinsic_marker( + "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" + ), + Some("v1:i1041") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1041" + )); +} +#[test] fn nanosleep_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0316" diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs index 565f8d0f47..0bee39eaa1 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests, grouped per intrinsic family. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs index 8494def465..d6c942687c 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs index 5f3d1cc7d5..88ec80f81b 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs index 21b7e977a4..7d92776321 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs index 631742faf9..2a74640dd2 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs index 0b15dff2e7..705edc5677 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `prmt` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs index df47440ddf..9de584943d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `redux` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs index c0bed7caf1..9f8dbf1ae5 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `register_mma` intrinsics. @@ -4440,6 +4440,42 @@ fn mma_m16n8k4_f32_tf32_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn mma_m16n8k64_mxf4_f32_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039"), + Some("v1:i1039") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::matrix::mma_m16n8k64_mxf4_f32" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::wmma::mma_m16n8k64_mxf4_f32" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::wmma::mma_m16n8k64_mxf4_f32" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::wmma::mma_m16n8k64_mxf4_f32"), + Some("v1:i1039") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1039" + )); +} +#[test] fn mma_m16n8k64_s32_s4_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs index c83a02f8c7..d9a69edb4e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs index 3ccd82f21f..3b9826558d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs index 810379af2c..29cddaa2a3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs index 016e6adff6..57c067083b 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs index 12b0d52620..9fd1dd0d24 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `sreg` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs index 57d1792af1..d4231a57ed 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs index 604b55d031..fc110ef0af 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `sync` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs index 2aab6e65a8..88f4c7c769 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `tcgen05` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs index a819be423d..c286ba6208 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs index fb8dfc0432..5b68a2a5be 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `vote` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs index 1fac7ad53d..9a19172d77 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs index 8a5dbddcf2..58ca44a029 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs index ae0ba0cb1e..3facbf75f8 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs index 89ac02894a..693497de06 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tests/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs index 8114ef4158..5f96afd655 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs index f0c6e4aa97..3fe1066099 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `vote` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs index c3b33fcb86..2b5650bc62 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs index 6efbb7eec0..536582e80a 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs index f0a267bfd8..c6440078a3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs index 2fd62b2fed..1f8cb1817d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/generated/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/memory.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/memory.rs index b8a3643d8c..9de8b59a6f 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/memory.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/memory.rs @@ -1311,8 +1311,9 @@ pub fn emit_cvta_generic_to_shared_offset( block_map: &[Ptr], loc: Location, result_width: u32, + cluster: bool, ) -> TranslationResult> { - use dialect_nvvm::ops::CvtaGenericToSharedOffsetOp; + use dialect_nvvm::ops::{CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp}; use pliron::builtin::types::Signedness; if args.len() != 1 { @@ -1348,7 +1349,11 @@ pub fn emit_cvta_generic_to_shared_offset( let result_ty = IntegerType::get(ctx, result_width, Signedness::Unsigned); let cvta_op = Operation::new( ctx, - CvtaGenericToSharedOffsetOp::get_concrete_op_info(), + if cluster { + CvtaGenericToClusterSharedOffsetOp::get_concrete_op_info() + } else { + CvtaGenericToSharedOffsetOp::get_concrete_op_info() + }, vec![result_ty.into()], vec![ptr_val], vec![], diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs index d88f37d4b0..8f640e87ef 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs @@ -968,6 +968,25 @@ fn translate_call( let (pattern_name, call_name, substs_args, type_substs) = extract_func_info(func, &loc)?; let loc = call_debug_location(pattern_name.as_deref(), loc); + // Keep CuTe calls as CuTe operations while their tensor and layout + // meaning is still visible. Ordinary source helpers continue through the + // normal Rust-body path. + if let Some(result) = crate::translator::cute::try_translate_cute_call( + ctx, + body, + func, + args, + destination, + &target.map(|block| block), + block_ptr, + prev_op, + value_map, + block_map, + &loc, + ) { + return result; + } + // Is the trait-method Self type SharedArray? Shared with // `values::classify_call` so intrinsic dispatch and destination-slot // classification can't drift. @@ -1076,7 +1095,7 @@ fn translate_call( ) ); }; - if factor == 1 || factor > 1024 { + if factor == 1 || (factor > 1024 && factor != u32::MAX - 1) { return input_err!( loc, TranslationErr::invalid_op(format!( @@ -2111,6 +2130,7 @@ fn extract_function_item_target( || instance.is_foreign_item() || !instance.has_body() || matches!(crate_name.as_str(), "cuda_device" | "cuda-device" | "libm") + || super::cute::function_item_requires_direct_dispatch(&fn_def) || generated_direct_call_only, })) } @@ -3342,6 +3362,7 @@ fn try_dispatch_intrinsic( block_map, loc, 64, + false, )?, )), "cuda_device::shared::cvta_generic_to_shared_u32" => Ok(Some( @@ -3357,9 +3378,26 @@ fn try_dispatch_intrinsic( block_map, loc, 32, + false, )?, )), + "cuda_device::shared::cvta_generic_to_cluster_shared_offset" => Ok(Some( + intrinsics::memory::emit_cvta_generic_to_shared_offset( + ctx, + body, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc, + 64, + true, + )?, + )), // Public SharedArray pointer conversions all narrow the shared-memory // base to a generic pointer. Recognition is shared with destination // address-space classification in translator::values. diff --git a/cuda-oxide/crates/mir-importer/src/translator/types.rs b/cuda-oxide/crates/mir-importer/src/translator/types.rs index 6e693acfaa..671a54ae19 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/types.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/types.rs @@ -364,14 +364,6 @@ fn translate_pointer_like( let elem = shared_array_element_type(ctx, &substs, "SharedArray")?; Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into()) } - rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, substs)) - if is_cuda_device_adt(&adt_def, "SharedSlot") => - { - // `*mut SharedSlot` is the field of `SharedPtr`: a pointer - // to a `T` in CTA-shared memory, by its type. - let elem = shared_array_element_type(ctx, &substs, "SharedSlot")?; - Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into()) - } rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, _substs)) if is_cuda_device_adt(&adt_def, "Barrier") => { @@ -670,6 +662,29 @@ pub fn translate_type( translate_pointer_like(ctx, &pointee, origin) } rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, substs)) => { + if let Some(result) = + crate::translator::cute::pipeline::try_translate_tma_load_pipeline_type( + ctx, rust_ty, + ) + { + return result; + } + if let Some(result) = + crate::translator::cute::scheduler::try_translate_work_tile_type(ctx, rust_ty) + { + return result; + } + if let Some(result) = + crate::translator::cute::block_scaled::try_translate_block_scaled_type(ctx, rust_ty) + { + return result; + } + if let Some(result) = + crate::translator::cute::tensor::try_translate_tensor_view_type(ctx, rust_ty) + { + return result; + } + // Get the trimmed name (just the type name without path) let trimmed_name = adt_def.trimmed_name(); @@ -726,14 +741,6 @@ pub fn translate_type( // ThreadIndex is a newtype around usize - translate to usize // The type safety is enforced at the Rust level, not the IR level Ok(get_usize_type(ctx).into()) - } else if is_cuda_device_adt(&adt_def, "SharedPtr") { - // `SharedPtr { ptr: *mut SharedSlot }` is a scalar-lowered - // newtype: the value is its field, a CTA-shared (`addrspace(3)`) - // pointer to `T`. Params, returns, locals and fields of type - // `SharedPtr` are therefore shared pointers by type, across - // function boundaries. - let elem = shared_array_element_type(ctx, &substs, "SharedPtr")?; - Ok(facts::mint_shared_ptr_type(ctx, elem, facts::abi_shared_ptr_field()).into()) } else if is_cuda_device_adt(&adt_def, "RangeToken") { Ok(dialect_iket::types::IketRangeTokenType::get(ctx).into()) } else if is_cuda_device_adt(&adt_def, "SharedArray") { @@ -764,7 +771,11 @@ pub fn translate_type( 64, pliron::builtin::types::Signedness::Unsigned, ); - Ok(llvm_export::types::ArrayType::get(ctx, i64_ty.into(), 16).into()) + // Keep the opaque storage in MIR until backend lowering. A + // typed transparent wrapper retains this as a nested field; + // injecting an LLVM array here bypasses MIR aggregate layout + // handling and breaks the explicit MLIR exporter. + Ok(dialect_mir::types::MirArrayType::get(ctx, i64_ty.into(), 16).into()) } else { // Generic ADT handling for user-defined structs and enums let variants = adt_def.variants(); @@ -875,10 +886,8 @@ pub fn translate_type( // Total struct size (bytes) let size: u64 = shape.size.bytes() as u64; let abi_kind = if adt_def.repr().flags.is_transparent - && matches!( - &shape.value_repr, - rustc_public::abi::ValueRepr::Scalar(_) - ) { + && matches!(&shape.abi, rustc_public::abi::ValueAbi::Scalar(_)) + { StructAbiKind::TransparentScalar } else { StructAbiKind::Aggregate diff --git a/cuda-oxide/crates/mir-importer/src/translator/values.rs b/cuda-oxide/crates/mir-importer/src/translator/values.rs index 450665aef0..357bb944f5 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/values.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/values.rs @@ -71,6 +71,12 @@ pub struct ValueMap { /// identity resolution stay entirely out of Off/LineTables builds without /// threading another flag through every operand helper. debug_variables: bool, + /// Compiler-only SSA values that deliberately never enter an alloca. + /// + /// High-level semantic handles use this route when their verifier needs + /// direct producer-to-consumer SSA. Ordinary MIR locals continue to use + /// `slots` and later mem2reg promotion. + direct_values: Vec>, /// Per-body unchecked-indexing policy, resolved once by /// [`super::body::translate_body`] from the `__unchecked_indexing_config` /// marker and the `CUDA_OXIDE_UNCHECKED_INDEXING` environment switch. @@ -86,6 +92,7 @@ impl ValueMap { slots: vec![None; num_locals], statement_debug_spills: vec![false; num_locals], debug_variables: false, + direct_values: vec![None; num_locals], unchecked_indexing: false, } } @@ -100,6 +107,20 @@ impl ValueMap { self.debug_variables } + /// Bind a compiler-only local directly to one SSA value. + pub(crate) fn set_direct_value(&mut self, local: mir::Local, value: Value) { + let idx: usize = local; + if idx < self.direct_values.len() { + self.direct_values[idx] = Some(value); + } + } + + /// Read a compiler-only local that bypasses the alloca model. + pub(crate) fn get_direct_value(&self, local: mir::Local) -> Option { + let idx: usize = local; + self.direct_values.get(idx).copied().flatten() + } + /// Record the resolved unchecked-indexing policy for this body. pub fn set_unchecked_indexing(&mut self, enabled: bool) { self.unchecked_indexing = enabled; @@ -537,7 +558,10 @@ impl SlotAddrSpaceMap { } let mir::TerminatorKind::Call { - func, destination, .. + func, + args, + destination, + .. } = &block.terminator.kind else { continue; @@ -545,7 +569,7 @@ impl SlotAddrSpaceMap { if !destination.projection.is_empty() { continue; } - let class = classify_call(func); + let class = classify_call(func, args, &classes); let local_idx: usize = destination.local; let declared = declared_addr_spaces.get(local_idx).copied().flatten(); if let Some(observation) = resolve(class, declared) @@ -615,6 +639,17 @@ fn merge_into( } } +/// Classify a single MIR operand's address space. +fn classify_operand(operand: &mir::Operand, classes: &[SlotAddrSpace]) -> WriteClass { + match operand { + mir::Operand::Copy(place) | mir::Operand::Move(place) if place.projection.is_empty() => { + propagate_from_local(place.local, classes) + } + mir::Operand::Constant(const_op) => classify_constant(const_op), + _ => WriteClass::Unclassified, + } +} + /// Classify the write produced by an `Assign(_, rvalue)` statement. /// /// The rule set is intentionally narrow: when in doubt, return @@ -624,28 +659,22 @@ fn merge_into( /// incomplete classification therefore leaves ordinary pointers generic. fn classify_rvalue(rvalue: &mir::Rvalue, classes: &[SlotAddrSpace]) -> WriteClass { match rvalue { - // `_y = _x` — propagate `_x`'s current classification. `Move` and - // `Copy` are indistinguishable for addrspace purposes. `CopyForDeref` - // behaves the same at this layer. - mir::Rvalue::Use(mir::Operand::Copy(place), _) - | mir::Rvalue::Use(mir::Operand::Move(place), _) - | mir::Rvalue::CopyForDeref(place) - if place.projection.is_empty() => - { + // `_y = _x` or `_y = CONSTANT` — propagate the operand's classification. + mir::Rvalue::Use(operand, _) => classify_operand(operand, classes), + mir::Rvalue::CopyForDeref(place) if place.projection.is_empty() => { propagate_from_local(place.local, classes) } - // `_y = CONSTANT` — a constant-operand pointer to a shared-memory - // static (e.g. `&mut TILE_A` where `TILE_A: SharedArray<...>`) - // lowers to a `mir.shared_alloc` in `translate_operand` whose - // result is `addrspace(3)`. The matching `WriteClass::Classified` - // here keeps the destination slot typed to match, avoiding the - // otherwise-inevitable `PtrToPtr` narrow-to-generic cast. - mir::Rvalue::Use(mir::Operand::Constant(const_op), _) => classify_constant(const_op), - // Every other rvalue shape (aggregates, arithmetic, casts, complex - // projections, `Ref`/`AddressOf`, …) we decline to reason about - // here. The matching `Call`-terminator classifier handles the - // pointer-producing intrinsics; the remaining cases can be - // tightened in a follow-up if a benchmark asks for it. + // `_y = Offset(_ptr, _count)` — pointer arithmetic stays in the base + // pointer's address space (`MirPtrOffsetOp` preserves `left_val`'s type). + mir::Rvalue::BinaryOp(mir::BinOp::Offset, base, _) => classify_operand(base, classes), + // `_y = _ptr as *const U` / `_ptr as *mut U` — pointer-to-pointer and + // mut-to-const casts change pointee type or mutability, not memory space. + mir::Rvalue::Cast( + mir::CastKind::PtrToPtr + | mir::CastKind::PointerCoercion(mir::PointerCoercion::MutToConstPointer), + operand, + _, + ) => classify_operand(operand, classes), _ => WriteClass::Unclassified, } } @@ -712,7 +741,11 @@ pub(super) fn is_constant_wrapper_type(ty: &rustc_public::ty::Ty) -> bool { /// whose emitter unconditionally produces a pointer in a specific address /// space is listed here; new intrinsics should add an entry on the same /// commit that adds their emitter. -fn classify_call(func: &mir::Operand) -> WriteClass { +fn classify_call( + func: &mir::Operand, + args: &[mir::Operand], + classes: &[SlotAddrSpace], +) -> WriteClass { let mir::Operand::Constant(const_op) = func else { return WriteClass::Unclassified; }; @@ -725,6 +758,14 @@ fn classify_call(func: &mir::Operand) -> WriteClass { let path = fn_def.name(); let on_shared_array = self_ty_is_shared_array(&substs); + // `arith_offset(ptr, count)` lowers to `MirPtrOffsetOp` with `ptr`'s + // `MirPtrType` (including its address space). + if path.ends_with("::arith_offset") + && let Some(base) = args.first() + { + return classify_operand(base, classes); + } + // --- addrspace 3 (shared) producers ------------------------------------- // // `SharedArray::index` / `IndexMut::index_mut` on a `SharedArray` diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs index b976d918c1..2b834ff1ae 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/active_mask.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `active_mask` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs index 6539303d91..91d29997dc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/clc.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `clc` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs index 54ea032e58..3acdf9d2fb 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `cluster_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs index 507fba08de..b42ad78914 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cluster_memory.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `cluster_memory` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs index 7e05acca3d..80403e2d20 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/cp_async.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `cp_async` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs index 42096e065a..614449b84e 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/debug_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `debug_control` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs index 8ca4269bd1..4ecaf508a2 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/dotprod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `dotprod` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs index 51544ab2a1..5927a56de9 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/elect.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `elect` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs index cc8c53c24a..00eca17508 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/execution_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `execution_control` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs index 93fa8e6798..296b6dd69c 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/extended_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `extended_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs index 0d99dc692e..fbdd7b11af 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/integer_minmax.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `integer_minmax` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs index be52262fd0..85202fb799 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/ldmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `ldmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs index 03aec87416..38c3c3570a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_basic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `mbarrier_basic` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs index f75d9b824c..c18e2051fa 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mbarrier_extended.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `mbarrier_extended` intrinsics. @@ -14,10 +14,12 @@ use crate::convert::intrinsics::{ common::cast_to_shared_addrspace, common::inline_asm_convergent, common::trunc_to_i1, }; use dialect_nvvm::ops::{ - FenceMbarrierInitReleaseClusterOp, FenceProxyAsyncGenericAcquireSharedClusterClusterOp, + FenceAcquireSharedClusterClusterOp, FenceMbarrierInitReleaseClusterOp, + FenceProxyAsyncGenericAcquireSharedClusterClusterOp, FenceProxyAsyncGenericReleaseSharedCtaClusterOp, FenceProxyAsyncSharedCtaOp, - MbarrierArriveClusterOp, MbarrierArriveExpectTxClusterOp, MbarrierArriveExpectTxSharedOp, - MbarrierTryWaitParityClusterOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, + MbarrierArriveClusterOp, MbarrierArriveClusterRelaxedOp, MbarrierArriveExpectTxClusterOp, + MbarrierArriveExpectTxSharedOp, MbarrierTryWaitParityClusterOp, + MbarrierTryWaitParityClusterRelaxedOp, MbarrierTryWaitParitySharedOp, MbarrierTryWaitSharedOp, NanosleepOp, }; use llvm_export::types as llvm_types; @@ -30,6 +32,37 @@ use pliron::op::Op; use pliron::result::Result; use pliron::value::DefiningEntity; +#[op_interface_impl] +impl MirToLlvmConversion for FenceAcquireSharedClusterClusterOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + let op = self.get_operation(); + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if !operands.is_empty() { + return pliron::input_err_noloc!( + "fence_acquire_shared_cluster_cluster requires 0 operands, got {}", + operands.len() + ); + } + let void_ty = llvm_types::VoidType::get(ctx); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + operands, + "fence.acquire.sync_restrict::shared::cluster.cluster;", + "~{memory}", + ); + rewriter.erase_operation(ctx, op); + Ok(()) + } +} + #[op_interface_impl] impl MirToLlvmConversion for FenceMbarrierInitReleaseClusterOp { fn convert( @@ -185,6 +218,37 @@ impl MirToLlvmConversion for MbarrierArriveClusterOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for MbarrierArriveClusterRelaxedOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + let op = self.get_operation(); + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.len() != 1 { + return pliron::input_err_noloc!( + "mbarrier_arrive_cluster_relaxed requires 1 operands, got {}", + operands.len() + ); + } + let void_ty = llvm_types::VoidType::get(ctx); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + operands, + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [$0];", + "l,~{memory}", + ); + rewriter.erase_operation(ctx, op); + Ok(()) + } +} + #[op_interface_impl] impl MirToLlvmConversion for MbarrierArriveExpectTxSharedOp { fn convert( @@ -360,6 +424,43 @@ impl MirToLlvmConversion for MbarrierTryWaitParityClusterOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for MbarrierTryWaitParityClusterRelaxedOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + let op = self.get_operation(); + let mut operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.len() != 2 { + return pliron::input_err_noloc!( + "mbarrier_try_wait_parity_cluster_relaxed requires 2 operands, got {}", + operands.len() + ); + } + operands[0] = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let result_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let asm = inline_asm_convergent( + ctx, + rewriter, + op, + result_ty.into(), + operands, + "{ .reg .pred %p0; mbarrier.try_wait.parity.relaxed.cluster.shared::cta.b64 %p0, [$1], $2; selp.b32 $0, 1, 0, %p0; }", + "=r,l,r,~{memory}", + ); + let asm_result = asm.deref(ctx).get_result(0); + let result = trunc_to_i1(ctx, rewriter, asm_result); + let DefiningEntity::Op(result_op) = result.defining_entity() else { + unreachable!() + }; + rewriter.replace_operation(ctx, op, result_op); + Ok(()) + } +} + #[op_interface_impl] impl MirToLlvmConversion for NanosleepOp { fn convert( diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs index e9867a4cb4..46e6d3bcb3 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/mod.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces for admitted CUDA intrinsic families. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs index 25e0ea09b0..c30cf9ecfc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/movmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `movmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs index 81d979e5f5..567e6efcc1 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_alu.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_alu` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs index 39b38cc87e..0d841cc1d5 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_atomic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_atomic` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs index 45cd3c17b2..a24c7b1201 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/packed_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs index 55df006d2b..47c635e498 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/prmt.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `prmt` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs index 5a9360d2cd..b62c7ad0fc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/redux.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `redux` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs index 8b9ebfa2fe..0787369606 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/register_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `register_mma` intrinsics. @@ -2137,6 +2137,23 @@ impl MirToLlvmConversion for RegisterMmaOp { "mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32 {$0, $1, $2, $3}, {$8, $9}, {$10}, {$4, $5, $6, $7};", "=f,=f,=f,=f,f,f,f,f,r,r,r", ), + ( + Some(&RegisterMmaShapeAttr::M16n8k64), + RegisterMmaOperationAttr::Multiply, + RegisterMmaKindAttr::Mxf4, + Some(&RegisterMmaAccumulatorAttr::F32), + Some(&RegisterMmaElementAttr::E2m1), + Some(&RegisterMmaElementAttr::E2m1), + Some(&RegisterMmaLayoutAttr::Row), + Some(&RegisterMmaLayoutAttr::Col), + Some(&RegisterMmaOverflowAttr::NotApplicable), + ) => ( + GeneratedMmaResultType::F32, + 4, + 16, + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X.f32.e2m1.e2m1.f32.ue8m0 {$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7}, $14, {$15, $16}, $17, {$18, $19};", + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r,r,h,h,r,h,h", + ), ( Some(&RegisterMmaShapeAttr::M16n8k64), RegisterMmaOperationAttr::Multiply, diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs index 487a60358e..168a5537ea 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_arithmetic.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_arithmetic` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs index 7f78ccbed3..6123726891 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_conversion.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_conversion` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs index b1f3baf883..bcb2a15e8d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/scalar_math.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_math` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs index d608302717..aab2bcf1b0 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sparse_mma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `sparse_mma` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs index 07548cf648..9c1c4ff210 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sreg.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `sreg` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs index b28daeaad8..aab5fd1b9e 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/stmatrix.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `stmatrix` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs index 1f213168c2..30799d126f 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/sync.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `sync` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs index 4960b066f9..8bb2bf6c96 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tcgen05.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tcgen05` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs index 50f31b2099..567ddbb37d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/tma.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tma` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs index af0c94c463..3252a6bd1a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/vote.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `vote` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs index 8d693df970..1f3d4d8d97 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_barrier.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `warp_barrier` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs index 6972ee7b7b..2acc2f26ef 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_match.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `warp_match` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs index ff65fb5c25..1da375e4d8 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/warp_shuffle.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `warp_shuffle` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs index 7dcfa3d705..4aae07bf38 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/generated_intrinsics/wgmma_control.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `wgmma_control` intrinsics. diff --git a/cuda-oxide/crates/mir-lower/src/convert/interface_impls.rs b/cuda-oxide/crates/mir-lower/src/convert/interface_impls.rs index 41e1ee694b..a593ef9256 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/interface_impls.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/interface_impls.rs @@ -37,15 +37,15 @@ use dialect_mir::ops::{ MirStorageLiveOp, MirStoreOp, MirSubOp, MirUndefOp, MirUnreachableOp, MirUnrollHintOp, }; use dialect_nvvm::ops::{ - AssertFailOp, CvtaGenericToSharedOffsetOp, InlinePtxOp, NvvmAtomicCmpxchgOp, NvvmAtomicFenceOp, - NvvmAtomicLoadOp, NvvmAtomicRmwOp, NvvmAtomicStoreOp, ReadPtxSregClusterIdxOp, - ReadPtxSregNclusterIdOp, VprintfOp, WgmmaMakeSmemDescOp, WgmmaMmaGroupM64N64K16F32Bf16Op, - WgmmaMmaGroupValuesM64N64K8F32Tf32Op, WgmmaMmaGroupValuesM64N64K16F32Bf16Op, - WgmmaMmaGroupValuesM64N64K16F32F16Op, WgmmaMmaGroupValuesM64N128K16F32Bf16Op, - WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32Bf16Op, - WgmmaMmaLoopValuesM64N64K16F32F16Op, WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, - WgmmaMmaM64N64K16F32F16Op, WgmmaMmaM64N128K16F32Bf16Op, - WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, + AssertFailOp, CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp, InlinePtxOp, + NvvmAtomicCmpxchgOp, NvvmAtomicFenceOp, NvvmAtomicLoadOp, NvvmAtomicRmwOp, NvvmAtomicStoreOp, + ReadPtxSregClusterIdxOp, ReadPtxSregNclusterIdOp, VprintfOp, WgmmaMakeSmemDescOp, + WgmmaMmaGroupM64N64K16F32Bf16Op, WgmmaMmaGroupValuesM64N64K8F32Tf32Op, + WgmmaMmaGroupValuesM64N64K16F32Bf16Op, WgmmaMmaGroupValuesM64N64K16F32F16Op, + WgmmaMmaGroupValuesM64N128K16F32Bf16Op, WgmmaMmaLoopPipelineValuesM64N64K16F32Bf16Op, + WgmmaMmaLoopValuesM64N64K16F32Bf16Op, WgmmaMmaLoopValuesM64N64K16F32F16Op, + WgmmaMmaM64N64K8F32Tf32Op, WgmmaMmaM64N64K16F32Bf16Op, WgmmaMmaM64N64K16F32F16Op, + WgmmaMmaM64N128K16F32Bf16Op, WgmmaMmaPipelineValuesM64N64K16F32Bf16Op, }; // ---- Arithmetic ops -------------------------------------------------------- @@ -1043,6 +1043,23 @@ impl MirToLlvmConversion for CvtaGenericToSharedOffsetOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CvtaGenericToClusterSharedOffsetOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + super::intrinsics::memory::convert_cvta_generic_to_cluster_shared_offset( + ctx, + rewriter, + self.get_operation(), + operands_info, + ) + } +} + // ---- NVVM WGMMA ops -------------------------------------------------------- #[op_interface_impl] diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs index 6b947c3c40..344f0acfc0 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs @@ -155,10 +155,28 @@ pub fn inline_asm_convergent( rewriter: &mut DialectConversionRewriter, current_op: Ptr, result_ty: pliron::r#type::TypeHandle, - inputs: Vec, + mut inputs: Vec, asm_template: &str, constraints: &str, ) -> Ptr { + if ((asm_template.contains("cvta.to.shared.u64 ") && asm_template.contains(", $0;")) + || asm_template.contains("tcgen05.commit.")) + && let Some(first) = inputs.first_mut() + { + let first_ty = first.get_type(ctx); + let addrspace = first_ty + .deref(ctx) + .downcast_ref::() + .map(|pt| pt.address_space()); + if let Some(space) = addrspace + && space != 0 + { + let generic_ty = llvm_types::PointerType::get(ctx, 0).into(); + let cast_op = llvm::AddrSpaceCastOp::new(ctx, *first, generic_ty); + rewriter.insert_operation(ctx, cast_op.get_operation()); + *first = cast_op.get_operation().deref(ctx).get_result(0); + } + } let inline_asm = llvm::InlineAsmOp::build( ctx, result_ty, @@ -242,6 +260,7 @@ mod tests { ClusterSharedCast, CallIntrinsic, InlineAsmConvergent, + InlineAsmConvergentGenericPtr(&'static str), InlineAsmSideEffect, TruncToI1, } @@ -311,6 +330,20 @@ mod tests { ); None } + HelperAction::InlineAsmConvergentGenericPtr(template) => { + let void_ty = llvm_types::VoidType::get(ctx); + let ptr = block.deref(ctx).get_argument(0); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + vec![ptr], + template, + "l,~{memory}", + ); + None + } HelperAction::InlineAsmSideEffect => { let void_ty = llvm_types::VoidType::get(ctx); let dst = block.deref(ctx).get_argument(0); @@ -731,4 +764,42 @@ mod tests { ); assert_eq!(integer_width(&ctx, result_ty), 1); } + + #[test] + fn inline_asm_convergent_normalizes_non_generic_first_pointer_for_generic_templates() { + for template in [ + "{ .reg .u64 %shared64; cvta.to.shared.u64 %shared64, $0; }", + "tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [$0], 3;", + ] { + let mut ctx = make_ctx(); + let shared_ptr_ty: TypeHandle = llvm_types::PointerType::get(&ctx, 3).into(); + let (module_ptr, entry) = build_test_func(&mut ctx, vec![shared_ptr_ty]); + append_trigger(&mut ctx, entry); + + assert!( + run_helper_action( + &mut ctx, + module_ptr, + HelperAction::InlineAsmConvergentGenericPtr(template), + ) + .is_none() + ); + + let casts = find_body_ops::(&ctx, module_ptr); + assert_eq!(casts.len(), 1); + let cast_result = casts[0].get_operation().deref(&ctx).get_result(0); + assert_eq!(pointer_addr_space(&ctx, cast_result.get_type(&ctx)), 0); + + let asms = find_body_ops::(&ctx, module_ptr); + assert_eq!(asms.len(), 1); + assert_eq!( + asms[0] + .get_operation() + .deref(&ctx) + .operands() + .collect::>(), + vec![cast_result] + ); + } + } } diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs index f5ae10e694..2d83d61489 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs @@ -12,7 +12,9 @@ //! boundary, not a Rust pointer-address observation (those genericize first, //! see `convert/ops/cast.rs`). -use crate::convert::intrinsics::common::cast_to_shared_addrspace; +use crate::convert::intrinsics::common::{ + cast_to_cluster_shared_addrspace, cast_to_shared_addrspace, +}; use llvm_export::op_interfaces::CastOpInterface; use llvm_export::ops as llvm; use pliron::builtin::types::{IntegerType, Signedness}; @@ -36,7 +38,59 @@ pub(crate) fn convert_cvta_generic_to_shared_offset( if operands.is_empty() { return pliron::input_err_noloc!("cvta_generic_to_shared_offset requires an operand"); } - let shared_ptr = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let space = operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|ty| ty.address_space()); + let shared_ptr = match space { + Some(3 | 7) => operands[0], // AS7 retains the remote CTA rank. + Some(0) => cast_to_shared_addrspace(ctx, rewriter, operands[0]), + _ => { + return pliron::input_err_noloc!( + "shared offset requires generic, shared or cluster shared pointer" + ); + } + }; + + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let Some(result_width) = result_ty + .deref(ctx) + .downcast_ref::() + .map(IntegerType::width) + else { + return pliron::input_err_noloc!( + "cvta_generic_to_shared_offset requires an integer result" + ); + }; + let result_ty = IntegerType::get(ctx, result_width, Signedness::Signless); + let ptr_to_int = llvm::PtrToIntOp::new(ctx, shared_ptr, result_ty.into()); + rewriter.insert_operation(ctx, ptr_to_int.get_operation()); + rewriter.replace_operation(ctx, op, ptr_to_int.get_operation()); + Ok(()) +} + +pub(crate) fn convert_cvta_generic_to_cluster_shared_offset( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.is_empty() { + return pliron::input_err_noloc!("cvta_generic_to_shared_offset requires an operand"); + } + let space = operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|ty| ty.address_space()); + if !matches!(space, Some(0 | 3 | 7)) { + return pliron::input_err_noloc!( + "cluster shared offset requires generic or shared pointer" + ); + } + let shared_ptr = cast_to_cluster_shared_addrspace(ctx, rewriter, operands[0]); let result_ty = op.deref(ctx).get_result(0).get_type(ctx); let Some(result_width) = result_ty @@ -54,3 +108,96 @@ pub(crate) fn convert_cvta_generic_to_shared_offset( rewriter.replace_operation(ctx, op, ptr_to_int.get_operation()); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::convert::ops::test_util::{ + append_mir_return, build_kernel, find_all, kernel_blocks, make_ctx, + }; + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::{CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp}; + use pliron::r#type::TypeHandle; + #[test] + fn cvta_preserves_local_and_cluster_offsets() { + for space in [0, 3, 7] { + for width in [32, 64] { + let mut ctx = make_ctx(); + let word: TypeHandle = IntegerType::get(&ctx, width, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, word, true, space).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr], vec![word]); + let arg = block.deref(&ctx).get_argument(0); + let cvta = Operation::new( + &mut ctx, + CvtaGenericToSharedOffsetOp::get_concrete_op_info(), + vec![word], + vec![arg], + vec![], + 0, + ); + cvta.insert_at_back(block, &ctx); + let value = cvta.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![value]); + crate::lower_mir_to_llvm(&mut ctx, module).unwrap(); + let blocks = kernel_blocks(&ctx, module); + let casts = find_all::(&ctx, &blocks); + let ints = find_all::(&ctx, &blocks); + assert_eq!(casts.len(), usize::from(space == 0)); + assert_eq!(ints.len(), 1); + let operand = ints[0].get_operation().deref(&ctx).get_operand(0); + let ty = operand.get_type(&ctx); + assert_eq!( + ty.deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + if space == 0 { 3 } else { space } + ); + if space == 0 { + assert_eq!(operand, casts[0].get_operation().deref(&ctx).get_result(0)); + } + } + } + } + #[test] + fn cluster_cvta_targets_cluster_window_from_generic_and_named_pointers() { + for space in [0, 3, 7] { + for width in [32, 64] { + let mut ctx = make_ctx(); + let word: TypeHandle = IntegerType::get(&ctx, width, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, word, true, space).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr], vec![word]); + let arg = block.deref(&ctx).get_argument(0); + let cvta = Operation::new( + &mut ctx, + CvtaGenericToClusterSharedOffsetOp::get_concrete_op_info(), + vec![word], + vec![arg], + vec![], + 0, + ); + cvta.insert_at_back(block, &ctx); + let value = cvta.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![value]); + crate::lower_mir_to_llvm(&mut ctx, module).unwrap(); + let blocks = kernel_blocks(&ctx, module); + let casts = find_all::(&ctx, &blocks); + let ints = find_all::(&ctx, &blocks); + assert_eq!(casts.len(), usize::from(space != 7)); + assert_eq!(ints.len(), 1); + let operand = ints[0].get_operation().deref(&ctx).get_operand(0); + let ty = operand.get_type(&ctx); + assert_eq!( + ty.deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + 7 + ); + if space != 7 { + assert_eq!(operand, casts[0].get_operation().deref(&ctx).get_result(0)); + } + } + } + } +} diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs index 66a51c81eb..0b1726cb53 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs @@ -85,17 +85,13 @@ fn pointer_address_space(ctx: &Context, value: pliron::value::Value) -> Option`, which lowers to an `addrspace(3)` pointer, while the raw -/// ABI entry still passes a generic pointer under the caller's promise of -/// CTA locality. The barrier is a caller promise on both layers, as for every -/// mbarrier operation. Generic and shared pointers are converted to the +/// Both destination and barrier are caller promises of CTA locality, as for +/// every mbarrier operation. Generic and shared pointers are converted to the /// shared window. A pointer known to be cluster shared (`addrspace(7)`, for /// example from `map_shared_rank`) may name another CTA's shared memory, so /// it is rejected rather than reinterpreted as a local offset; such copies /// use the cluster form. The check sees only uncast operands: a pointer cast -/// or `SharedPtr::from_raw` applied to a cluster pointer is the caller's -/// unsafe contract. +/// applied to a cluster pointer is the caller's unsafe contract. pub(crate) fn convert_g2s_cta( ctx: &mut Context, rewriter: &mut DialectConversionRewriter, diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs index 94aced0846..1f4d7fed22 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs @@ -1043,8 +1043,8 @@ mod tests { let body = kernel_blocks(&ctx, module); assert_eq!( count_ops::(&ctx, &body), - 2, - "construction must genericize the pointer and extraction must restore shared space" + 3, + "construction must genericize the pointer, extraction must restore shared space, and return must genericize for .func ABI" ); assert_eq!( count_ops::(&ctx, &body), diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs index e60a2878d1..356208f099 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs @@ -505,8 +505,8 @@ mod tests { assert_eq!( count_ops::(&ctx, &body), - 2, - "union insertion must genericize AS3 and extraction must restore AS3" + 3, + "union insertion must genericize AS3, extraction must restore AS3, and return must genericize for .func ABI" ); assert_eq!( count_ops::(&ctx, &body), @@ -610,8 +610,8 @@ mod tests { assert_eq!( count_ops::(&ctx, &body), - 2, - "nested union insertion must genericize AS3 and extraction must restore AS3" + 3, + "nested union insertion must genericize AS3, extraction must restore AS3, and return must genericize for .func ABI" ); assert_eq!(count_ops::(&ctx, &body), 0); assert_eq!(count_ops::(&ctx, &body), 0); diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs index bc19168b0c..c5ae9f0325 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs @@ -64,7 +64,7 @@ use crate::convert::target_stable_storage::coerce_target_stable_value; use crate::convert::types::{ StructLayoutInfo, TransparentScalarAbiInfo, build_struct_slot_map, convert_function_type, - convert_type, is_kernel_func, is_zero_sized_type, packed_shared_internal_abi_info, + convert_type, is_kernel_func, is_zero_sized_type, shared_return_abi_info, transparent_scalar_abi_info, }; use crate::helpers; @@ -680,9 +680,7 @@ pub fn convert( let scalar_ty = abi.scalar_ty; transparent_result_abi = Some(abi); scalar_ty - } else if let Some(abi) = - packed_shared_internal_abi_info(ctx, mir_ty).map_err(anyhow_to_pliron)? - { + } else if let Some(abi) = shared_return_abi_info(ctx, mir_ty).map_err(anyhow_to_pliron)? { let storage_ty = abi.storage_ty; packed_shared_result_abi = Some(abi); storage_ty diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs index 3b3b5b799f..90c8721126 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs @@ -1523,11 +1523,19 @@ mod tests { ); // The exposed integer is a generic address, so re-entering the shared - // space must go inttoptr → generic, then addrspacecast → shared. A - // bare inttoptr into addrspace(3) would reinterpret the generic - // address as a space-local offset. - assert_cast_lowered_to::(&ctx, module, "llvm.addrspacecast"); + // space must go inttoptr -> generic, then addrspacecast -> shared, + // followed by the .func return-boundary addrspacecast back to generic. let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "expected cast addrspacecast plus .func return-boundary addrspacecast" + ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "mir.cast must be replaced during lowering" + ); let int_to_ptrs = find_all::(&ctx, &body); let [int_to_ptr] = int_to_ptrs.as_slice() else { panic!("integer-to-shared transmute must lower to exactly one llvm.inttoptr"); @@ -1564,8 +1572,14 @@ mod tests { let module = lower_single_cast(&mut ctx, source, destination, MirCastKindAttr::Transmute); - assert_cast_lowered_to::(&ctx, module, "llvm.addrspacecast"); let body = kernel_blocks(&ctx, module); + let expected_casts = if shared_is_source { 1 } else { 2 }; + assert_eq!( + count_ops::(&ctx, &body), + expected_casts, + "expected cast addrspacecast (plus return-boundary cast when returning shared)" + ); + assert_eq!(count_ops::(&ctx, &body), 0); assert_eq!( count_ops::(&ctx, &body), 0, @@ -1743,6 +1757,7 @@ mod tests { llvm_export::types::address_space::SHARED, llvm_export::types::address_space::CONSTANT, llvm_export::types::address_space::LOCAL, + 7, // cluster shared ] { let mut ctx = make_ctx(); let pointee_ty = int_ty(&mut ctx, 32, Signedness::Signless); @@ -1922,6 +1937,7 @@ mod tests { llvm_export::types::address_space::SHARED, llvm_export::types::address_space::CONSTANT, llvm_export::types::address_space::LOCAL, + 7, // cluster shared ] { let mut ctx = make_ctx(); let usize_ty = int_ty(&mut ctx, 64, Signedness::Unsigned); @@ -1944,17 +1960,28 @@ mod tests { // The exposed integer is a generic address, so re-entry must be // inttoptr-to-generic followed by addrspacecast into the space, - // mirroring the PointerExposeAddress direction. - assert_cast_lowered_to::(&ctx, module_ptr, "llvm.addrspacecast"); + // mirroring the PointerExposeAddress direction. When returning an + // addrspace(3) pointer, the .func return boundary adds a second + // addrspacecast back to generic. + let body = kernel_blocks(&ctx, module_ptr); + let expected_casts = if address_space == llvm_export::types::address_space::SHARED { + 2 + } else { + 1 + }; assert_eq!( - count_ops::(&ctx, &kernel_blocks(&ctx, module_ptr)), + count_ops::(&ctx, &body), + expected_casts, + "expected cast addrspacecast (plus return-boundary cast for shared)" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!( + count_ops::(&ctx, &body), 1, "address space {address_space} must re-enter through a generic inttoptr" ); - let casts = find_all::(&ctx, &kernel_blocks(&ctx, module_ptr)); - let [cast] = casts.as_slice() else { - panic!("expected exactly one llvm.addrspacecast"); - }; + let casts = find_all::(&ctx, &body); + let cast = &casts[0]; let result_ty = cast .get_operation() .deref(&ctx) @@ -1982,12 +2009,16 @@ mod tests { MirCastKindAttr::PtrToPtr, ); - assert_cast_lowered_to::(&ctx, module_ptr, "llvm.addrspacecast"); + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "ptr -> shared cast plus .func return-boundary cast must emit two llvm.addrspacecast ops" + ); + assert_eq!(count_ops::(&ctx, &body), 0); - let casts = find_all::(&ctx, &kernel_blocks(&ctx, module_ptr)); - let [cast] = casts.as_slice() else { - panic!("ptr -> ptr addrspace cast must lower to exactly one llvm.addrspacecast"); - }; + let casts = find_all::(&ctx, &body); + let cast = &casts[0]; let result_ty = cast .get_operation() .deref(&ctx) diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs index bf31f820cf..0acd93bb8b 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs @@ -24,7 +24,9 @@ //! not copied). Successor pointers are already valid — no block map lookup needed. use crate::convert::target_stable_storage::coerce_target_stable_value; -use crate::convert::types::packed_shared_internal_abi_info; +use crate::convert::types::{ + is_kernel_func, packed_shared_internal_abi_info, shared_return_abi_info, +}; use dialect_mir::types::MirStructType; use llvm_export::ops as llvm; use pliron::basic_block::BasicBlock; @@ -81,8 +83,17 @@ pub(crate) fn convert_return( } Some(scalar) } else { + let is_kernel = op + .deref(ctx) + .get_parent_op(ctx) + .is_some_and(|parent| is_kernel_func(ctx, parent)); let packed_shared_abi = if let Some(mir_ty) = mir_ty { - match packed_shared_internal_abi_info(ctx, mir_ty) { + let abi_res = if is_kernel { + packed_shared_internal_abi_info(ctx, mir_ty) + } else { + shared_return_abi_info(ctx, mir_ty) + }; + match abi_res { Ok(abi) => abi, Err(error) => { return pliron::input_err_noloc!( @@ -304,6 +315,9 @@ pub(crate) fn convert_goto( } let llvm_br = llvm::BrOp::new(ctx, dest, final_args); + if dialect_mir::ops::MirGotoOp::new(op).is_loop_unroll_disable(ctx) { + llvm::set_loop_unroll_disable(ctx, llvm_br.get_operation()); + } crate::convert::preserve_location(ctx, op, llvm_br.get_operation()); rewriter.insert_operation(ctx, llvm_br.get_operation()); rewriter.erase_operation(ctx, op); diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs b/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs index 3c60bbf1b8..55c1a5112b 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs @@ -349,6 +349,37 @@ pub(crate) fn packed_shared_internal_abi_info( })) } +/// Recognize an internal device-function return type that must cross the `.func` +/// boundary as a target-stable 64-bit generic pointer (or generic-pointer carrier). +/// +/// Under `--nvptx-short-ptr` (`p3:32:32:32`), LLVM 23's `NVPTXAsmPrinter` emits +/// `.func (.param .b64 func_retval0)` in the function header for a bare +/// `ptr addrspace(3)` return value while emitting `.param .b32 retval0` at the +/// call site. Returning `ptr` (`addrspace(0)`) across the `.func` boundary and +/// coercing back to `ptr addrspace(3)` at the call site matches `nvcc` and +/// avoids the PTX return-slot width mismatch. +pub(crate) fn shared_return_abi_info( + ctx: &mut Context, + mir_ty: TypeHandle, +) -> Result, anyhow::Error> { + if let Some(abi) = packed_shared_internal_abi_info(ctx, mir_ty)? { + return Ok(Some(abi)); + } + let is_shared_ptr = mir_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|ptr_ty| ptr_ty.address_space() == llvm_types::address_space::SHARED); + if !is_shared_ptr { + return Ok(None); + } + let semantic_ty = convert_type(ctx, mir_ty)?; + let storage_ty = llvm_types::PointerType::get_generic(ctx).into(); + Ok(Some(PackedSharedInternalAbiInfo { + semantic_ty, + storage_ty, + })) +} + /// Convert a type that crosses a function boundary as one LLVM value. /// /// A struct with a natural-layout divergence is legal by value only when @@ -621,7 +652,7 @@ pub fn convert_function_type( let ty = if is_transparent_scalar { transparent_scalar_llvm_type(ctx, mir_ret_ty)? } else if !is_kernel_entry { - match packed_shared_internal_abi_info(ctx, mir_ret_ty)? { + match shared_return_abi_info(ctx, mir_ret_ty)? { Some(abi) => abi.storage_ty, None => convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")?, } diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs index af6089f3cc..f68c950342 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs @@ -117,7 +117,7 @@ pub(crate) use enum_layout::{ }; pub(crate) use func_abi::{ TransparentScalarAbiInfo, convert_grid_constant_storage_type, packed_shared_internal_abi_info, - transparent_scalar_abi_info, transparent_scalar_field, + shared_return_abi_info, transparent_scalar_abi_info, transparent_scalar_field, }; pub use func_abi::{convert_function_type, is_kernel_func}; // Surface parity with the pre-split `types.rs`: these four are currently diff --git a/cuda-oxide/crates/mir-transforms/src/scalarize_borrowed_aggregate_reads.rs b/cuda-oxide/crates/mir-transforms/src/scalarize_borrowed_aggregate_reads.rs index dbc2140224..3bd6a9d403 100644 --- a/cuda-oxide/crates/mir-transforms/src/scalarize_borrowed_aggregate_reads.rs +++ b/cuda-oxide/crates/mir-transforms/src/scalarize_borrowed_aggregate_reads.rs @@ -35,8 +35,10 @@ //! //! The pre-mem2reg phase fails closed on pointer provenance and mutation. It //! rejects additional stores, volatile loads, calls, pointer casts, pointer -//! PHIs/selects, unknown users, non-array fields, and projections in the entry -//! block before the initializer can be proven to dominate them. Derived +//! PHIs/selects and unknown users. Same-block argument snapshots additionally +//! support nested projections and whole-aggregate copies, only after a unique +//! initializing store. The older cross-block path retains its one-field array +//! restriction. Derived //! pointers remain mutable `Erased` storage addresses because projections //! preserve the alloca's carrier; the complete use-graph proof, rather than a //! false immutable type, establishes that this particular access is read-only. @@ -139,6 +141,11 @@ pub fn canonicalize_read_only_aggregate_arguments( let mut rewritten_loads = 0usize; for alloca in allocas { + let snapshot_reads = rewrite_entry_snapshot(ctx, alloca); + rewritten_loads += snapshot_reads; + if snapshot_reads != 0 { + continue; + } let Some(plan) = analyze_alloca(ctx, alloca) else { continue; }; @@ -150,6 +157,169 @@ pub fn canonicalize_read_only_aggregate_arguments( } } +#[derive(Clone)] +enum SnapshotProjection { + Field(u32, TypeHandle), + Element(Value, TypeHandle), +} + +/// A bounded extension for nested, by-value argument snapshots. Every access +/// must stay in the initializer block, after its unique store, and the entire +/// pointer-use graph must consist of projections and nonvolatile reads. This +/// cannot widen reads from shared/global memory or allow a pointer to escape. +fn rewrite_entry_snapshot(ctx: &mut Context, alloca: Ptr) -> usize { + fn argument_root(ctx: &Context, value: Value, args: &[Value]) -> bool { + if args.contains(&value) { + return true; + } + let Some(op) = value.defining_op() else { + return false; + }; + if Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() + { + return argument_root(ctx, op.deref(ctx).get_operand(0), args); + } + false + } + type Read = (Ptr, Vec); + fn walk( + ctx: &Context, + pointer: Value, + initializer: Ptr, + positions: &HashMap, usize>, + path: &[SnapshotProjection], + reads: &mut Vec, + projections: &mut Vec>, + ) -> Option<()> { + let ty = pointer.get_type(ctx); + if !is_mutable_erased_alloca_projection(ty.deref(ctx).downcast_ref::()?) { + return None; + } + for usage in pointer.uses(ctx) { + let user = usage.user_op(); + if user == initializer && path.is_empty() { + continue; + } + if usage.find_index(ctx) != 0 || positions.get(&user)? <= positions.get(&initializer)? { + return None; + } + if let Some(load) = Operation::get_op::(user, ctx) { + if load.is_volatile(ctx) { + return None; + } + reads.push((user, path.to_vec())); + continue; + } + if Operation::get_op::(user, ctx).is_none() + && Operation::get_op::(user, ctx).is_none() + { + return None; + } + let result = user.deref(ctx).get_result(0); + let result_ty = result.get_type(ctx); + let pointee = result_ty.deref(ctx).downcast_ref::()?.pointee; + let projection = if let Some(field) = Operation::get_op::(user, ctx) { + SnapshotProjection::Field(field.get_attr_field_index(ctx)?.0, pointee) + } else if Operation::get_op::(user, ctx).is_some() { + SnapshotProjection::Element(user.deref(ctx).get_operand(1), pointee) + } else { + return None; + }; + let mut path = path.to_vec(); + path.push(projection); + projections.push(user); + walk( + ctx, + result, + initializer, + positions, + &path, + reads, + projections, + )?; + } + Some(()) + } + let Some(block) = alloca.deref(ctx).get_parent_block() else { + return 0; + }; + let positions: HashMap<_, _> = block + .deref(ctx) + .iter(ctx) + .enumerate() + .map(|(i, o)| (o, i)) + .collect(); + let args: Vec<_> = block.deref(ctx).arguments().collect(); + let root = alloca.deref(ctx).get_result(0); + let stores: Vec<_> = root + .uses(ctx) + .into_iter() + .filter_map(|u| Operation::get_op::(u.user_op(), ctx)) + .collect(); + if stores.len() != 1 || stores[0].is_volatile(ctx) { + return 0; + } + let initializer = stores[0].get_operation(); + if !positions.contains_key(&initializer) || initializer.deref(ctx).get_operand(0) != root { + return 0; + } + let value = stores[0].value_opd(ctx); + if !argument_root(ctx, value, &args) { + return 0; + } + let mut reads = vec![]; + let mut projections = vec![]; + if walk( + ctx, + root, + initializer, + &positions, + &[], + &mut reads, + &mut projections, + ) + .is_none() + { + return 0; + } + let count = reads.len(); + let mut rewriter = IRRewriter::::default(); + for (load, path) in reads { + let mut replacement = value; + for projection in path { + let (info, ty, operands) = match projection { + SnapshotProjection::Field(_, ty) => ( + MirExtractFieldOp::get_concrete_op_info(), + ty, + vec![replacement], + ), + SnapshotProjection::Element(index, ty) => ( + MirExtractArrayElementOp::get_concrete_op_info(), + ty, + vec![replacement, index], + ), + }; + let extract = Operation::new(ctx, info, vec![ty], operands, vec![], 0); + if let SnapshotProjection::Field(index, _) = projection { + MirExtractFieldOp::new(extract).set_attr_index(ctx, FieldIndexAttr(index)); + } + extract + .deref_mut(ctx) + .set_loc(load.deref(ctx).loc().clone()); + extract.insert_before(ctx, load); + replacement = extract.deref(ctx).get_result(0); + } + let old = load.deref(ctx).get_result(0); + old.replace_all_uses_with(ctx, &replacement); + rewriter.erase_operation(ctx, load); + } + for projection in projections.into_iter().rev() { + rewriter.erase_operation(ctx, projection); + } + count +} + fn collect_ops(ctx: &Context, root: Ptr, output: &mut Vec>) { output.push(root); let regions: Vec<_> = root.deref(ctx).regions().collect(); @@ -984,6 +1154,83 @@ mod tests { .count() } + #[test] + fn same_block_argument_snapshot_reads_are_forwarded_only_after_unique_initialization() { + for (extra, volatile, write, before, escape, expected) in [ + (false, false, false, false, false, 0), + (true, false, false, false, false, 1), + (false, true, false, false, false, 1), + (false, false, true, false, false, 1), + (false, false, false, true, false, 1), + (false, false, false, false, true, 1), + ] { + let mut ctx = Context::new(); + let fixture = build_fixture(&mut ctx, 32, None, extra, volatile, write); + let entry = fixture.alloca.deref(&ctx).get_parent_block().unwrap(); + let mut ops = vec![]; + collect_ops(&ctx, fixture.module, &mut ops); + let goto = *ops + .iter() + .find(|&&o| Operation::get_op::(o, &ctx).is_some()) + .unwrap(); + let store = *ops + .iter() + .find(|&&o| Operation::get_op::(o, &ctx).is_some()) + .unwrap(); + let body = *ops + .iter() + .find(|&&o| Operation::get_op::(o, &ctx).is_some()) + .unwrap(); + let body = body.deref(&ctx).get_parent_block().unwrap(); + let moved: Vec<_> = body + .deref(&ctx) + .iter(&ctx) + .filter(|&o| Operation::get_op::(o, &ctx).is_none()) + .collect(); + for op in moved { + op.unlink(&ctx); + op.insert_before(&ctx, if before { store } else { goto }); + } + assert_eq!(store.deref(&ctx).get_parent_block(), Some(entry)); + let load = *ops + .iter() + .find(|&&o| Operation::get_op::(o, &ctx).is_some()) + .unwrap(); + let value = load.deref(&ctx).get_result(0); + let ty = value.get_type(&ctx); + let use_value = Operation::new( + &mut ctx, + dialect_mir::ops::MirAddOp::get_concrete_op_info(), + vec![ty], + vec![value, value], + vec![], + 0, + ); + use_value.insert_after(&ctx, load); + if escape { + let pointer = fixture.alloca.deref(&ctx).get_result(0); + let ty = pointer.get_type(&ctx); + let cast = Operation::new( + &mut ctx, + MirCastOp::get_concrete_op_info(), + vec![ty], + vec![pointer], + vec![], + 0, + ); + MirCastOp::new(cast).set_attr_cast_kind(&ctx, MirCastKindAttr::PtrToPtr); + cast.insert_before(&ctx, goto); + } + canonicalize_read_only_aggregate_arguments(fixture.module, &mut ctx, false); + assert_eq!( + count::(&ctx, fixture.module), + expected, + "extra={extra} volatile={volatile} write={write} before={before}" + ); + pliron::operation::verify_operation(fixture.module, &ctx).unwrap(); + } + } + #[test] fn bounded_rem_rewrites_large_array_with_small_candidate_set() { let mut ctx = Context::new(); diff --git a/cuda-oxide/crates/mir-transforms/src/unroll.rs b/cuda-oxide/crates/mir-transforms/src/unroll.rs index 1616439232..c9546cf1b0 100644 --- a/cuda-oxide/crates/mir-transforms/src/unroll.rs +++ b/cuda-oxide/crates/mir-transforms/src/unroll.rs @@ -319,6 +319,28 @@ pub fn unroll_annotated_loops( }; let kind = unroll_kind(factor); + if factor == u32::MAX - 1 { + let _ = take_loop_hints(ctx, &hints, &info, loop_id); + let lp = &info.loops()[loop_id]; + let header = lp.header; + let mut backedges = Vec::new(); + for &block in &lp.blocks { + let tail = block.deref(ctx).get_tail(); + if let Some(term) = tail + && Operation::get_op::(term, ctx).is_some() + && term.deref(ctx).successors().any(|succ| succ == header) + { + backedges.push(term); + } + } + for term in backedges { + MirGotoOp::new(term).set_loop_unroll_disable(ctx); + } + function_changed = true; + changed = true; + continue; + } + let Some(ph) = info.preheader(ctx, region, loop_id) else { // The author asked for unrolling; never silently do nothing. let location = take_loop_hints(ctx, &hints, &info, loop_id); @@ -1319,6 +1341,11 @@ fn partial_unroll( } let last_latch = copies[factor_usize - 1].latch_term; rewire_goto(ctx, last_latch, main_h, &running); + MirGotoOp::new(last_latch).set_loop_unroll_disable(ctx); + let rem_latch = s.latch.deref(ctx).get_tail(); + if let Some(rem_latch) = rem_latch { + MirGotoOp::new(rem_latch).set_loop_unroll_disable(ctx); + } // The main-loop counter is mh_iv = init + (factor*step)*t, so it is always // init plus a multiple of factor*step. That lets us replace counter-derived diff --git a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock index 30afd8262e..512dd99178 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock @@ -204,27 +204,70 @@ dependencies = [ name = "cuda-oxide-codegen" version = "0.1.0" dependencies = [ + "cuda-oxide-mlir-export", "cuda-target-spec", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", "iket-lower", + "libloading", "libnvvm-sys", "llvm-export", "mir-lower", "mir-transforms", "nvvm-transforms", + "object", "pliron", "ptx-parse", "rustc-demangle", + "serde_json", + "sha2", "tempfile", "thiserror", ] +[[package]] +name = "cuda-oxide-mlir-export" +version = "0.1.0" +dependencies = [ + "cute-layout", + "dialect-cute", + "dialect-mir", + "dialect-nvvm", + "llvm-export", + "pliron", + "pliron-derive", + "pliron-mlir-interop", + "reserved-oxide-symbols", + "serde", + "serde_json", + "sha2", + "thiserror", +] + [[package]] name = "cuda-target-spec" version = "0.1.0" +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "dialect-cute" +version = "0.1.0" +dependencies = [ + "combine", + "cute-layout", + "dialect-mir", + "linkme", + "once_cell", + "pliron", + "pliron-derive", + "thiserror", +] + [[package]] name = "dialect-iket" version = "0.1.0" @@ -497,6 +540,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -652,6 +697,17 @@ dependencies = [ "thiserror", ] +[[package]] +name = "pliron-mlir-interop" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d58884474a71b0e9216d69be5c99dd1b945541f38a3acd76046d7c192a991a" +dependencies = [ + "pliron", + "pliron-llvm", + "thiserror", +] + [[package]] name = "proc-macro2" version = "1.0.107" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml index c20211c939..67b39a8c36 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml @@ -12,7 +12,7 @@ readme = "README.md" [lib] crate-type = ["dylib"] -# Keep out of workspace - requires special nightly build +# Keep out of workspace - links against rustc_private (rustc-dev) [workspace] [dependencies] diff --git a/cuda-oxide/crates/rustc-codegen-cuda/README.md b/cuda-oxide/crates/rustc-codegen-cuda/README.md index 372dc3d8a8..d4c9c405c6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/README.md @@ -34,8 +34,10 @@ fn main() { ▼ rustc_codegen_cuda ├── collect #[kernel] functions + transitive callees - ├── device MIR → dialect-mir → mem2reg → unroll → LLVM dialect → LLVM IR → PTX - └── host MIR → rustc_codegen_llvm (standard path) + ├── device MIR → dialect-mir → common preparation + verification + │ ├── ordinary backend → LLVM IR → PTX + │ └── CUTLASS backend → MLIR → cubin + └── host MIR → rustc_codegen_llvm (standard path) │ ▼ host binary + PTX file(s) @@ -51,6 +53,20 @@ The backend implements rustc's `CodegenBackend` trait. When rustc calls `codegen ## Usage +For the instruction-level CUTLASS path, select both the compiler and its +primitives profile: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ +CUDA_OXIDE_MLIR_PROFILE=cutlass-primitives-mlir22 \ +cargo oxide build --arch sm_100a +``` + +This profile accepts the reviewed MIR/NVVM subset and rejects CuTe operations +and types, including types nested inside pointers. It uses the managed CUTLASS +4.7 compiler. Leaving `CUDA_OXIDE_MLIR_PROFILE` unset retains the full-CuTe +profile. Native compilation and native MLIR dump defaults remain unchanged. + ```bash # Preferred: use the cargo-oxide wrapper cargo oxide run vecadd @@ -87,6 +103,10 @@ These are set automatically by `cargo oxide`. For manual invocations, all four a | `CUDA_OXIDE_DUMP_LLVM` | Dump the LLVM dialect module | | `CUDA_OXIDE_SHOW_RUSTC_MIR` | Dump raw rustc MIR | | `CUDA_OXIDE_EMIT_NVVM_IR` | Emit NVVM IR for libNVVM | +| `CUDA_OXIDE_DEVICE_BACKEND` | Select `native` or `cutlass-mlir` | +| `CUDA_OXIDE_CUTLASS_COMPILER` | Absolute `libCutlassCompiler.so` override | +| `CUDA_OXIDE_MLIR_PROFILE` | Select the pinned CUTLASS consumer profile | +| `CUDA_OXIDE_MLIR_OUTPUT` | Optional exact path for prepared MLIR | | `CUDA_OXIDE_DEVICE_CODEGEN_CRATE` | Comma-separated device owner crate filter | | `CUDA_OXIDE_DEVICE_ARCH` | Detected GPU arch; advisory (see below) | | `CUDA_OXIDE_DEBUG` | Device DWARF level override | @@ -120,6 +140,23 @@ underscores. Host LLVM codegen still runs for every crate. An excluded target must not call its generated module loader; the filter suppresses its device artifact but does not give sibling targets separate runtime bundle names. +### Official CUTLASS backend + +The ordinary MIR/NVVM/LLVM backend remains the default. To compile high-level +CuTe kernels through NVIDIA's official CUTLASS 4.7 compiler, first install the +pinned toolchain and then select it for the build: + +```bash +cargo oxide toolchain install cutlass +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build blockscale_gemm_cute --arch sm_120a +``` + +`cargo oxide` resolves the managed compiler automatically. An explicit +absolute `CUDA_OXIDE_CUTLASS_COMPILER` path takes precedence. The CUTLASS +continuation produces a validated cubin and uses the same embedded-artifact +launch machinery as the ordinary backend. + ## Source Layout ```text @@ -192,6 +229,6 @@ cargo oxide pipeline - **Closures work**. Both `move` closures (capture by value) and non-move closures (capture by reference via HMM) can be passed to kernels. - **Unwind paths are unreachable**. The backend ignores all MIR unwind edges, so `panic=abort` and custom sysroots are unnecessary. If a panic condition is hit at runtime, the GPU traps. -### Nightly Rust +### Pinned Rust Toolchain -The backend uses `#![feature(rustc_private)]` and pins to a specific nightly via `rust-toolchain.toml`. The workspace currently uses `nightly-2026-10-11`. +The backend uses `#![feature(rustc_private)]` (unlocked on stable via `RUSTC_BOOTSTRAP=1`) and pins to a specific toolchain release via `rust-toolchain.toml`. The workspace currently uses `1.99.1`. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md b/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md index 94b339cbda..f42a90129f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md @@ -25,7 +25,6 @@ Run `cuda-oxide/scripts/check-error-example-status.sh` to verify both are in syn | `error_host_arch_intrinsic` | diagnostics-fixture | `core::arch::` intrinsic reached | | `error_host_target_feature` | diagnostics-fixture | `#[target_feature]` fn reached | | `error_kernel_shared_param` | diagnostics-fixture | AS3/AS5 pointer as kernel parameter | -| `error_kernel_shared_ptr_param` | diagnostics-fixture | AS3 pointer inside a kernel parameter struct | | `error_missing_device_attr` | diagnostics-fixture | `thread::index_*` stub (#76) | | `error_set_discriminant_uninhabited` | diagnostics-fixture | Invalid enum variant selection | | `error_union_shared_pointer_addr` | diagnostics-fixture | `&mut`/in-place write to a shared-pointer union field | diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.lock index ce0fdb7a46..8adff48ecb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/src/main.rs index f3dab8edbd..20c46ab9c7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/src/main.rs @@ -37,7 +37,7 @@ //! ``` use cuda_core::CudaContext; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_device::{cuda_module, kernel, thread}; // ============================================================================= diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.lock index 298682ebe7..997c11b0b4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/src/main.rs index c2c124323e..9052739af7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/src/main.rs @@ -47,7 +47,7 @@ #![allow(static_mut_refs)] #![allow(clippy::assign_op_pattern)] // Expanded assignment preserves the addressof repro CFG. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, device, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.lock index f5335376d4..c8d07bc89e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/src/main.rs index de2bd536ea..2b5b6bd30a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/src/main.rs @@ -26,7 +26,7 @@ //! //! Run: `cargo oxide run aligned_field_loads` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.lock index 4c5f7cb96a..6b293e903d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/src/main.rs index 2bae17b7e0..a7948e54ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/src/main.rs @@ -33,7 +33,7 @@ //! //! Run: `cargo oxide run aligned_field_stores` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.lock index ef1ecaa2a8..1e66f8eccd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/src/main.rs index 686f91339e..3f2a47b3dc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/src/main.rs @@ -32,7 +32,7 @@ //! ./crates/rustc-codegen-cuda/examples/array_constants/verify-code-shape.sh use core::mem::MaybeUninit; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.lock index 3d737175ef..32eb93cbd0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/src/main.rs index 9b7a1c2917..d430175d91 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/src/main.rs @@ -26,7 +26,7 @@ //! //! Run: cargo oxide run array_for_loop -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.lock index d3a3501bfa..d792ec8dda 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/src/main.rs index 3a56f6fe09..31e610a29f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/src/main.rs @@ -24,7 +24,7 @@ //! //! Once array index writes are implemented, all tests should pass. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.lock index 86fba05561..e95a06828a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/src/main.rs index 8e47e960a2..707cc6aed1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/src/main.rs @@ -15,7 +15,7 @@ //! //! Run: cargo oxide run array_init -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.lock index 802b3b906e..0e7977791b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/src/main.rs index c8773d936d..72d92b7e66 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/src/main.rs @@ -34,7 +34,7 @@ //! //! Exits 0 on PASS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.lock index 3adf8ffba8..675cf05c43 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/src/main.rs index df1a653686..02c41b3909 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/src/main.rs @@ -10,7 +10,7 @@ //! ("attempted to instantiate uninhabited type"), so the kernel must trap. //! If the launch succeeds instead, the panic path was compiled away. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.lock index 7028e34dbe..848b244b9b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.lock @@ -66,6 +66,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -145,7 +165,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +180,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -172,18 +192,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -194,7 +216,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -204,6 +228,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -224,6 +249,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -390,6 +419,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", @@ -823,9 +853,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.3" +version = "1.54.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" dependencies = [ "pin-project-lite", "tokio-macros", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/README.md index 30238df523..52fd17da1b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/README.md @@ -106,12 +106,12 @@ to access the currently assigned CUDA stream and perform async memory operations ```rust fn h2d(host_data: Vec) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - unsafe { - let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memcpy_htod_async(dptr, host_data.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) - } + let host_data = Arc::new(host_data); + ctx.retain(host_data.clone()).unwrap(); + let buffer = + unsafe { DeviceBuffer::from_host_async_unchecked(ctx.get_cuda_stream(), &host_data) } + .unwrap(); + value(DeviceBox::from_buffer(buffer)) }) } ``` @@ -165,5 +165,5 @@ are required on `and_then` closure parameters: The `Zippable` trait must also be explicitly imported for `zip!` to work: ```rust -use cuda_async::simt::device_operation::Zippable; +use cuda_async::device_operation::Zippable; ``` diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/src/main.rs index de892af8f0..6eb3d3e502 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/src/main.rs @@ -35,18 +35,15 @@ //! Build and run with: //! cargo oxide run async_mlp -use cuda_async::simt::device_box::DeviceBox; -use cuda_async::simt::device_context::init_device_contexts; -use cuda_async::simt::device_operation::{self, DeviceOperation, Zippable, value}; +use cuda_async::device_box::DeviceBox; +use cuda_async::device_context::init_device_contexts; +use cuda_async::device_operation::{self, DeviceOperation, Zippable, value}; use cuda_async::zip; -use cuda_core::simt::LaunchConfig; -use cuda_core::simt::memory::{ - malloc_async, memcpy_dtoh_async, memcpy_htod_async, memset_d8_async, -}; +use cuda_core::DeviceBuffer; +use cuda_core::LaunchConfig; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; use std::future::IntoFuture; -use std::mem; use std::sync::Arc; // ============================================================================= @@ -136,51 +133,35 @@ mod kernels { /// and carries the assigned CUDA stream for this operation. fn h2d(host_data: Vec) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let n = host_data.len(); - let num_bytes = n * mem::size_of::(); - unsafe { - let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memcpy_htod_async(dptr, host_data.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) - } + let host_data = Arc::new(host_data); + // Retain before DMA: dropping or forgetting the operation must not + // release the source while the device still reads it. + ctx.retain(host_data.clone()).unwrap(); + let buffer = + unsafe { DeviceBuffer::from_host_async_unchecked(ctx.get_cuda_stream(), &host_data) } + .unwrap(); + value(DeviceBox::from_buffer(buffer)) }) } -/// Allocates zero-initialized device memory. -/// `ctx` is provided by the scheduler at execution time. +/// Zero-initialization and later kernels share the same stream and storage. fn zeros(n: usize) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let num_bytes = n * mem::size_of::(); - unsafe { - let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memset_d8_async(dptr, 0, num_bytes, stream.cu_stream()).unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) - } + value(DeviceBox::from_buffer( + DeviceBuffer::zeroed(ctx.get_cuda_stream(), n).unwrap(), + )) }) } -/// Copies device memory back to host (D2H). -/// `ctx` is provided by the scheduler at execution time. +/// Exposing an ordinary Vec to host combinators requires completion before +/// returning it. The shared transfer helper also handles partial-error cleanup. fn d2h(dev: DeviceBox<[f32]>) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let n = dev.len(); - let num_bytes = n * mem::size_of::(); - let mut host = vec![0.0f32; n]; - unsafe { - memcpy_dtoh_async( - host.as_mut_ptr(), - dev.cu_deviceptr(), - num_bytes, - stream.cu_stream(), - ) - .unwrap(); - } - // dev is kept alive in this closure until the stream synchronizes - let _ = &dev; - value(host) + value( + dev.into_buffer() + .to_host_vec(ctx.get_cuda_stream()) + .unwrap(), + ) }) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.lock index 2174866807..6a04c4e688 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.lock @@ -66,6 +66,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -145,7 +165,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +180,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -172,18 +192,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -194,7 +216,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -204,6 +228,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -224,6 +249,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -390,6 +419,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", @@ -823,9 +853,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.3" +version = "1.54.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" dependencies = [ "pin-project-lite", "tokio-macros", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/src/main.rs index d79daadf61..bbf9173e18 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/src/main.rs @@ -43,11 +43,11 @@ mod kernels { #[tokio::main] async fn main() -> Result<(), Box> { - use cuda_async::simt::device_box::DeviceBox; - use cuda_async::simt::device_context::init_device_contexts; - use cuda_async::simt::device_operation::DeviceOperation; - use cuda_core::simt::LaunchConfig; - use cuda_core::simt::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; + use cuda_async::device_box::DeviceBox; + use cuda_async::device_context::init_device_contexts; + use cuda_async::device_operation::DeviceOperation; + use cuda_core::LaunchConfig; + use cuda_core::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; use std::mem; println!("=== Async Vector Addition (cuda-async) ===\n"); @@ -68,27 +68,26 @@ async fn main() -> Result<(), Box> { // 3. Allocate device memory and copy host data. // We use cuda_core memory APIs with the context's default stream. - let (a_dev, b_dev, mut c_dev) = - cuda_async::simt::device_context::with_cuda_context(0, |ctx| { - let stream = ctx.default_stream(); - let num_bytes = N * mem::size_of::(); + let (a_dev, b_dev, mut c_dev) = cuda_async::device_context::with_cuda_context(0, |ctx| { + let stream = ctx.default_stream(); + let num_bytes = N * mem::size_of::(); - unsafe { - let a_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - let b_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - let c_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); + unsafe { + let a_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); + let b_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); + let c_dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memcpy_htod_async(a_dptr, a_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); - memcpy_htod_async(b_dptr, b_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); + memcpy_htod_async(a_dptr, a_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); + memcpy_htod_async(b_dptr, b_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); - stream.synchronize().unwrap(); + stream.synchronize().unwrap(); - let a_dev = DeviceBox::<[f32]>::from_raw_parts(a_dptr, N, 0); - let b_dev = DeviceBox::<[f32]>::from_raw_parts(b_dptr, N, 0); - let c_dev = DeviceBox::<[f32]>::from_raw_parts(c_dptr, N, 0); - (a_dev, b_dev, c_dev) - } - })?; + let a_dev = DeviceBox::<[f32]>::from_raw_parts(a_dptr, N, stream.clone()); + let b_dev = DeviceBox::<[f32]>::from_raw_parts(b_dptr, N, stream.clone()); + let c_dev = DeviceBox::<[f32]>::from_raw_parts(c_dptr, N, stream.clone()); + (a_dev, b_dev, c_dev) + } + })?; // 4. Launch the kernel asynchronously. // `vecadd_async` returns an AsyncKernelLaunch (a DeviceOperation). @@ -106,7 +105,7 @@ async fn main() -> Result<(), Box> { // 5. Copy results back. let mut c_host = vec![0.0f32; N]; - cuda_async::simt::device_context::with_cuda_context(0, |ctx| { + cuda_async::device_context::with_cuda_context(0, |ctx| { let stream = ctx.default_stream(); unsafe { memcpy_dtoh_async( diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.lock index b9a5b8237d..6566602a05 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/bin/bench.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/bin/bench.rs index ba9f8d35c8..3a24a09025 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/bin/bench.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/bin/bench.rs @@ -12,7 +12,7 @@ use std::sync::Arc; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, DeviceAtomicF16, DeviceAtomicF32}; use cuda_device::{DisjointSlice, kernel, thread}; @@ -210,7 +210,7 @@ fn reset( stream: &CudaStream, ) -> Result<(), Box> { unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( buffer.cu_deviceptr(), 0, buffer.num_bytes(), diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/main.rs index 94a77623bb..83939b562e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/src/main.rs @@ -9,7 +9,7 @@ #![feature(f16)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, BlockAtomicF16, DeviceAtomicF16, SystemAtomicF16}; use cuda_device::{DisjointSlice, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock index 88b7fea84b..6e827fc7f3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/README.md index a99a7bdca3..eb9de3f9a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/README.md @@ -14,7 +14,7 @@ Atomic operations require **LLVM 22 or newer** for correct syncscope generation. Without it, scopes (`.gpu`, `.cta`, `.sys`) will be missing from generated PTX. -The pinned Rust toolchain (`nightly-2026-10-11`) ships LLVM 23 with NVPTX +The pinned Rust toolchain (`1.99.1`) ships LLVM 23 with NVPTX enabled via the `llvm-tools` component, so the default onboarding path already satisfies this requirement: @@ -83,7 +83,7 @@ CUDA_OXIDE_VERBOSE=1 cargo oxide run atomics ### Floating-point add lowering -With the pinned `nightly-2026-10-11` toolchain's LLVM 23 NVPTX defaults, +With the pinned `1.99.1` toolchain's LLVM 23 NVPTX defaults, `atomic_f32_fetch_add_test` uses a compare-and-swap (CAS) loop for global memory. Generic-address-space `f32` adds also use a loop; shared `f32` and `f64` adds still use native instructions on supported targets. For example: diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/src/main.rs index 774cd77142..62b53bb42f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/src/main.rs @@ -46,7 +46,7 @@ //! cargo oxide run atomics use core::sync::atomic::Ordering; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::atomic::{ AtomicOrdering, BlockAtomicU32, DeviceAtomicF32, DeviceAtomicF64, DeviceAtomicI32, @@ -545,19 +545,15 @@ mod kernels { pub fn core_atomic_ordering_probe(counter: *mut u32, mut out: DisjointSlice) { let gid = thread::index_1d(); let pointer = counter; - // nightly-2026-08-28 added a `const VOLATILE: bool` generic to the - // load/store intrinsics; plain atomics pass `false`. let current = unsafe { - core::intrinsics::atomic_load::( + core::intrinsics::atomic_load::( pointer, ) }; unsafe { - core::intrinsics::atomic_store::< - u32, - { core::intrinsics::AtomicOrdering::Release }, - false, - >(pointer, current) + core::intrinsics::atomic_store::( + pointer, current, + ) }; let swapped = unsafe { core::intrinsics::atomic_xchg::( diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock index 8b838a20e4..2cec46e09f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/src/main.rs index f450bff9f0..3ebcec49f7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/src/main.rs @@ -15,7 +15,7 @@ //! Build and run with: //! cargo oxide run barrier -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::barrier::{ Barrier, mbarrier_arrive, mbarrier_arrive_no_complete, mbarrier_init, mbarrier_inval, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.lock index de59e51102..f162611ac5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/src/main.rs index 52f20059ce..f993c4dc52 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/src/main.rs @@ -3,7 +3,7 @@ //! End-to-end example exercising all eight packed bf16x2 arithmetic intrinsics. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::bf16x2::{ abs_bf16x2, add_bf16x2, fma_relu_bf16x2, max_bf16x2, min_bf16x2, mul_bf16x2, neg_bf16x2, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock index 251e9a2a67..57c42f1c92 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/src/main.rs index cf0a5af7bc..50486794a7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/src/main.rs @@ -3,7 +3,7 @@ //! Behavior example for the packed bf16x2 FMA intrinsic. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::bf16x2::fma_bf16x2; use cuda_device::convert::cvt_bf16x2_f32; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock index 670e54de0c..75df3b0a1a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "cuda-host", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/src/main.rs index 6225168543..ea5f112864 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/src/main.rs @@ -24,7 +24,7 @@ * cargo oxide run bool_phi_cmp */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.lock index c07608a163..c2ab84c661 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "carrying_mul_add" version = "0.1.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/src/main.rs index 3c91191070..31a6c243ce 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/src/main.rs @@ -30,7 +30,7 @@ #![feature(signed_bigint_helpers)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.lock index 7354ec93d3..3a0f53e07d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cast_chunks" version = "0.2.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/src/main.rs index 72b3ec3e78..3e61123662 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/src/main.rs @@ -15,7 +15,7 @@ //! //! Run: `cargo oxide run cast_chunks` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::cuda_module; @@ -28,6 +28,11 @@ pub struct Quad(pub [f32; 4]); // Plain POD aggregate, no pointers: safe to memcpy to/from the device. unsafe impl cuda_core::DeviceCopy for Quad {} +// Four `f32` with no padding: every bit pattern is a valid `Quad`, which is +// what lets `cast_chunks` build one from flat `f32` bytes. +unsafe impl cuda_core::bytemuck::Zeroable for Quad {} +unsafe impl cuda_core::bytemuck::Pod for Quad {} + #[cuda_module] mod kernels { use super::Quad; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.lock index 5a0f4f1248..ccb62861b1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cast_tests" version = "0.1.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/src/main.rs index 0fc21de2dd..1f907f8260 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/src/main.rs @@ -24,7 +24,7 @@ //! the correct LLVM instruction. Struct↔ptr conversions (fat pointers, newtypes) //! are handled generically via `emit_pointer_cast` in mir-lower. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.lock index 3dac7b1997..5090e773f6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cbrt_smoke" version = "0.1.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/src/main.rs index 4c84f7469a..4a8507f41b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/src/main.rs @@ -33,7 +33,7 @@ //! //! Exits 0 on PASS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.lock index aef438f58d..6a477056e0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/src/main.rs index dfdc160be1..50c7e5dfe0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/src/main.rs @@ -21,7 +21,7 @@ //! cargo oxide run checked_arith //! CUDA_OXIDE_NO_OPT=1 cargo oxide run checked_arith -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock index b2358c8f96..2370ffbd94 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/src/main.rs index ddd6d92684..eb8c4191ef 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/src/main.rs @@ -20,7 +20,7 @@ //! **Hardware Requirements:** Blackwell (B200, GB200) with SM 100+ use core::ptr::addr_of_mut; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::barrier::{ Barrier, fence_proxy_async_shared_cta, mbarrier_arrive_expect_tx, mbarrier_init, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock index 428561d84b..96c92cd9d2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/src/main.rs index d9d2a5c18f..4a0f9ed9ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/src/main.rs @@ -274,7 +274,7 @@ mod kernels { // ============================================================================ fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== Thread Block Cluster Tests (sm_90+) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.lock index 33acd7147e..def5b02172 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/src/main.rs index b774c7856b..7e5981601d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/src/main.rs @@ -21,7 +21,7 @@ //! //! Run: cargo oxide run compiler_features -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::shared::cvta_generic_to_shared_offset; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.lock index c4fa3cf464..d4adfa615c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/src/main.rs index dd9af04d1b..314b037234 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/src/main.rs @@ -17,7 +17,7 @@ //! //! Run: cargo oxide run complex_mul -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.lock index 6ea623f850..856d9bedcd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/src/main.rs index fbbe5150b6..c16dcef2b1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/src/main.rs @@ -21,7 +21,7 @@ //! //! Run: cargo oxide run const_aggregate -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.lock index 91b2876e08..78d37d279f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/src/main.rs index 3e46631a87..a2d4b18efc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/src/main.rs @@ -4,7 +4,7 @@ //! Regression for generic branches controlled by an associated constant. use core::marker::PhantomData; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DeviceCopy}; use cuda_device::{DisjointSlice, DynamicSharedArray, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.lock index b15a635b4c..fd8ba2a23c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/src/main.rs index 05f649dde7..2afea93825 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/src/main.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{cuda_module, device, kernel, launch_bounds, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.lock index 8ae82696df..01e0611b88 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/src/main.rs index dae9eaa082..4e156013a9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/src/main.rs @@ -45,7 +45,7 @@ //! //! Run: `cargo oxide run const_table_lookup` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.lock index 53037c381d..16c06023f0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/src/main.rs index 464417d905..9d2f0044f0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/src/main.rs @@ -37,7 +37,7 @@ //! //! Run: `cargo oxide run const_tuple_table_lookup` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock index ab06a74fe1..b435c6bc2b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/src/main.rs index 13699d90db..974c80fb92 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/src/main.rs @@ -19,7 +19,7 @@ //! Run with: //! CUDA_OXIDE_DUMP_MIR=1 cargo oxide run constant_index_from_end -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.lock index 441dc70016..37afdf5238 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/src/main.rs index f575c5f970..87e7b1807e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/src/main.rs @@ -14,7 +14,7 @@ //! Build and run with: //! cargo oxide run constant_memory -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ConstantMemory, DisjointSlice, constant, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.lock index cbd2c4a773..fd901086ef 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/src/main.rs index 5cc830e458..e065d0b3fd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/src/main.rs @@ -21,7 +21,7 @@ //! Build and run with: //! cargo oxide run constant_memory_coeffs -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ConstantMemory, DisjointSlice, constant, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.lock index 8656b34c60..8f6b2748a7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/src/main.rs index 4bb64fe912..45b8ffdb3d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/src/main.rs @@ -8,7 +8,7 @@ //! Build and run with: //! cargo oxide run constant_memory_simple -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ConstantMemory, DisjointSlice, constant, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.lock index e1b980e909..249e9c9ff1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/src/main.rs index 5a8d5783fc..f61e5716b1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/src/main.rs @@ -42,7 +42,7 @@ //! //! Run: `cargo oxide run constant_memory_table` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.lock index cd1195cbd1..6365863597 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/src/main.rs index 1b5cf1e66a..47ef78eb45 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/src/main.rs @@ -20,7 +20,7 @@ //! op/type matrix. The generated `coop_groups_demo.ptx` is also handy //! for inspecting how each kernel actually lowers. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D}; use cuda_device::cooperative_groups::{ ThreadGroup, WarpCollective, block_reduce, block_scan, coalesced_threads, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.lock index 0ef4472db0..551a566bf6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/src/main.rs index 2017fb65ed..b0d73e86bf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/src/main.rs @@ -9,7 +9,7 @@ //! read-only methods. Runtime indexing into its small array fields must not //! force the aggregate into NVPTX local memory. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.lock index 03fc3fb565..b85f374d71 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/src/main.rs index ac8bda154d..90efeb06ca 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/src/main.rs @@ -19,7 +19,7 @@ //! Usage: //! cargo oxide run copy_nonoverlapping -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock index ee38589b3f..d89a1dc228 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/src/main.rs index 088cb66596..8ec3f3a380 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/src/main.rs @@ -26,7 +26,7 @@ //! Build and run with: //! cargo oxide run counted_barrier -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::barrier::{barrier_cta_arrive, barrier_cta_sync}; use cuda_device::{DisjointSlice, SharedArray, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.lock index 8884ed45e2..5dab62fd2d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/src/main.rs index ee84d89c02..dd4fe23738 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/src/main.rs @@ -14,7 +14,7 @@ //! Build and run with: //! cargo oxide run cp_async_small -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::async_copy::{ cp_async_ca_4, cp_async_ca_8, cp_async_ca_16, cp_async_ca_zfill_4, cp_async_ca_zfill_8, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock index 4c768d4fe9..a10dcfb4dd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -137,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -180,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -200,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/src/main.rs index b2a9d6b174..bb7c873a5d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/src/main.rs @@ -17,7 +17,7 @@ //! Build and run with: //! cargo oxide run cp_async_zfill -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::async_copy::{cp_async_ca_zfill_4, cp_async_commit_group, cp_async_wait_all}; use cuda_device::{DisjointSlice, SharedArray, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock index 51381eb264..2b58cf9f10 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -136,7 +156,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -148,18 +168,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -170,7 +192,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -179,6 +203,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -199,6 +224,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -263,6 +292,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock index 58f9f47452..03e6dff9e0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -274,6 +303,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/src/main.rs index a1de359f63..e62b3144b5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/src/main.rs @@ -22,7 +22,7 @@ //! //! Run: cargo oxide run cross_crate_embedded -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use embedded_kernel_lib::kernels; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock index 4142ea1aa4..045d746f67 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -139,7 +159,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -151,18 +171,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -173,7 +195,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -182,6 +206,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -202,6 +227,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -266,6 +295,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/src/main.rs index 36388bf446..34d9b76a7f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/src/main.rs @@ -32,7 +32,7 @@ //! cargo oxide run cross_crate_kernel //! ``` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; // Import the public kernel functions and their generated `_ptx_name` diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock index 12297db68c..3446b32c32 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock index 8f2b797d96..b6ab52f7f5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -138,7 +158,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,18 +170,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -172,7 +194,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -181,6 +205,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -201,6 +226,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/src/main.rs index b10a5fdbd0..0de522d2d5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/src/main.rs @@ -10,7 +10,7 @@ //! defined LOCALLY in this crate always compiled, which isolates the axis to "cross-crate + not //! inlinable", not the function body. See PR "pass -Zalways-encode-mir ...". -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.lock index 08be6237f6..ac756b3d60 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.lock index ed02b84ccf..8666eb0aee 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/src/main.rs index 5150484d3b..cff252211b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/src/main.rs @@ -48,7 +48,7 @@ //! cargo oxide run cuda_module_in_lib //! ``` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; use module_kernels::kernels; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock index 971388c6c3..00e87923ff 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/src/main.rs index ecdacf1082..d1162c1de3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/src/main.rs @@ -17,7 +17,7 @@ //! Build and run with: //! cargo oxide run cuda_module_nested -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.lock index fe8a69f59d..190b133ad5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.lock @@ -76,6 +76,26 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "bytes" version = "1.11.1" @@ -169,7 +189,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -184,7 +204,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -196,27 +216,33 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cutile" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-async", @@ -238,7 +264,7 @@ dependencies = [ [[package]] name = "cutile-compiler" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-async", @@ -258,7 +284,7 @@ dependencies = [ [[package]] name = "cutile-ir" -version = "0.4.0" +version = "0.5.0" dependencies = [ "half", "indexmap", @@ -268,7 +294,7 @@ dependencies = [ [[package]] name = "cutile-macro" -version = "0.4.0" +version = "0.5.0" dependencies = [ "convert_case", "cutile-compiler", @@ -471,6 +497,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", @@ -1017,9 +1044,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.3" +version = "1.54.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" dependencies = [ "bytes", "libc", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.toml index 7632b78ccf..d6804cd667 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/Cargo.toml @@ -16,8 +16,8 @@ ptx-dir = "simt" resolver = "2" [dependencies] -cutile = { path = "../../../../../cutile-rs/cutile", version = "0.4.0" } -cutile-compiler = { path = "../../../../../cutile-rs/cutile-compiler", version = "0.4.0" } +cutile = { path = "../../../../../cutile-rs/cutile", version = "0.5.0" } +cutile-compiler = { path = "../../../../../cutile-rs/cutile-compiler", version = "0.5.0" } cuda-async = { path = "../../../../../cuda-async" } cuda-core = { path = "../../../../../cuda-core" } tokio = { version = "1", features = ["full"] } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/README.md index d68dd21911..f9631d8910 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/README.md @@ -32,9 +32,9 @@ input tensor softmax_out tensor | | cuda-oxide SIMT kernel loaded from generated PTX - | threshold_scale_f32(softmax_out -> gated_out) + | threshold_scale_f32(softmax_out -> same allocation) v -gated_out tensor +processed tensor (same allocation) ``` The Tile stage is intentionally tile-shaped: it uses row reductions and @@ -84,35 +84,14 @@ PASS: cutile-rs Tile softmax -> cuda-oxide SIMT threshold/scale passed ## Same stream, shared arrays -The important cutile-rs mechanism is `DeviceOp::then`. - -```text -tile_softmax::row_softmax(...) - .then(|softmax_out| ThresholdScaleKernel { ... }) -``` - -`then` executes both operations with the same cutile-rs `ExecutionContext`. -That context owns the CUDA stream, so the Tile launch and cuda-oxide PTX -launch are submitted to the same stream in order. - -The shared arrays are just cutile-rs `Tensor` allocations. The custom SIMT -`DeviceOp` takes the Tile-produced tensor and passes its raw `CUdeviceptr` to -the cuda-oxide kernel: - -```text -softmax_out: Tensor - | - | device_pointer().cu_deviceptr() - v -threshold_scale_f32(..., input_ptr, output_ptr) -``` - -So the interop contract is deliberately small: - -```text -same process -same CUDA context -same CUDA stream -same device allocations -PTX-loaded cuda-oxide kernel launched as a cutile-rs DeviceOp -``` +Both stages use `cuda_async::device_operation::DeviceOp` and its shared execution context. +`then` submits them to the same stream in order. The example transfers the Tile +result with `Tensor::into_buffer`, passes it through the shared +`AsyncKernelLaunchBuilder` / `OwnedAsyncKernelLaunch` API, and returns it with +`Tensor::from_buffer`. It needs no custom `DeviceOp` implementation. + +The SIMT kernel reads and overwrites each element in place. Assertions check +both allocation identity and device address before and after the transition, +then check every result against the CPU calculation. The only device-to-host +copy is the final verification readback. Launch descriptors retain the shared +owner and write access until GPU completion, including early future drop. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.lock index 35aaf4c667..249e9cbf7d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs index 6d931332da..0e16174e39 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/src/main.rs @@ -10,16 +10,13 @@ //! 2. cuda-oxide SIMT kernel thresholds and scales the softmax result. //! 3. Both kernels run on the same CUDA stream and share device tensors. -use cuda_async::device_context::{load_module_from_ptx, with_default_device_policy}; -use cuda_async::device_future::DeviceFuture; -use cuda_async::device_operation::{DeviceOp, ExecutionContext}; -use cuda_async::error::DeviceError; -use cuda_async::launch::AsyncKernelLaunch; -use cuda_core::{Function, LaunchConfig}; +use cuda_async::device_context::load_module_from_ptx; +use cuda_async::device_operation::DeviceOp; +use cuda_async::launch::{AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; +use cuda_core::LaunchConfig; use cutile::api::{copy_host_vec_to_device, zeros}; use cutile::error::Error; use cutile::tensor::{IntoPartition, Reshape, Tensor, ToHostVec}; -use std::future::IntoFuture; use std::sync::Arc; const CUDA_OXIDE_SIMT_PTX_PATH: &str = concat!( @@ -50,56 +47,6 @@ mod tile_softmax { } } -struct ThresholdScaleKernel { - function: Arc, - n: u32, - threshold: f32, - scale: f32, - input: Arc>, - output: Tensor, -} - -impl DeviceOp for ThresholdScaleKernel { - type Output = (Arc>, Tensor); - - unsafe fn execute( - self, - ctx: &ExecutionContext, - ) -> Result<::Output, DeviceError> { - let mut launcher = AsyncKernelLaunch::new(self.function); - launcher.push_arg(self.n); - launcher.push_arg(self.threshold); - launcher.push_arg(self.scale); - unsafe { - launcher - .push_device_ptr(self.input.device_pointer().cu_deviceptr()) - .push_device_ptr(self.output.device_pointer().cu_deviceptr()); - } - launcher.set_launch_config(LaunchConfig { - grid_dim: (self.n.div_ceil(256), 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }); - unsafe { launcher.execute(ctx)? }; - Ok((self.input, self.output)) - } -} - -impl IntoFuture for ThresholdScaleKernel { - type Output = Result<(Arc>, Tensor), DeviceError>; - type IntoFuture = DeviceFuture<(Arc>, Tensor), ThresholdScaleKernel>; - - fn into_future(self) -> Self::IntoFuture { - match with_default_device_policy(|policy| { - let stream = policy.next_stream()?; - Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream))) - }) { - Ok(Ok(future)) => future, - Ok(Err(e)) | Err(e) => DeviceFuture::failed(e), - } - } -} - fn softmax_rows(input: &[f32], rows: usize, cols: usize) -> Vec { let mut out = vec![0.0f32; input.len()]; for row in 0..rows { @@ -158,26 +105,37 @@ async fn run() -> Result<(), Error> { .expect("input reshape") .into(); let softmax_out = zeros::(&[ROWS, COLS]).await?; - let gated_out: Tensor = zeros::(&[ROWS, COLS]).await?; + let allocation = softmax_out.as_buffer()?.storage().allocation_id(); + let address = softmax_out.device_pointer().cu_deviceptr(); // `then` is the key inter-kernel interop piece here: cutile-rs executes // both DeviceOps with the same ExecutionContext, so the Tile launch and // the cuda-oxide PTX launch are submitted to the same CUDA stream. let pipeline = tile_softmax::row_softmax(softmax_out.partition([BM, BN]), input.clone()).then( move |(softmax_part, _input)| { - let softmax_out: Arc> = softmax_part.unpartition().into(); - ThresholdScaleKernel { - function: threshold_scale, - n: N as u32, - threshold, - scale, - input: softmax_out, - output: gated_out, - } + // Transfer the Tile result to the shared typed buffer. No allocation + // or device copy is needed for the SIMT stage. + let buffer = softmax_part.unpartition().into_buffer(); + assert_eq!(buffer.storage().allocation_id(), allocation); + assert_eq!(buffer.cu_deviceptr(), address); + let mut builder = AsyncKernelLaunchBuilder::new(threshold_scale); + builder.retain_buffer(buffer.storage().clone(), true); + builder.push_args((N as u32, threshold, scale, address, address)); + // SAFETY: the SIMT function's five parameters match this packet; + // each thread reads then overwrites its own f32 within 0..N. Raw + // pointers permit this in-place operation. Storage and exclusion + // are retained independently of the returned typed buffer. + let launch = + unsafe { builder.finalize_unchecked(LaunchConfig::for_num_elems(N as u32)) }; + OwnedAsyncKernelLaunch::new(launch, buffer) }, ); - let (_softmax_out, gated_out) = pipeline.await?; + let buffer = pipeline.await?; + assert_eq!(buffer.storage().allocation_id(), allocation); + assert_eq!(buffer.cu_deviceptr(), address); + let gated_out = Tensor::from_buffer(buffer, vec![ROWS as i32, COLS as i32]); + assert_eq!(gated_out.device_pointer().cu_deviceptr(), address); let got = gated_out.to_host_vec().await?; let softmax_expected = softmax_rows(&input_host, ROWS, COLS); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.lock index 818576913e..1347156dfc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/src/main.rs index 3152f7fae6..87a695465e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/src/main.rs @@ -15,7 +15,7 @@ //! //! Run: cargo oxide run cvt_f16x2 -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::convert::cvt_f16x2_f32; use cuda_device::{DisjointSlice, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.lock index 94d0618251..f38cc20a4e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/src/main.rs index f2fe30c8fd..b9a825ecfb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/src/main.rs @@ -31,7 +31,7 @@ //! //! Run: cargo oxide run cvt_packed -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::convert::{ cvt_f16x2_f32, cvt_f32_bf16x2_hi, cvt_f32_bf16x2_lo, cvt_f32_f16x2_hi, cvt_f32_f16x2_lo, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.lock index f5df8eec21..da2f203013 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/src/main.rs index deff084d07..e1beda3ea9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/src/main.rs @@ -41,7 +41,7 @@ //! ./debug-tests oob_index triggers a Rust bounds-check trap //! ./debug-tests null_deref triggers a null-pointer dereference -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ ConstantMemory, DisjointSlice, SharedArray, constant, debug, gpu_assert, gpu_printf, kernel, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock index 06c9f62bfd..5581889fe3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/src/main.rs index deceab0b39..9d0d9053a9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/src/main.rs @@ -170,7 +170,7 @@ mod kernels { // ============================================================================= fn main() -> Result<(), Box> { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== GPU Debug & Utility Intrinsics Test (Unified) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.lock index a163c07916..515323c9e5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/src/main.rs index bab5eab83c..88fcdee2af 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/src/main.rs @@ -8,7 +8,7 @@ //! Run under cuda-gdb with: //! `CUDA_OXIDE_DEBUG=full cargo oxide debug debug_pointer_locals` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock index a71448a33a..1183c5ba44 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs index 5b4f1fe1cc..38b14208d5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/src/main.rs @@ -322,7 +322,7 @@ mod kernels { fn main() { println!("=== Unified Compilation: Closures ===\n"); - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.lock index a158e03b8d..5235490398 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/src/main.rs index 0deb545121..b12586cbc0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/src/main.rs @@ -511,7 +511,7 @@ fn build_pipeline() -> Result { // Uses cuda-driver to load the merged cubin, launch kernels, and verify results. // ============================================================================= -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use std::sync::Arc; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.lock index 35ee0cd474..c238e58027 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/src/main.rs index f3703af6cb..0583a5fdc6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/src/main.rs @@ -8,7 +8,7 @@ //! Build and run with: //! cargo oxide run device_global -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::kernel; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.lock index 6407c721d4..281fb5257c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/src/main.rs index fa446fe363..60f511b016 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/src/main.rs @@ -18,7 +18,7 @@ //! driver reports as a launch failure. A launch that instead *succeeds* would //! mean the panic path fell through rather than trapping. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.lock index 16f5eeea64..986b41be11 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/src/main.rs index 11b1f27fee..1bb917f4c9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/src/main.rs @@ -24,7 +24,7 @@ //! wrote a constant, or that addressed the wrong element, fails the check //! rather than passing quietly. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, Runtime2DIndex, cuda_module, device, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock index dd634310c2..c9ee19ef4e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/src/main.rs index 85625fe267..76b37495e9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/src/main.rs @@ -15,7 +15,7 @@ //! ``` //! //! Run: cargo oxide run disjoint_slice_len -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock index 6eae6d938e..1bc6591a03 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/src/main.rs index e502cda0a1..c48212973d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/src/main.rs @@ -25,7 +25,7 @@ //! Build and run with: //! cargo oxide run dotprod --arch sm_80 -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::dotprod::{dp2a_s32, dp2a_u32, dp4a_s32, dp4a_u32}; use cuda_device::{DisjointSlice, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.lock index e089b4e94a..ff104c654c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/src/main.rs index d59612e2d3..974d89d346 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/src/main.rs @@ -14,7 +14,7 @@ //! //! Run: cargo oxide run dpx_minmax_chains -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.lock index 266b7f6694..02248f38f5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/src/main.rs index 316fd51308..38296aab33 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/src/main.rs @@ -22,7 +22,7 @@ use core::mem::{ManuallyDrop, MaybeUninit}; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.lock index 9c6fb2bc3e..9d57db5b04 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/src/main.rs index 0960a045ae..d752a6de32 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/src/main.rs @@ -14,7 +14,7 @@ //! Usage: //! cargo oxide run dst_metadata -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock index d00f6456ab..505aab8101 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/src/main.rs index b0603a0ed2..3df807c48f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/src/main.rs @@ -26,7 +26,7 @@ //! Build and run with: //! cargo oxide run dynamic_smem -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, DynamicSharedArray, gpu_printf, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.lock index 24b95937f2..783993e150 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/src/main.rs index a1a83d0b50..0eb94779f9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/src/main.rs @@ -74,7 +74,7 @@ mod kernels { // ============================================================================= fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== elect.sync warp leader election (sm_90+) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.lock index 6139e1050f..b11fea5b99 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/src/main.rs index 785853b78b..4507a2610e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/src/main.rs @@ -31,7 +31,7 @@ //! Run with: //! cargo oxide run enum_abi_multi_payload -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer, DeviceCopy}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.lock index 4985ca6d9f..6309e2f48f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/src/main.rs index 7d3d200096..ef6ec3c203 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/src/main.rs @@ -34,7 +34,7 @@ //! Run with: //! cargo oxide run enum_array_match -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.lock index 8af4e5d135..f4ff213c7c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/src/main.rs index 941a825e11..40049127ef 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/src/main.rs @@ -9,7 +9,7 @@ //! static subobject, direct-tagged enums, and relocation-carrying enums nested //! inside struct, tuple, and array constants. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::kernel; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.lock index c952ef0989..c939c78abc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/src/main.rs index 348dafc884..0239ba88b5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/src/main.rs @@ -34,7 +34,7 @@ //! Each kernel reads its value back after mutating, so a write that landed in //! a copy shows up as an unchanged element rather than passing quietly. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, device, kernel, thread}; use std::time::Instant; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock index 3a6c6a60cc..1c4b0fce44 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/src/main.rs index 0b20e5abe7..b5339d52ae 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/src/main.rs @@ -13,7 +13,7 @@ //! cargo oxide run enum_state_transitions //! CUDA_OXIDE_NO_OPT=1 cargo oxide run enum_state_transitions -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock index 984dd1bb34..258c87ae2a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/src/main.rs index 4e37292a26..a101ec7f1a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/src/main.rs @@ -41,7 +41,7 @@ //! //! Run: `cargo oxide run enum_table_lookup` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock index aaa54e5982..277c334ba0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.lock index ce15431ac5..9b3985fd66 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock index 0fa3b0a4ee..60620c0e2c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock index f06797a06b..d09584d81a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock index 4c5a1cecc5..b7bca8cb7b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -268,6 +297,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_callable/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_callable/Cargo.lock index edd815e111..9520e13b15 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_callable/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_callable/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -268,6 +297,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_fn_pointer/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_fn_pointer/Cargo.lock index d5446db5d0..6356b965df 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_fn_pointer/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_fn_pointer/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -268,6 +297,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_unknown_id/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_unknown_id/Cargo.lock index 9b25129d83..e4e1aa470c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_unknown_id/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_unknown_id/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -268,6 +297,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock index 1a466544e6..ca89f31898 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock index 7191053847..25c7e6f0d6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock index 3219ee0d34..6d3cb3b190 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock index 44916f1a99..46663f9dd4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml deleted file mode 100644 index 15fee6ddd1..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "error_kernel_shared_ptr_param" -version = "0.1.0" -edition = "2024" -authors = ["NVIDIA Corporation"] -license = "Apache-2.0" - -[workspace] - -[dependencies] -cuda-core = { path = "../../../../../cuda-core" } -cuda-device = { path = "../../../cuda-device" } -cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md deleted file mode 100644 index 0fcfdc5d95..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md +++ /dev/null @@ -1,32 +0,0 @@ -# `error_kernel_shared_ptr_param` - -Negative test for a `#[kernel]` whose parameter is a struct that holds a -`SharedPtr`, a pointer that lowers to `addrspace(3)` by its type. - -A kernel receives its parameters in `.param` space, filled by the host at -launch. Shared memory is allocated per block by the device, so the host holds -no address in it to pass. `error_kernel_shared_param` pins the refusal of a -bare shared pointer parameter. This fixture pins the same refusal for a -pointer carried inside an aggregate: before the exporter looked through struct -fields and array elements, `Args { p: SharedPtr, n: u32 }` compiled, the -entry declared a 16-byte `.param`, and the kernel applied `cvta.shared.u64` to -the host-supplied bytes as if they were a shared address. - -The compiler must reject the signature at compile time instead: - -```bash -cargo oxide build error_kernel_shared_ptr_param -``` - -Expected diagnostic (pinned by `cuda-oxide/scripts/smoketest.sh`): - -```text -holds a pointer into shared memory -``` - -The diagnostic names the field path inside the parameter (`.0` here), which -keeps it distinct from the bare parameter's `is a pointer into shared memory`. -Structs of global (`addrspace(1)`) and constant (`addrspace(4)`) pointers stay -allowed as kernel parameters: the host allocates in those spaces, so it has an -address to supply. A `SharedPtr` belongs on the device side, created inside -the kernel from a `static mut SharedArray` or from `DynamicSharedArray`. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs deleted file mode 100644 index 74f8b80c79..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs +++ /dev/null @@ -1,51 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! Negative test: a `#[kernel]` parameter struct that holds a `SharedPtr`. -//! -//! A kernel receives its parameters in `.param` space, filled by the host at -//! launch. Shared memory is allocated per block by the device, so the host -//! holds no shared-memory address it could pass. A bare `SharedPtr` -//! parameter is refused for that reason (`error_kernel_shared_param` covers -//! the bare case with `*mut Barrier`), and a struct or array that carries one -//! must be refused the same way. Before the exporter looked inside -//! aggregates, `Args { p: SharedPtr, n: u32 }` compiled and the kernel -//! converted the host-supplied bytes with `cvta.shared.u64` as if they were a -//! shared address. -//! -//! Usage: -//! cargo oxide run error_kernel_shared_ptr_param -//! -//! Expected: the build FAILS with a diagnostic containing (pinned): -//! -//! ```text -//! holds a pointer into shared memory -//! ``` - -use cuda_device::{SharedPtr, kernel}; - -/// Kernel arguments that carry a typed shared pointer next to a plain field. -#[derive(Clone, Copy)] -pub struct Args { - pub p: SharedPtr, - pub n: u32, -} - -/// # Safety -/// -/// Never launched: the device compilation is expected to fail before a host -/// binary that could launch it exists. -#[kernel] -pub unsafe fn shared_ptr_param_kernel(args: Args) { - // BUG UNDER TEST: `args.p` points into shared memory, which the host - // cannot supply at launch. The refusal fires on the entry signature - // alone; the body deliberately does nothing else, so the diagnostic - // under test is the only error this fixture can produce. - let _ = args; -} - -fn main() { - println!("This negative example should fail during device compilation."); -} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock index 649e784194..a5fad8161a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock index 5c81d90910..ca4fc82ac4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock index 746e3e1063..91d772ac18 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.lock index ac804086d6..cd11060441 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.lock index 33cd0aa21e..e396de7c2c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/src/main.rs index 83538ee75e..5239cf4ee4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/src/main.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::float::ex2_approx_f16; use cuda_device::{DisjointSlice, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock index 93a2a2860b..5b5db6ebba 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/src/main.rs index 1003adab19..c8751b906f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/src/main.rs @@ -21,7 +21,7 @@ #![feature(core_intrinsics)] #![allow(internal_features)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.lock index ad50057c7b..9e0b233df7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/src/main.rs index 49ba3eb1ea..320096b78c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/src/main.rs @@ -16,7 +16,7 @@ //! not be classified as generic merely because the MIR `FnDef` operand carries //! substitutions or because the FQDN contains `<`, `>`, or `::`. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock index a3de61cf65..53d4282e18 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/src/main.rs index 4850e874fb..75400d096b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/src/main.rs @@ -20,7 +20,7 @@ //! //! Run: cargo oxide run extern_crate_closure -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.lock index 676af798d0..4b2a6016ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/src/main.rs index 309ce7548d..0dbde672b9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/src/main.rs @@ -24,7 +24,7 @@ //! Exits 0 on SUCCESS, 1 on FAIL. use core::marker::PhantomData; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.lock index 1981d414f2..79ea490936 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/src/main.rs index 0d78067872..4cd16438ef 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/src/main.rs @@ -12,7 +12,7 @@ //! //! Run: cargo oxide run f16_stress -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.lock index 31579841ea..8fa5a8ab7f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/src/main.rs index bc20264607..b1f34607de 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/src/main.rs @@ -3,7 +3,7 @@ //! End-to-end example for packed f16x2 arithmetic intrinsics. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::f16x2::{ abs_f16x2, add_f16x2, fma_f16x2, fma_ftz_f16x2, fma_ftz_relu_f16x2, fma_ftz_sat_f16x2, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock index 9bc296863f..a773592cb5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/src/main.rs index b40854d4c0..49cf92233f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/src/main.rs @@ -3,7 +3,7 @@ //! End-to-end example for the SM100 packed f32x2 arithmetic family. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::f32x2::{ add_f32x2, add_ftz_f32x2, fma_f32x2, fma_ftz_f32x2, mul_f32x2, mul_ftz_f32x2, pack_f32x2, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock index 71b0d21466..736d92314e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/src/main.rs index f6c46412ac..40df592afe 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/src/main.rs @@ -6,7 +6,7 @@ // Minimal repro: float `!=` should be UNORDERED (x != x is TRUE for NaN). // Rust PartialEq::ne on floats is unordered. cuda-oxide lowered it to fcmp ONE // (ordered, FALSE for NaN), so x != x folded to false -> NaN handling broken. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.lock index b198305b8e..c6682dfb5c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/src/main.rs index ddbe90fad2..531e83470d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/src/main.rs @@ -16,7 +16,7 @@ //! //! Run: cargo oxide run field_array_assign -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.lock index 80dfa026eb..49f564f82c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/src/main.rs index a8f31d8d14..22eee8a9d2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/src/main.rs @@ -31,7 +31,7 @@ //! //! Exits 0 on PASS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.lock index a3233b3c21..f90e1d5b47 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/src/main.rs index 958a72a1a5..392fcdc79f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/src/main.rs @@ -29,7 +29,7 @@ //! Build and run with: //! cargo oxide run fmaxmin_smoke -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.lock index 2a909c83ac..f048a50004 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/src/main.rs index f726fb5862..a1e592d7a0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/src/main.rs @@ -17,7 +17,7 @@ //! Run with: //! cargo oxide run fn_ptr_naming -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.lock index 86eb3f25e8..dc39963e78 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/src/main.rs index 09d2a2fb32..a8f158a27d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/src/main.rs @@ -31,7 +31,7 @@ //! Build and run with: //! cargo oxide run fp_special_literal_smoke -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.lock index fc1c163579..3c9265aca4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs index 08e0784125..d3f90fd774 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs @@ -18,7 +18,9 @@ //! `call_once_decoy` proves that an ordinary function whose name contains //! trait-like text is not mistaken for a callable-trait shim. -use cuda_core::simt::LaunchConfig; +#![feature(never_type)] + +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.lock index cbfd3853ea..ee8c375e2f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs index 55426c2f0a..19f6030320 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/src/main.rs @@ -304,7 +304,7 @@ mod kernels { // ============================================================================= fn main() -> Result<(), Box> { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== Future APIs Test (Unified) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock index d83a865a01..1bfb93c818 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/src/main.rs index 23f1f9108a..248d5fe16f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/src/main.rs @@ -13,7 +13,7 @@ //! Build and run with: //! cargo oxide run gemm -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.lock index 4ce39a1971..b601daec7a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/src/main.rs index f8e9a8d45b..577b817a93 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/src/main.rs @@ -36,7 +36,7 @@ //! Build and run with: //! cargo oxide run gemm_sol -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32}; use cuda_device::barrier::{ @@ -5821,7 +5821,7 @@ fn run_benchmark_persistent( for _ in 0..WARMUP { let z = 0u32; unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( dev_tile_counter.cu_deviceptr(), &z as *const u32, std::mem::size_of::(), @@ -5851,7 +5851,7 @@ fn run_benchmark_persistent( for _ in 0..ITERS { let z = 0u32; unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( dev_tile_counter.cu_deviceptr(), &z as *const u32, std::mem::size_of::(), diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.lock index ad956f3af4..1d75326b3d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -266,6 +295,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/README.md index 7def4f5c1d..2680e36d2c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/README.md @@ -7,12 +7,30 @@ The goal is simple: show how close a Rust/cuda-oxide kernel can get to the closest supported live cuBLASLt GEMM on the same GPU, without weakening the workload or correctness check. +## Measured B200 throughput vs cuBLASLt + +```text +Visual Bar Order on B200 Panel (0 – 2.2 PFLOP/s axis): + + 4096³ cuda-oxide ██████████████████████████████████████████ 1.857 (beats top-8!) + cuBLASLt top-8 █████████████████████████████████████████ 1.840 + Exhaustive* ███████████████████████████████████████████ 1.927 + + 8192³ cuda-oxide ███████████████████████████████████████████████ 2.085 (beats BOTH!) + cuBLASLt top-8 ███████████████████████████████████████████████ 2.080 + Exhaustive* ███████████████████████████████████████████████ 2.084 + + 16384³ cuda-oxide ████████████████████████████████████████████████ 2.121 (beats BOTH!) + cuBLASLt top-8 ██████████████████████████████████████████████ 2.065 + Exhaustive* ███████████████████████████████████████████████ 2.103 +``` + ## Final design ```text KernelFamily: gemm_sol_final/output_tile@1 -├── M256xN256 -> existing small resource envelope -└── M512xN256 -> existing large resource envelope +├── M256xN256 -> 4096 / 8192 entry +└── M512xN256 -> 16384 entry Auto selector: 4096 / 8192 -> M256xN256 @@ -20,8 +38,11 @@ Auto selector: ``` Both entries use CLC work scheduling, a two-CTA cluster, `cta_group::2` UMMA, -four TMA/shared-memory pipeline stages, two TMEM accumulator halves, -compiler-directed K-loop unrolling, and L2-aware output-tile ordering. +a seven-stage TMA/shared-memory pipeline (`224 KiB` of contiguous A/B stages), +two TMEM accumulator halves, register-direct 128-bit epilogue stores (`U32x4`), +wave-tail tile balancing (`M256×N192` / `M256×N224`), overlapped CLC +work-stealing, compiler-directed K-loop unrolling, and L2-aware output-tile +ordering. ## What moved performance @@ -105,6 +126,7 @@ From the cuda-oxide repository root: GEMM_SOL_MODE=validate cargo oxide run gemm_sol_final # Fixed 4096/8192/16384 benchmark; 16K dispatches to M512xN256. +# Defaults: 2000 warmup, median of five 301-launch CUDA-event batches. GEMM_SOL_MODE=bench cargo oxide run gemm_sol_final # Print the same dispatch decisions without opening a CUDA context. @@ -135,6 +157,35 @@ bash build.sh The target writes BF16 while the closest supported reference uses FP16 input, FP32 compute, and FP16 output. The output conversion differs, but both paths -write two bytes per element. The helper requests cuBLASLt's first heuristic -candidate, so its result is a reproducible live comparison rather than an -exhaustive autotune of every cuBLASLt algorithm. +write two bytes per element. The helper defaults to timing and selecting among cuBLASLt's top eight +heuristic candidates. Run both that mode and the wider configuration search: + +```bash +./cublaslt_bench --mode both +``` + +The exhaustive-grid mode tests every returned algorithm ID and its exposed +tile, stage, cluster, custom-option, swizzle, and supported reduction settings, +with the finite split-K grid `0,1,2,3,4,5,6,8,12,16,32`. Unlike NVIDIA's +illustrative custom-find sample, it does not stop after 16 IDs or 100 successful +configurations. This is an exhaustive search of the stated grid, not a claim +that every possible cuBLASLt implementation or split-K integer was explored. +It retains the heuristic candidates in final selection. + +Both the Rust kernel and the helper default to 2,000 warmup launches, then five +independent CUDA-event batches of 301 launches. Reported throughput uses the +median batch mean, with tuning and validation excluded. Timing inputs are +zeroed in both paths. Each selected cuBLASLt algorithm additionally passes 64 +nonzero CPU dot-product checks at each measured size; the Rust full-output +validator checks both kernel variants at 4096³. Those 4K checks do not establish +full-output correctness at 8K or 16K. + +Rust timing overrides are `GEMM_SOL_WARMUP`, `GEMM_SOL_ITERS`, and +`GEMM_SOL_REPEATS` (positive odd repeat count); they are forwarded to the live +helper. `GEMM_CUBLAS_MODE=top8|exhaustive|both` selects its search mode. +`GEMM_CUBLAS_SKIP=1` omits the helper for a separately collected comparison. +The helper accepts corresponding `--warmup`, `--iters`, and `--repeats` +options, `--n` for one square size, and emits `RESULT` JSON lines including raw +batch timings, selected configuration, workspace, and validation scope. Results +are reported per shape without aggregating the three shapes into a geomean. +See [bench/README.md](bench/README.md) for search and output-layout details. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/README.md index 1b3bf2e9ce..1d28a85ed4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/README.md @@ -5,7 +5,7 @@ Blackwell GPUs. | Script | What it measures | Dependencies | |------------------------|----------------------------------------|---------------------------| -| `cublaslt_bench.c` | Live heuristic cuBLASLt GEMM reference | CUDA toolkit (C compiler) | +| `cublaslt_bench.c` | Top-8 and exhaustive-grid tuned cuBLASLt reference | CUDA toolkit (C compiler) | | `cublas_sol_bench.py` | Legacy `cublasGemmEx` comparison | numpy, CUDA toolkit | | `cutlass_sol_bench.py` | CUTLASS CuTe DSL FP16 GEMM throughput | nvidia-cutlass-dsl, torch | @@ -20,16 +20,19 @@ pinned stack, cuBLASLt does not expose that exact mixed input/output combination, so the harness parses the benchmark's FP16 A/B, FP32-compute, FP16-output section. That is the closest supported path with the same input type and two-byte output -width, but its final conversion differs from the kernel. The benchmark also -prints a BF16-input/BF16-output section for context. +width, but its final conversion differs from the kernel. The helper's TN +column-major output also differs in storage order from the Rust row-major +output; all measured problems are square. Both compute A×Bᵀ with contiguous K +inputs. `--output bf16` probes exact output-type support; cuBLASLt 13.8.0 on the +B200 returned `CUBLAS_STATUS_NOT_SUPPORTED` for FP16 inputs and BF16 output. Historical throughput from another GPU, driver, or toolkit is not a valid baseline. Build and run the benchmark on the same host as every kernel trial; -the accepted final measurements are recorded in `../README.md`. +the originating B300 checkpoint is recorded in `../README.md`. ## Requirements -- **GPU**: Datacenter Blackwell. Verified on the current B300 (sm_103a). +- **GPU**: Datacenter Blackwell. Measured on B200 (sm_100a); originating results used B300 (sm_103a). - **CUDA Toolkit**: 13.0+ at build time, plus a compatible NVIDIA driver at runtime - **Python**: 3.12+ (for Python benchmarks only) @@ -73,9 +76,45 @@ pip install nvidia-cutlass-dsl torch # for cutlass_sol_bench.py ./cublaslt_bench ``` -Tests BF16 and FP16 regular GEMM with FP32 compute at 4K, 8K, and 16K -using TN format and heuristic algorithm selection. The harness consumes only -the FP16 section described above. +Tests FP16 A/B, FP32 compute, FP16 output at 4K, 8K, and 16K in TN format. +The default `--mode top8` measures all eight heuristic candidates, then selects +by the median of five 31-launch refinement batches after ten warmups per batch. + +```bash +# Top-8 and the wider search, for every fixed shape: +./cublaslt_bench --mode both +# One shape with an explicit final timing protocol: +./cublaslt_bench --mode exhaustive --n 8192 --warmup 2000 --iters 301 --repeats 5 +``` + +The wider search follows NVIDIA's +[custom-find sample](https://github.com/NVIDIA/CUDALibrarySamples/blob/main/cuBLASLt/Common/LtMatmulCustomFind.h), +without its 16-ID or 100-result truncation. It enumerates all returned algorithm +IDs, exposed tile/stage lists, cluster shapes, custom options, swizzles, and +supported reductions over the finite split-K grid +`0,1,2,3,4,5,6,8,12,16,32`. Both 0 and 1 mean no split; current algorithms may +require 1. Only configurations that pass `cublasLtMatmulAlgoCheck` and fit the +32 MiB workspace limit are timed. The search is exhaustive over this documented +grid, not over every possible split-K integer. The inner MMA shape uses its +library default, as in NVIDIA's custom-find sample. + +Grid screening uses one warmup and five launches per valid configuration. +The fastest 16 screening candidates plus **all** heuristic candidates undergo +five 31-launch refinement batches; the lowest median wins. The winner is +validated on nonzero signed 1/8 data using 64 CPU dot products at distributed +coordinates, then remeasured from fresh zeroed inputs with the final protocol. +The final number is not the minimum tuning sample. Every matmul return status +is checked; failed final timing or correctness checks prevent a result. + +`RESULT` JSON lines preserve the five raw final batch means, selected algorithm +configuration, workspace usage, search counts, and validation scope. cuBLASLt +only needs the winning configuration after tuning, so search time is excluded +from steady-state throughput. The Rust consumer defaults to top-8 tuning and +forwards any `GEMM_SOL_WARMUP`, `GEMM_SOL_ITERS`, or `GEMM_SOL_REPEATS` override. +Set `GEMM_CUBLAS_MODE` to `top8`, `exhaustive`, or `both` to change its mode; +`GEMM_CUBLAS_SKIP=1` supports collecting the two implementations separately. +When the helper runs in `both` mode, the Rust summary consumes the top-8 row; +the standalone `RESULT` records retain both independently measured results. ### cuBLAS (legacy) SoL @@ -101,7 +140,12 @@ software-pipelined K-loop and TMA loads. ## Notes - All benchmarks use GPU-side timing (CUDA events), not wall-clock. -- Warmup: 10 iterations. Timed: 100 iterations. +- Rust and cuBLASLt defaults: 2,000 warmup launches followed by five independent + batches of 301 launches, reported as median batch mean. Timing uses zero data. +- GPU clocks remain under the driver's control; record GPU model, SM count, + driver/toolkit/library versions, power limit, and other active GPU workloads. +- Historical B300 results and legacy Python scripts retain their original + 10-warmup/100-iteration protocol; do not mix them with current measurements. - The `venv/` directory is gitignored — each machine creates its own. - The compiled `cublaslt_bench` binary is gitignored — rebuild from `cublaslt_bench.c` on each machine. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/cublaslt_bench.c b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/cublaslt_bench.c index ec42a0f1cb..5a6bd6979d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/cublaslt_bench.c +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/bench/cublaslt_bench.c @@ -1,163 +1,465 @@ /* * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 + * + * Top-8 tuning and the NVIDIA LtMatmulCustomFind search grid, without the + * sample's 16-algorithm / 100-success limits. Split-K is a finite grid, not + * every integer: 0,1,2,3,4,5,6,8,12,16,32. Search time is excluded from results. + * https://github.com/NVIDIA/CUDALibrarySamples/tree/main/cuBLASLt */ - +#include +#include +#include +#include +#include #include #include -#include -#include - -#define CHECK_CUDA(call) do { \ - cudaError_t err = (call); \ - if (err != cudaSuccess) { \ - fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \ - exit(1); \ - } \ -} while(0) - -#define CHECK_CUBLAS(call) do { \ - cublasStatus_t status = (call); \ - if (status != CUBLAS_STATUS_SUCCESS) { \ - fprintf(stderr, "cuBLAS error at %s:%d: status=%d (%s)\n", __FILE__, __LINE__, \ - (int)status, cublasLtGetStatusString(status)); \ - exit(1); \ - } \ -} while(0) - -static double bench_cublaslt(cublasLtHandle_t handle, cudaStream_t stream, - int M, int N, int K, - cudaDataType_t dataType, const char *label) { - size_t elem_bytes = 2; - size_t workspace_size = 32 * 1024 * 1024; - - void *d_A, *d_B, *d_C, *d_workspace; - CHECK_CUDA(cudaMalloc(&d_A, (size_t)M * K * elem_bytes)); - CHECK_CUDA(cudaMalloc(&d_B, (size_t)K * N * elem_bytes)); - CHECK_CUDA(cudaMalloc(&d_C, (size_t)M * N * elem_bytes)); - CHECK_CUDA(cudaMalloc(&d_workspace, workspace_size)); - CHECK_CUDA(cudaMemset(d_A, 0, (size_t)M * K * elem_bytes)); - CHECK_CUDA(cudaMemset(d_B, 0, (size_t)K * N * elem_bytes)); - CHECK_CUDA(cudaMemset(d_C, 0, (size_t)M * N * elem_bytes)); - - cublasLtMatmulDesc_t matmulDesc; - CHECK_CUBLAS(cublasLtMatmulDescCreate(&matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); - - cublasOperation_t transa = CUBLAS_OP_T; - cublasOperation_t transb = CUBLAS_OP_N; - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa))); - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(matmulDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb))); - - cublasLtMatrixLayout_t layoutA, layoutB, layoutC, layoutD; - CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&layoutA, dataType, K, M, K)); - CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&layoutB, dataType, K, N, K)); - CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&layoutC, dataType, M, N, M)); - CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&layoutD, dataType, M, N, M)); - - cublasLtMatmulPreference_t preference; - CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&preference)); - CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( - preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, - &workspace_size, sizeof(workspace_size))); - - cublasLtMatmulHeuristicResult_t heuristicResult; - int returnedAlgoCount = 0; - cublasStatus_t heurStatus = cublasLtMatmulAlgoGetHeuristic( - handle, matmulDesc, layoutA, layoutB, layoutC, layoutD, - preference, 1, &heuristicResult, &returnedAlgoCount); - - double tflops = 0.0; - if (heurStatus == CUBLAS_STATUS_SUCCESS && returnedAlgoCount > 0) { - float alpha = 1.0f, beta = 0.0f; - cudaEvent_t start, stop; - CHECK_CUDA(cudaEventCreate(&start)); - CHECK_CUDA(cudaEventCreate(&stop)); - - for (int i = 0; i < 10; i++) { - cublasLtMatmul(handle, matmulDesc, - &alpha, d_A, layoutA, d_B, layoutB, &beta, - d_C, layoutC, d_C, layoutD, - &heuristicResult.algo, d_workspace, workspace_size, stream); +#include +#define CHECK_CUDA(call) \ + do { \ + cudaError_t e = (call); \ + if (e != cudaSuccess) { \ + fprintf(stderr, "CUDA: %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(e)); \ + exit(1); \ + } \ + } while (0) +#define CHECK_BLAS(call) \ + do { \ + cublasStatus_t s = (call); \ + if (s != CUBLAS_STATUS_SUCCESS) { \ + fprintf(stderr, "cuBLAS: %s:%d: %s\n", __FILE__, __LINE__, \ + cublasLtGetStatusString(s)); \ + exit(1); \ + } \ + } while (0) +#define FINALISTS 16 +static const size_t workspace_size = 32 * 1024 * 1024; +static int warmup = 2000, iters = 301, repeats = 5; +static const char *mode = "top8"; +static cudaDataType_t output_type = CUDA_R_16F; +static const char *output_name = "fp16"; +typedef struct { + cublasLtHandle_t h; + cudaStream_t stream; + cudaEvent_t start, stop; + cublasLtMatmulDesc_t op; + cublasLtMatrixLayout_t a, b, c; + void *A, *B, *C, *workspace; + int n; +} Problem; +typedef struct { + cublasLtMatmulAlgo_t algo; + double us; + size_t workspace; +} Candidate; +static cublasStatus_t launch(Problem *p, const cublasLtMatmulAlgo_t *a) { + float alpha = 1, beta = 0; + return cublasLtMatmul(p->h, p->op, &alpha, p->A, p->a, p->B, p->b, &beta, p->C, p->c, p->C, + p->c, a, p->workspace, workspace_size, p->stream); +} +static double measure(Problem *p, const cublasLtMatmulAlgo_t *a, int nw, int ni) { + for (int i = 0; i < nw; i++) + if (launch(p, a) != CUBLAS_STATUS_SUCCESS) + return DBL_MAX; + CHECK_CUDA(cudaStreamSynchronize(p->stream)); + CHECK_CUDA(cudaEventRecord(p->start, p->stream)); + for (int i = 0; i < ni; i++) + if (launch(p, a) != CUBLAS_STATUS_SUCCESS) + return DBL_MAX; + CHECK_CUDA(cudaEventRecord(p->stop, p->stream)); + CHECK_CUDA(cudaEventSynchronize(p->stop)); + float ms; + CHECK_CUDA(cudaEventElapsedTime(&ms, p->start, p->stop)); + return ms * 1000.0 / ni; +} +static int compare(const void *a, const void *b) { + double x = *(const double *)a, y = *(const double *)b; + return (x > y) - (x < y); +} +static double median(double *values, int count) { + qsort(values, count, sizeof(double), compare); + return values[count / 2]; +} +static int cap(const cublasLtMatmulAlgo_t *a, cublasLtMatmulAlgoCapAttributes_t attr) { + int v = 0; + CHECK_BLAS(cublasLtMatmulAlgoCapGetAttribute(a, attr, &v, sizeof(v), NULL)); + return v; +} +static int *cap_list(const cublasLtMatmulAlgo_t *a, cublasLtMatmulAlgoCapAttributes_t attr, + int *count) { + size_t bytes = 0; + CHECK_BLAS(cublasLtMatmulAlgoCapGetAttribute(a, attr, NULL, 0, &bytes)); + *count = bytes / sizeof(int); + int *v = calloc(*count ? *count : 1, sizeof(int)); + if (!v) + exit(1); + if (*count) + CHECK_BLAS(cublasLtMatmulAlgoCapGetAttribute(a, attr, v, bytes, NULL)); + else + *count = 1; + return v; +} +static int set(cublasLtMatmulAlgo_t *a, cublasLtMatmulAlgoConfigAttributes_t attr, int v) { + return cublasLtMatmulAlgoConfigSetAttribute(a, attr, &v, sizeof(v)) == CUBLAS_STATUS_SUCCESS; +} +static int get(const cublasLtMatmulAlgo_t *a, cublasLtMatmulAlgoConfigAttributes_t attr) { + int v = 0; + CHECK_BLAS(cublasLtMatmulAlgoConfigGetAttribute(a, attr, &v, sizeof(v), NULL)); + return v; +} +static void describe(const cublasLtMatmulAlgo_t *a) { + uint16_t cluster = 0; + CHECK_BLAS(cublasLtMatmulAlgoConfigGetAttribute(a, CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID, + &cluster, sizeof(cluster), NULL)); + printf("\"algo_id\":%d,\"tile\":%d,\"stages\":%d,\"cluster\":%u,\"split_k\":%d,\"reduction\":%" + "d,\"swizzle\":%d,\"custom\":%d", + get(a, CUBLASLT_ALGO_CONFIG_ID), get(a, CUBLASLT_ALGO_CONFIG_TILE_ID), + get(a, CUBLASLT_ALGO_CONFIG_STAGES_ID), cluster, get(a, CUBLASLT_ALGO_CONFIG_SPLITK_NUM), + get(a, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME), + get(a, CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING), get(a, CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION)); +} +static void keep(Candidate *best, int *count, Candidate c) { + int pos = *count; + if (pos == FINALISTS) { + if (c.us >= best[pos - 1].us) + return; + pos--; + } else + (*count)++; + while (pos > 0 && c.us < best[pos - 1].us) { + best[pos] = best[pos - 1]; + pos--; + } + best[pos] = c; +} +static int trial(Problem *p, const cublasLtMatmulAlgo_t *a, Candidate *best, int *count) { + cublasLtMatmulHeuristicResult_t r; + if (cublasLtMatmulAlgoCheck(p->h, p->op, p->a, p->b, p->c, p->c, a, &r) != + CUBLAS_STATUS_SUCCESS || + r.state != CUBLAS_STATUS_SUCCESS || r.workspaceSize > workspace_size) + return 0; + double us = measure(p, a, 1, 5); + if (us == DBL_MAX) + return 0; + Candidate c = {*a, us, r.workspaceSize}; + keep(best, count, c); + return 1; +} +/* Nonzero, exactly representable signed 1/8 inputs; independent CPU dot-product + * checks at 64 distributed output coordinates. This is outside timed regions. */ +static uint16_t small_half(int v) { + if (!v) + return 0; + int sign = v < 0 ? 0x8000 : 0; + v = abs(v); + return sign | (v == 1 ? 0x3000 : v == 2 ? 0x3400 : 0x3600); +} +static float from16(uint16_t x) { + if (output_type == CUDA_R_16BF) { + uint32_t b = (uint32_t)x << 16; + float f; + memcpy(&f, &b, 4); + return f; + } + int exp = (x >> 10) & 31, mant = x & 1023; + float f = exp == 31 ? (mant ? NAN : INFINITY) + : exp ? ldexpf(1 + mant / 1024.0f, exp - 15) + : ldexpf(mant / 1024.0f, -14); + return x & 0x8000 ? -f : f; +} +static int av(int row, int k) { + return (int)(((uint64_t)row * 13 + (uint64_t)k * 7 + (uint64_t)(k / 11) * 3) % 7) - 3; +} +static int bv(int col, int k) { + return (int)(((uint64_t)col * 11 + (uint64_t)k * 5 + (uint64_t)(k / 13) * 2) % 7) - 3; +} +static void validate(Problem *p, const cublasLtMatmulAlgo_t *a) { + size_t count = (size_t)p->n * p->n, bytes = count * 2; + uint16_t *host = malloc(bytes); + if (!host) + exit(1); + for (int row = 0; row < p->n; row++) + for (int k = 0; k < p->n; k++) + host[(size_t)row * p->n + k] = small_half(av(row, k)); + CHECK_CUDA(cudaMemcpy(p->A, host, bytes, cudaMemcpyHostToDevice)); + for (int col = 0; col < p->n; col++) + for (int k = 0; k < p->n; k++) + host[(size_t)col * p->n + k] = small_half(bv(col, k)); + CHECK_CUDA(cudaMemcpy(p->B, host, bytes, cudaMemcpyHostToDevice)); + CHECK_BLAS(launch(p, a)); + CHECK_CUDA(cudaStreamSynchronize(p->stream)); + CHECK_CUDA(cudaMemcpy(host, p->C, bytes, cudaMemcpyDeviceToHost)); + for (int i = 0; i < 64; i++) { + int row = (i * 7919 + 17) % p->n, col = (i * 6701 + 31) % p->n; + double expected = 0; + for (int k = 0; k < p->n; k++) + expected += av(row, k) * bv(col, k) / 64.0; + double actual = from16(host[(size_t)col * p->n + row]); + double tolerance = + fmax(0.03125, fabs(expected) * (output_type == CUDA_R_16BF ? 0.004 : 0.0005)); + if (!isfinite(actual) || fabs(actual - expected) > tolerance) { + fprintf(stderr, "VALIDATION FAILED n=%d row=%d col=%d actual=%g expected=%g\n", p->n, + row, col, actual, expected); + exit(1); } - CHECK_CUDA(cudaStreamSynchronize(stream)); - - int iters = 100; - CHECK_CUDA(cudaEventRecord(start, stream)); - for (int i = 0; i < iters; i++) { - cublasLtMatmul(handle, matmulDesc, - &alpha, d_A, layoutA, d_B, layoutB, &beta, - d_C, layoutC, d_C, layoutD, - &heuristicResult.algo, d_workspace, workspace_size, stream); + } + free(host); + CHECK_CUDA(cudaMemset(p->A, 0, bytes)); + CHECK_CUDA(cudaMemset(p->B, 0, bytes)); + CHECK_CUDA(cudaMemset(p->C, 0, bytes)); + CHECK_CUDA(cudaStreamSynchronize(p->stream)); +} +static Candidate refine(Problem *p, Candidate *best, int count) { + Candidate winner = best[0]; + winner.us = DBL_MAX; + for (int j = 0; j < count; j++) { + double ts[5]; + for (int r = 0; r < 5; r++) + ts[r] = measure(p, &best[j].algo, 10, 31); + double us = median(ts, 5); + if (us < winner.us) { + winner = best[j]; + winner.us = us; } - CHECK_CUDA(cudaEventRecord(stop, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - float elapsed_ms; - CHECK_CUDA(cudaEventElapsedTime(&elapsed_ms, start, stop)); - double avg_ms = elapsed_ms / iters; - double flops = 2.0 * (double)M * N * K; - tflops = (flops / (avg_ms / 1000.0)) / 1e12; - printf(" %-24s %5dx%5dx%5d %8.4f ms %8.1f TFLOPS\n", - label, M, N, K, avg_ms, tflops); - - CHECK_CUDA(cudaEventDestroy(start)); - CHECK_CUDA(cudaEventDestroy(stop)); - } else { - printf(" %-24s %5dx%5dx%5d NO ALGORITHM (status=%d)\n", - label, M, N, K, (int)heurStatus); } - - cublasLtMatmulPreferenceDestroy(preference); - cublasLtMatrixLayoutDestroy(layoutA); - cublasLtMatrixLayoutDestroy(layoutB); - cublasLtMatrixLayoutDestroy(layoutC); - cublasLtMatrixLayoutDestroy(layoutD); - cublasLtMatmulDescDestroy(matmulDesc); - CHECK_CUDA(cudaFree(d_A)); - CHECK_CUDA(cudaFree(d_B)); - CHECK_CUDA(cudaFree(d_C)); - CHECK_CUDA(cudaFree(d_workspace)); - - return tflops; -} - -int main() { - int dev; - CHECK_CUDA(cudaGetDevice(&dev)); + return winner; +} +static void report(Problem *p, Candidate chosen, const char *label, long checked, long valid, + int algo_count) { + validate(p, &chosen.algo); + double *times = calloc(repeats, sizeof(double)); + /* Long warmup once, then repeated independent event batches. */ + for (int i = 0; i < warmup; i++) + CHECK_BLAS(launch(p, &chosen.algo)); + CHECK_CUDA(cudaStreamSynchronize(p->stream)); + for (int r = 0; r < repeats; r++) { + times[r] = measure(p, &chosen.algo, 0, iters); + if (!isfinite(times[r]) || times[r] <= 0 || times[r] == DBL_MAX) { + fprintf(stderr, "Final timing failed; no result will be reported.\n"); + exit(1); + } + } + printf("RESULT " + "{\"implementation\":\"cublaslt\",\"mode\":\"%s\",\"n\":%d,\"input\":\"fp16\"," + "\"output\":\"%s\",\"compute\":\"fp32\",\"workspace_limit\":%zu,\"workspace_used\":%zu," + "\"warmup\":%d,\"iters\":%d,\"repeats\":%d,\"samples_us\":[", + label, p->n, output_name, workspace_size, chosen.workspace, warmup, iters, repeats); + for (int r = 0; r < repeats; r++) + printf("%s%.9f", r ? "," : "", times[r]); + double us = median(times, repeats); + double tf = 2.0 * p->n * p->n * p->n / us / 1e6; + printf("],\"median_us\":%.9f,\"tflops\":%.9f,\"configs_checked\":%ld,\"configs_timed\":%ld," + "\"algo_ids\":%d,\"validation\":\"64-nonzero-cpu-dot-products-pass\",", + us, tf, checked, valid, algo_count); + describe(&chosen.algo); + printf("}\n"); + if (!strcmp(label, "top8") || !strcmp(mode, "exhaustive")) + printf(" FP16 FP32 compute %dx%dx%d %.9f ms %.9f TFLOPS\n", p->n, p->n, p->n, us / 1000, + tf); + free(times); + fflush(stdout); +} +static void bench(cublasLtHandle_t h, int n) { + Problem p = {0}; + p.h = h; + p.n = n; + CHECK_CUDA(cudaStreamCreate(&p.stream)); + CHECK_CUDA(cudaEventCreate(&p.start)); + CHECK_CUDA(cudaEventCreate(&p.stop)); + size_t bytes = (size_t)n * n * 2; + CHECK_CUDA(cudaMalloc(&p.A, bytes)); + CHECK_CUDA(cudaMalloc(&p.B, bytes)); + CHECK_CUDA(cudaMalloc(&p.C, bytes)); + CHECK_CUDA(cudaMalloc(&p.workspace, workspace_size)); + CHECK_CUDA(cudaMemset(p.A, 0, bytes)); + CHECK_CUDA(cudaMemset(p.B, 0, bytes)); + CHECK_CUDA(cudaMemset(p.C, 0, bytes)); + CHECK_BLAS(cublasLtMatmulDescCreate(&p.op, CUBLAS_COMPUTE_32F, CUDA_R_32F)); + cublasOperation_t t = CUBLAS_OP_T, normal = CUBLAS_OP_N; + CHECK_BLAS(cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSA, &t, sizeof(t))); + CHECK_BLAS( + cublasLtMatmulDescSetAttribute(p.op, CUBLASLT_MATMUL_DESC_TRANSB, &normal, sizeof(normal))); + CHECK_BLAS(cublasLtMatrixLayoutCreate(&p.a, CUDA_R_16F, n, n, n)); + CHECK_BLAS(cublasLtMatrixLayoutCreate(&p.b, CUDA_R_16F, n, n, n)); + CHECK_BLAS(cublasLtMatrixLayoutCreate(&p.c, output_type, n, n, n)); + cublasLtMatmulPreference_t pref; + CHECK_BLAS(cublasLtMatmulPreferenceCreate(&pref)); + CHECK_BLAS(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, sizeof(workspace_size))); + cublasLtMatmulHeuristicResult_t heur[8] = {0}; + int returned = 0; + cublasStatus_t status = + cublasLtMatmulAlgoGetHeuristic(h, p.op, p.a, p.b, p.c, p.c, pref, 8, heur, &returned); + if (status != CUBLAS_STATUS_SUCCESS || !returned) { + fprintf(stderr, "No supported FP16->%s heuristic for %d (status=%d count=%d)\n", + output_name, n, status, returned); + exit(2); + } + Candidate top[FINALISTS]; + int ntop = 0; + long topvalid = 0; + for (int i = 0; i < returned; i++) + topvalid += trial(&p, &heur[i].algo, top, &ntop); + if (!ntop) { + fprintf(stderr, "No valid top8 candidates\n"); + exit(2); + } + Candidate topwinner = refine(&p, top, ntop); + if (strcmp(mode, "exhaustive")) + report(&p, topwinner, "top8", returned, topvalid, 0); + if (strcmp(mode, "top8")) { + int capacity = 64, count = 0, *ids = NULL; + do { + capacity *= 2; + ids = realloc(ids, capacity * sizeof(int)); + if (!ids) + exit(1); + CHECK_BLAS(cublasLtMatmulAlgoGetIds(h, CUBLAS_COMPUTE_32F, CUDA_R_32F, CUDA_R_16F, + CUDA_R_16F, output_type, output_type, capacity, ids, + &count)); + } while (count == capacity); + Candidate finalists[FINALISTS]; + int nf = 0; + for (int i = 0; i < ntop; i++) + keep(finalists, &nf, top[i]); + const int splits[] = {0, 1, 2, 3, 4, 5, 6, 8, 12, 16, 32}; + long checked = 0, valid = 0; + for (int ai = 0; ai < count; ai++) { + cublasLtMatmulAlgo_t a; + if (cublasLtMatmulAlgoInit(h, CUBLAS_COMPUTE_32F, CUDA_R_32F, CUDA_R_16F, CUDA_R_16F, + output_type, output_type, ids[ai], + &a) != CUBLAS_STATUS_SUCCESS) + continue; + int tiles_count, stages_count; + int *tiles = cap_list(&a, CUBLASLT_ALGO_CAP_TILE_IDS, &tiles_count), + *stages = cap_list(&a, CUBLASLT_ALGO_CAP_STAGES_IDS, &stages_count); + int split = cap(&a, CUBLASLT_ALGO_CAP_SPLITK_SUPPORT), + redmask = cap(&a, CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK), + swmax = cap(&a, CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT), + custommax = cap(&a, CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX); + for (int ti = 0; ti < tiles_count; ti++) + for (int si = 0; si < stages_count; si++) + for (uint16_t cl = 0; cl < CUBLASLT_CLUSTER_SHAPE_END; cl++) + for (int cu = 0; cu <= custommax; cu++) + for (int sw = 0; sw <= swmax; sw++) + for (int ki = 0; ki < (split ? 11 : 2); ki++) + for (int red = ki > 1 ? 1 : 0; + red < (int)CUBLASLT_REDUCTION_SCHEME_MASK; + red = red ? red * 2 : 8) { + if (ki > 1 && !(red & redmask)) + continue; + if (!set(&a, CUBLASLT_ALGO_CONFIG_TILE_ID, tiles[ti]) || + !set(&a, CUBLASLT_ALGO_CONFIG_STAGES_ID, stages[si]) || + !set(&a, CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, cu) || + !set(&a, CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, sw) || + !set(&a, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, splits[ki]) || + !set(&a, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, red)) + continue; + if (cublasLtMatmulAlgoConfigSetAttribute( + &a, CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID, &cl, + sizeof(cl)) != CUBLAS_STATUS_SUCCESS) + continue; + checked++; + valid += trial(&p, &a, finalists, &nf); + } + fprintf(stderr, "SEARCH n=%d algo=%d/%d id=%d checked=%ld timed=%ld best_us=%.3f\n", n, + ai + 1, count, ids[ai], checked, valid, finalists[0].us); + fflush(stderr); + free(tiles); + free(stages); + } + /* Remeasure every heuristic candidate as well as the screened grid + * finalists: noisy screening must never evict the heuristic fallback. */ + Candidate combined[FINALISTS + 8]; + memcpy(combined, finalists, nf * sizeof(Candidate)); + memcpy(combined + nf, top, ntop * sizeof(Candidate)); + report(&p, refine(&p, combined, nf + ntop), "exhaustive-grid", checked, valid, count); + free(ids); + } + CHECK_BLAS(cublasLtMatmulPreferenceDestroy(pref)); + CHECK_BLAS(cublasLtMatrixLayoutDestroy(p.a)); + CHECK_BLAS(cublasLtMatrixLayoutDestroy(p.b)); + CHECK_BLAS(cublasLtMatrixLayoutDestroy(p.c)); + CHECK_BLAS(cublasLtMatmulDescDestroy(p.op)); + CHECK_CUDA(cudaFree(p.A)); + CHECK_CUDA(cudaFree(p.B)); + CHECK_CUDA(cudaFree(p.C)); + CHECK_CUDA(cudaFree(p.workspace)); + CHECK_CUDA(cudaEventDestroy(p.start)); + CHECK_CUDA(cudaEventDestroy(p.stop)); + CHECK_CUDA(cudaStreamDestroy(p.stream)); +} +static int positive(const char *s) { + char *end; + long v = strtol(s, &end, 10); + if (*end || v <= 0 || v > 10000000) { + fprintf(stderr, "Invalid positive integer: %s\n", s); + exit(2); + } + return v; +} +int main(int argc, char **argv) { + int only = 0; + const char *env = getenv("GEMM_CUBLAS_MODE"); + if (env) + mode = env; + for (int i = 1; i < argc; i++) { + if (i + 1 == argc) { + fprintf(stderr, + "Usage: %s [--mode top8|exhaustive|both] [--n N] [--warmup N] [--iters N] " + "[--repeats odd-N] [--output fp16|bf16]\n", + argv[0]); + return 2; + } + const char *key = argv[i++], *value = argv[i]; + if (!strcmp(key, "--mode")) + mode = value; + else if (!strcmp(key, "--n")) + only = positive(value); + else if (!strcmp(key, "--warmup")) + warmup = positive(value); + else if (!strcmp(key, "--iters")) + iters = positive(value); + else if (!strcmp(key, "--repeats")) + repeats = positive(value); + else if (!strcmp(key, "--output") && !strcmp(value, "fp16")) { + output_type = CUDA_R_16F; + output_name = "fp16"; + } else if (!strcmp(key, "--output") && !strcmp(value, "bf16")) { + output_type = CUDA_R_16BF; + output_name = "bf16"; + } else { + fprintf(stderr, "Invalid option %s %s\n", key, value); + return 2; + } + } + if ((strcmp(mode, "top8") && strcmp(mode, "exhaustive") && strcmp(mode, "both")) || + !(repeats & 1)) { + fprintf(stderr, "Invalid mode or even repeat count\n"); + return 2; + } + int dev, driver, runtime; struct cudaDeviceProp prop; + CHECK_CUDA(cudaGetDevice(&dev)); CHECK_CUDA(cudaGetDeviceProperties(&prop, dev)); - printf("GPU: %s, sm_%d%d, %d SMs\n", prop.name, prop.major, prop.minor, prop.multiProcessorCount); - printf("cublasLt version: %zu\n\n", cublasLtGetVersion()); - - cublasLtHandle_t handle; - CHECK_CUBLAS(cublasLtCreate(&handle)); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - int sizes[][3] = { - {4096, 4096, 4096}, - {8192, 8192, 8192}, - {16384, 16384, 16384}, - }; - int num_sizes = sizeof(sizes) / sizeof(sizes[0]); - - printf("==========================================================\n"); - printf(" cublasLtMatmul Benchmark (TN format, FP32 compute)\n"); - printf("==========================================================\n"); - - printf("\n--- BF16 ---\n"); - for (int i = 0; i < num_sizes; i++) - bench_cublaslt(handle, stream, sizes[i][0], sizes[i][1], sizes[i][2], CUDA_R_16BF, "BF16 FP32 compute"); - - printf("\n--- FP16 ---\n"); - for (int i = 0; i < num_sizes; i++) - bench_cublaslt(handle, stream, sizes[i][0], sizes[i][1], sizes[i][2], CUDA_R_16F, "FP16 FP32 compute"); - - printf("\n==========================================================\n"); - - cublasLtDestroy(handle); - CHECK_CUDA(cudaStreamDestroy(stream)); - + CHECK_CUDA(cudaDriverGetVersion(&driver)); + CHECK_CUDA(cudaRuntimeGetVersion(&runtime)); + printf("GPU: %s, sm_%d%d, %d SMs\ncuBLASLt: %zu, CUDA runtime: %d, driver API: %d\n", prop.name, + prop.major, prop.minor, prop.multiProcessorCount, cublasLtGetVersion(), runtime, driver); + printf("Protocol: zero-data timing; CUDA events around launch batches; median of %d batches; " + "%d warmup, %d launches/batch; tuning excluded; 32 MiB workspace\n", + repeats, warmup, iters); + printf("--- FP16 ---\n"); + fflush(stdout); + cublasLtHandle_t h; + CHECK_BLAS(cublasLtCreate(&h)); + if (only) + bench(h, only); + else { + bench(h, 4096); + bench(h, 8192); + bench(h, 16384); + } + CHECK_BLAS(cublasLtDestroy(h)); return 0; } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/kernels.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/kernels.rs index 8b8e7eb50a..0e1feff94e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/kernels.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/kernels.rs @@ -1,48 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -// Size-specialized device kernels for the canonical gemm_sol_final example. -// include!d into main.rs so the #[cuda_module] macro sees an inline module -// (it rejects file modules). Host validation and benchmarking remain in main.rs. #[cuda_module] mod kernels { use super::*; - /// Target kernel: CLC + TMA multicast with 4 SMEM stages and no MCAST_BAR. - /// - /// Blackwell GEMM kernel: CLC + cta_group::2 + 4-stage SMEM pipeline. - /// - /// Architecture: 6 warps per CTA, cluster size 2 (CTA pairs). - /// - Warp 4 (TMA): loads A/B tiles from global via TMA into shared memory. - /// - Warp 5 (MMA): consumes SMEM tiles via pair-UMMA (tcgen05, cta_group::2). - /// - Warps 0-3 (Epilogue): read accumulators from TMEM, convert f32->bf16, store to global. - /// - /// Pipeline stages: 4 SMEM stages (TMA_BAR0..3 / MMA_BAR0..3), 2 accumulator stages - /// (ACCUM_FULL0/1, ACCUM_EMPTY0/1), plus TILE_READY for TMA→MMA/Epilogue tile handoff. - /// - /// Key synchronization protocol (cta_group::2 barrier aliasing): - /// Both CTAs in a pair issue TMA loads, but the barrier pointer is masked with - /// PEER_BIT_MASK (0xFEFFFFF8) before being passed to the TMA instruction. This - /// clears bit 24, redirecting both CTAs' completion signals to the leader CTA's - /// (rank 0) barrier. Consequently: - /// - Only the leader sets expect_tx (doubled: both CTAs' bytes land on one barrier). - /// - Only the leader waits on TMA barriers in the MMA warp. - /// - The follower's MMA warp skips TMA waits since pair-UMMA is leader-issued. - /// MMA barriers still use normal multicast (tcgen05_commit_multicast_cg2 with - /// CTA_MASK_PAIR=0b11), so both CTAs receive MMA completion signals. - /// - /// CLC work-stealing: rank 0 issues clc_try_cancel_multicast; both CTAs receive the - /// response via CLC_BAR. Tile indices are derived by dividing the CLC first_ctaid_x - /// by the cluster size (2), NOT using raw CTA IDs. - /// - /// # Shape contract - /// - /// `n` must be a multiple of 256. The fixed host grid is expressed in 128-column - /// work IDs, and this kernel drains its upper half to form 256-column tiles. - /// - /// `k` must be a multiple of 256. With a K tile size of 64, this makes `k_iters` - /// a multiple of four, so every output tile starts on pipeline stage 0. The MMA - /// consumer relies on that reset to keep its unrolled stage selection constant. #[kernel] #[cluster_launch(2, 1, 1)] pub unsafe fn gemm_sol_clc_multicast_4_stage_pipeline( @@ -54,41 +16,28 @@ mod kernels { tiles_m: u32, _tiles_n: u32, ) { - unsafe { - static mut SMEM_A0: SharedArray = SharedArray::UNINIT; - static mut SMEM_A1: SharedArray = SharedArray::UNINIT; - static mut SMEM_A2: SharedArray = SharedArray::UNINIT; - static mut SMEM_A3: SharedArray = SharedArray::UNINIT; - static mut SMEM_B0: SharedArray = SharedArray::UNINIT; - static mut SMEM_B1: SharedArray = SharedArray::UNINIT; - static mut SMEM_B2: SharedArray = SharedArray::UNINIT; - static mut SMEM_B3: SharedArray = SharedArray::UNINIT; - static mut SMEM_OUT: SharedArray = SharedArray::UNINIT; + unsafe { + static mut SMEM_BUF: SharedArray = SharedArray::UNINIT; static mut TMEM_ADDR: SharedArray = SharedArray::UNINIT; - static mut TILE_INFO: SharedArray = SharedArray::UNINIT; - - // 4-stage TMA <-> MMA pipeline - static mut TMA_BAR0: Barrier = Barrier::UNINIT; - static mut TMA_BAR1: Barrier = Barrier::UNINIT; - static mut TMA_BAR2: Barrier = Barrier::UNINIT; - static mut TMA_BAR3: Barrier = Barrier::UNINIT; - static mut MMA_BAR0: Barrier = Barrier::UNINIT; - static mut MMA_BAR1: Barrier = Barrier::UNINIT; - static mut MMA_BAR2: Barrier = Barrier::UNINIT; - static mut MMA_BAR3: Barrier = Barrier::UNINIT; + static mut TILE_INFO: SharedArray = SharedArray::UNINIT; + + // Up to 8-stage TMA <-> MMA pipeline + static mut TMA_BAR: SharedArray = SharedArray::UNINIT; + static mut MMA_BAR: SharedArray = SharedArray::UNINIT; static mut ACCUM_FULL0: Barrier = Barrier::UNINIT; static mut ACCUM_FULL1: Barrier = Barrier::UNINIT; static mut ACCUM_EMPTY0: Barrier = Barrier::UNINIT; static mut ACCUM_EMPTY1: Barrier = Barrier::UNINIT; - static mut TILE_READY: Barrier = Barrier::UNINIT; + static mut TILE_READY: SharedArray = SharedArray::UNINIT; static mut CLC_RESPONSE: SharedArray = SharedArray::UNINIT; static mut CLC_BAR: Barrier = Barrier::UNINIT; const A_TILE_BYTES: u32 = 128 * 64 * 2; - const B_PANEL_BYTES: u32 = 64 * 64 * 2; - const B_TILE_BYTES: u32 = 128 * 64 * 2; // 128 B rows per CTA (split by rank) + const B_PANEL32_BYTES: u32 = 32 * 64 * 2; + const B_TILE_BYTES_256: u32 = 128 * 64 * 2; + const B_TILE_BYTES_192: u32 = 96 * 64 * 2; const SBO_BYTES: u32 = 1024; const LBO_BYTES: u32 = 16; const SWIZZLE_128B: u8 = 2; @@ -98,14 +47,10 @@ mod kernels { const NUM_ACCUM_STAGES: u32 = 2; const ACCUM_STAGE_COLS: u32 = 256; const CTA_MASK_PAIR: u16 = 0b11; - // Clears bit 24 (CTA rank within pair) + alignment bits 2:0 of a shared - // memory barrier address, redirecting TMA completions to the leader CTA's - // barrier. const PEER_BIT_MASK: u32 = 0xFEFFFFF8; - // L2 cache-blocking: the 256-column output tile doubles the physical - // width of each N tile, so preserve the accepted cache-band widths by - // halving the prior logical group sizes. - let swizzle_g: u32 = if tiles_m <= 16 { 2 } else { 8 }; + let num_stages: usize = 7; + let b_smem_offset: usize = 7 * 16384; + let swizzle_g: u32 = if tiles_m == 16 { 2 } else { 8 }; let n = n as u32; let k = k as u32; @@ -117,153 +62,180 @@ mod kernels { let my_rank = cluster::cluster_ctaidX(); let self_mask: u16 = 1u16 << (my_rank as u16); if tid == 0 { - mbarrier_init(&raw mut TMA_BAR0, 1); - mbarrier_init(&raw mut TMA_BAR1, 1); - mbarrier_init(&raw mut TMA_BAR2, 1); - mbarrier_init(&raw mut TMA_BAR3, 1); - mbarrier_init(&raw mut MMA_BAR0, 1); - mbarrier_init(&raw mut MMA_BAR1, 1); - mbarrier_init(&raw mut MMA_BAR2, 1); - mbarrier_init(&raw mut MMA_BAR3, 1); + let mut s = 0usize; + while s < 7 { + mbarrier_init((&raw mut TMA_BAR as *mut Barrier).add(s), 1); + mbarrier_init((&raw mut MMA_BAR as *mut Barrier).add(s), 1); + s += 1; + } mbarrier_init(&raw mut ACCUM_FULL0, 1); mbarrier_init(&raw mut ACCUM_FULL1, 1); - mbarrier_init(&raw mut ACCUM_EMPTY0, 256); - mbarrier_init(&raw mut ACCUM_EMPTY1, 256); - mbarrier_init(&raw mut TILE_READY, 1); + mbarrier_init(&raw mut ACCUM_EMPTY0, 8); + mbarrier_init(&raw mut ACCUM_EMPTY1, 8); + let mut ts = 0usize; + while ts < 4 { + mbarrier_init((&raw mut TILE_READY as *mut Barrier).add(ts), 1); + ts += 1; + } mbarrier_init(&raw mut CLC_BAR, 1); fence_proxy_async_shared_cta(); } - thread::sync_threads(); - - // Pre-arrive all MMA stage barriers so producer can start immediately. - if tid == 0 { - mbarrier_arrive(&raw const MMA_BAR0); - mbarrier_arrive(&raw const MMA_BAR1); - mbarrier_arrive(&raw const MMA_BAR2); - mbarrier_arrive(&raw const MMA_BAR3); - } - thread::sync_threads(); - if warp_id == 0 { tcgen05_alloc_cg2(&raw mut TMEM_ADDR as *mut u32, 512); } - thread::sync_threads(); + cluster::cluster_sync(); let tmem_addr = *(&raw const TMEM_ADDR as *const u32); let elect_one_cta = my_rank == 0; - let idesc = Tcgen05InstructionDescriptor::builder() + let idesc_256 = Tcgen05InstructionDescriptor::builder() .shape(Tcgen05MmaShape::M256_N256) .element_type(Tcgen05ElementType::F16) .accumulator_type(Tcgen05AccumulatorType::F32) .build() .raw(); + let idesc_192 = Tcgen05InstructionDescriptor::builder() + .shape(Tcgen05MmaShape { m: 256, n: 192 }) + .element_type(Tcgen05ElementType::F16) + .accumulator_type(Tcgen05AccumulatorType::F32) + .build() + .raw(); + let idesc_224 = Tcgen05InstructionDescriptor::builder() + .shape(Tcgen05MmaShape { m: 256, n: 224 }) + .element_type(Tcgen05ElementType::F16) + .accumulator_type(Tcgen05AccumulatorType::F32) + .build() + .raw(); let k_iters = k / 64; - cluster::cluster_sync(); if warp_id == TMA_WARP { - let is_lane0 = lane_id == 0; - let mut tile_seq: u32 = 0; - let mut clc_iter: u32 = 0; - - // The fixed host grid still contains one raw CLC work ID per - // 256x128 tile. The wide kernel needs half as many 256x256 tiles, - // so reinterpret the first half of the raw ID space and drain the - // remainder through CLC without publishing work to the consumers. - let cluster_base_id = thread::blockIdx_x() - my_rank; - let mut raw_tile_idx = cluster_base_id / 2; - let wide_tiles_n = _tiles_n / 2; - let wide_total = tiles_m * wide_tiles_n; - let resp_ptr = &raw mut CLC_RESPONSE as *mut u64; + if lane_id == 0 { + let mut stage: usize = 0; + let mut mma_parity: u32 = 1; + let mut tile_seq: usize = 0; + let mut clc_iter: u32 = 0; + + let cluster_base_id = thread::blockIdx_x() - my_rank; + let mut raw_tile_idx = cluster_base_id / 2; + let wide_tiles_n = _tiles_n / 2; + let wide_total = tiles_m * wide_tiles_n; + let num_waves = (wide_total + 73) / 74; + let work_total = num_waves * 74; + let num_convert = work_total - wide_total; + let use_n224 = tiles_m >= 32; + let full_tiles = if use_n224 { + wide_total - 7 * num_convert + } else { + wide_total - 3 * num_convert + }; + let head_cols = if use_n224 { + wide_tiles_n - 7 + } else if tiles_m == 16 { + wide_tiles_n - 9 + } else { + wide_tiles_n - 3 + }; + let head_tiles = tiles_m * head_cols; + let resp_ptr = &raw mut CLC_RESPONSE as *mut u64; + let aliased_tma_bar0 = (((&raw mut TMA_BAR as *mut Barrier) as usize) + & (PEER_BIT_MASK as usize)) + as *mut Barrier; + + loop { + let clc_parity = clc_iter & 1; + if raw_tile_idx < work_total { + let (tile_m, col_bf16_base, tile_kind) = if raw_tile_idx < head_tiles { + let group_tiles = swizzle_g * tiles_m; + let group = raw_tile_idx / group_tiles; + let in_group = raw_tile_idx % group_tiles; + let n_start = group * swizzle_g; + let band_w = if swizzle_g < head_cols - n_start { + swizzle_g + } else { + head_cols - n_start + }; + let raw_m = in_group / band_w; + let tm = if (group & 1) != 0 { + tiles_m - 1 - raw_m + } else { + raw_m + }; + let tn = n_start + in_group % band_w; + (tm, tn * 256, 0u32) + } else if use_n224 { + if raw_tile_idx < full_tiles { + let rem = raw_tile_idx - head_tiles; + let t = rem / 7; + let sub = rem % 7; + (t, (head_cols + sub) * 256, 0u32) + } else { + let tail = raw_tile_idx - full_tiles; + let t = (full_tiles - head_tiles) / 7 + (tail / 8); + let sub = tail & 7; + (t, head_cols * 256 + sub * 224, 2u32) + } + } else if raw_tile_idx < full_tiles { + let rem = raw_tile_idx - head_tiles; + let t = rem / 3; + let sub = rem % 3; + let band = t / tiles_m; + let tm = t % tiles_m; + let tn = head_cols + band * 3 + sub; + (tm, tn * 256, 0u32) + } else { + let tail = raw_tile_idx - full_tiles; + let t = (full_tiles - head_tiles) / 3 + (tail / 4); + let sub = tail & 3; + let band = t / tiles_m; + let tm = t % tiles_m; + let tn_base = head_cols + band * 3; + (tm, tn_base * 256 + sub * 192, 1u32) + }; - loop { - if raw_tile_idx < wide_total { - // L2 cache-blocking swizzle over logical 256-column tiles. - let group_tiles = swizzle_g * tiles_m; - let group = raw_tile_idx / group_tiles; - let in_group = raw_tile_idx % group_tiles; - let n_start = group * swizzle_g; - let band_w = if swizzle_g < wide_tiles_n - n_start { - swizzle_g - } else { - wide_tiles_n - n_start - }; - let tile_m = in_group / band_w; - let tile_n = n_start + in_group % band_w; - - if is_lane0 { - *(&raw mut TILE_INFO as *mut u32).add(0) = tile_m; - *(&raw mut TILE_INFO as *mut u32).add(1) = tile_n; - *(&raw mut TILE_INFO as *mut u32).add(2) = 1; - mbarrier_arrive(&raw const TILE_READY); - } + let packed_col = (col_bf16_base >> 5) | (tile_kind << 30); - let m_offset = (tile_m * 256 + my_rank * 128) as i32; - let b_n_offset = (tile_n * 256 + my_rank * 128) as i32; + let t_slot = tile_seq & 3; + let info_base = (&raw mut TILE_INFO as *mut u32).add(t_slot * 4); + *info_base.add(0) = tile_m; + *info_base.add(1) = packed_col; + *info_base.add(2) = 1; + mbarrier_arrive( + (&raw const TILE_READY as *const Barrier).add(t_slot), + ); + tile_seq += 1; - let mut k_idx: u32 = 0; - while k_idx < k_iters { - let global_k = tile_seq * k_iters + k_idx; - let stage = global_k & 3; - let mma_parity = (global_k >> 2) & 1; - - let ( - smem_a_ptr, - smem_b_ptr, - tma_bar_const, - tma_bar_mut, - mma_bar_const, - ): ( - *mut u8, - *mut u8, - *const Barrier, - *mut Barrier, - *const Barrier, - ) = match stage { - 0 => ( - &raw mut SMEM_A0 as *mut u8, - &raw mut SMEM_B0 as *mut u8, - &raw const TMA_BAR0 as *const Barrier, - &raw mut TMA_BAR0 as *mut Barrier, - &raw const MMA_BAR0 as *const Barrier, - ), - 1 => ( - &raw mut SMEM_A1 as *mut u8, - &raw mut SMEM_B1 as *mut u8, - &raw const TMA_BAR1 as *const Barrier, - &raw mut TMA_BAR1 as *mut Barrier, - &raw const MMA_BAR1 as *const Barrier, - ), - 2 => ( - &raw mut SMEM_A2 as *mut u8, - &raw mut SMEM_B2 as *mut u8, - &raw const TMA_BAR2 as *const Barrier, - &raw mut TMA_BAR2 as *mut Barrier, - &raw const MMA_BAR2 as *const Barrier, - ), - _ => ( - &raw mut SMEM_A3 as *mut u8, - &raw mut SMEM_B3 as *mut u8, - &raw const TMA_BAR3 as *const Barrier, - &raw mut TMA_BAR3 as *mut Barrier, - &raw const MMA_BAR3 as *const Barrier, - ), + let m_offset = (tile_m * 256 + my_rank * 128) as i32; + let b_n_offset = if tile_kind == 1 { + (col_bf16_base + my_rank * 96) as i32 + } else if tile_kind == 2 { + (col_bf16_base + my_rank * 112) as i32 + } else { + (col_bf16_base + my_rank * 128) as i32 + }; + let tx_bytes = if tile_kind == 1 { + (A_TILE_BYTES + B_TILE_BYTES_192) * 2 + } else { + (A_TILE_BYTES + B_TILE_BYTES_256) * 2 }; - // Both CTAs receive the MMA completion, so both producer - // warps can safely wait before reusing this SMEM stage. - while !mbarrier_try_wait_parity(mma_bar_const, mma_parity) {} + let mut k_idx: u32 = 0; + #[unroll(4)] + while k_idx < k_iters { + let smem_a_ptr = + (&raw mut SMEM_BUF as *mut u8).add(stage * 16384); + let smem_b_ptr = (&raw mut SMEM_BUF as *mut u8) + .add(b_smem_offset + stage * 16384); + let tma_bar_const = + (&raw const TMA_BAR as *const Barrier).add(stage); + let aliased_bar = aliased_tma_bar0.add(stage); + let mma_bar_const = + (&raw const MMA_BAR as *const Barrier).add(stage); + + while !mbarrier_try_wait_parity(mma_bar_const, mma_parity) {} - if is_lane0 { if elect_one_cta { - mbarrier_arrive_expect_tx( - tma_bar_const, - 1, - (A_TILE_BYTES + B_TILE_BYTES) * 2, - ); + mbarrier_arrive_expect_tx(tma_bar_const, 1, tx_bytes); } - let aliased_bar = - ((tma_bar_mut as u32) & PEER_BIT_MASK) as *mut Barrier; let k_base = (k_idx * 64) as i32; cp_async_bulk_tensor_2d_g2s_multicast_cg2( smem_a_ptr, @@ -273,8 +245,6 @@ mod kernels { aliased_bar, self_mask, ); - // The fixed host descriptor exposes a 64x64 B panel. - // Concatenate two panels to form this CTA's 128 columns. cp_async_bulk_tensor_2d_g2s_multicast_cg2( smem_b_ptr, b_tma, @@ -283,169 +253,209 @@ mod kernels { aliased_bar, self_mask, ); - cp_async_bulk_tensor_2d_g2s_multicast_cg2( - smem_b_ptr.add(B_PANEL_BYTES as usize), - b_tma, - k_base, - b_n_offset + 64, - aliased_bar, - self_mask, + if tile_kind == 1 { + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add(B_PANEL32_BYTES as usize), + b_tma, + k_base, + b_n_offset + 32, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 2) as usize), + b_tma, + k_base, + b_n_offset + 64, + aliased_bar, + self_mask, + ); + } else if tiles_m == 16 { + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add(B_PANEL32_BYTES as usize), + b_tma, + k_base, + b_n_offset + 32, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 2) as usize), + b_tma, + k_base, + b_n_offset + 64, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 3) as usize), + b_tma, + k_base, + b_n_offset + 96, + aliased_bar, + self_mask, + ); + } + if k_idx + 4 == k_iters { + mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); + if elect_one_cta { + clc_try_cancel_multicast( + resp_ptr as *mut u8, + &raw mut CLC_BAR, + ); + } + } + + stage += 1; + if stage == num_stages { + stage = 0; + mma_parity ^= 1; + } + k_idx += 1; + } + } else { + mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); + if elect_one_cta { + clc_try_cancel_multicast( + resp_ptr as *mut u8, + &raw mut CLC_BAR, ); } - - k_idx += 1; } - tile_seq += 1; - } - // One cancellation query consumes one CLC barrier phase even - // when its raw work ID is outside the reduced wide-tile range. - let clc_parity = clc_iter & 1; - if is_lane0 { - mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); - if elect_one_cta { - clc_try_cancel_multicast(resp_ptr as *mut u8, &raw mut CLC_BAR); - } - } - while !mbarrier_try_wait_parity(&raw const CLC_BAR, clc_parity) {} - clc_iter += 1; - - let resp_lo = *resp_ptr; - let resp_hi = *resp_ptr.add(1); - let is_canceled = clc_query_is_canceled(resp_lo, resp_hi); - if is_canceled == 0 { - if is_lane0 { - *(&raw mut TILE_INFO as *mut u32).add(2) = 0; - mbarrier_arrive(&raw const TILE_READY); + while !mbarrier_try_wait_parity(&raw const CLC_BAR, clc_parity) {} + clc_iter += 1; + + let resp_lo = *resp_ptr; + let resp_hi = *resp_ptr.add(1); + let is_canceled = clc_query_is_canceled(resp_lo, resp_hi); + if is_canceled == 0 { + let t_slot = tile_seq & 3; + let info_base = (&raw mut TILE_INFO as *mut u32).add(t_slot * 4); + *info_base.add(2) = 0; + mbarrier_arrive( + (&raw const TILE_READY as *const Barrier).add(t_slot), + ); + break; } - break; - } - let first_stolen = clc_query_get_first_ctaid_x(resp_lo, resp_hi); - raw_tile_idx = first_stolen / 2; + let first_stolen = clc_query_get_first_ctaid_x(resp_lo, resp_hi); + raw_tile_idx = first_stolen / 2; + } } } if warp_id == MMA_WARP { - let is_lane0 = lane_id == 0; - let mut tile_iter: u32 = 0; - let mut tile_parity: u32 = 0; - - loop { - while !mbarrier_try_wait_parity(&raw const TILE_READY, tile_parity) {} - tile_parity ^= 1; - - let has_work = *(&raw const TILE_INFO as *const u32).add(2); - if has_work == 0 { - break; - } + if elect_one_cta && lane_id == 0 { + let mut tile_iter: u32 = 0; + let mut stage: usize = 0; + let mut tma_parity: u32 = 0; + let smem_a_base0 = + cvta_generic_to_shared_offset(&raw const SMEM_BUF as *const u8); + let smem_b_base0 = smem_a_base0 + b_smem_offset as u64; + let a_desc0 = + build_smem_descriptor(smem_a_base0, LBO_BYTES, SBO_BYTES, SWIZZLE_128B); + let b_desc0 = + build_smem_descriptor(smem_b_base0, LBO_BYTES, SBO_BYTES, SWIZZLE_128B); + + loop { + let t_slot = (tile_iter & 3) as usize; + let t_parity = (tile_iter >> 2) & 1; + while !mbarrier_try_wait_parity( + (&raw const TILE_READY as *const Barrier).add(t_slot), + t_parity, + ) {} + + let info_base = (&raw const TILE_INFO as *const u32).add(t_slot * 4); + let has_work = *info_base.add(2); + if has_work == 0 { + break; + } + let packed_col = *info_base.add(1); + let tile_kind = packed_col >> 30; + let idesc = if tile_kind == 1 { + idesc_192 + } else if tile_kind == 2 { + idesc_224 + } else { + idesc_256 + }; - let accum_stage = tile_iter % NUM_ACCUM_STAGES; - let tmem_stage_offset = accum_stage * ACCUM_STAGE_COLS; + let accum_stage = tile_iter % NUM_ACCUM_STAGES; + let tmem_stage_offset = accum_stage * ACCUM_STAGE_COLS; - if elect_one_cta && tile_iter >= NUM_ACCUM_STAGES { - let empty_parity = ((tile_iter - NUM_ACCUM_STAGES) / NUM_ACCUM_STAGES) & 1; - if accum_stage == 0 { - while !mbarrier_try_wait_parity(&raw const ACCUM_EMPTY0, empty_parity) { - } - } else { - while !mbarrier_try_wait_parity(&raw const ACCUM_EMPTY1, empty_parity) { + if tile_iter >= NUM_ACCUM_STAGES { + let empty_parity = + ((tile_iter - NUM_ACCUM_STAGES) / NUM_ACCUM_STAGES) & 1; + if accum_stage == 0 { + while !mbarrier_try_wait_parity( + &raw const ACCUM_EMPTY0, + empty_parity, + ) {} + } else { + while !mbarrier_try_wait_parity( + &raw const ACCUM_EMPTY1, + empty_parity, + ) {} } } - } - let tile_k_base = tile_iter * k_iters; - let mut k_idx: u32 = 0; - // Unroll one full pipeline cycle. The launch contract guarantees - // k_iters % 4 == 0, so the producer's global stage and this local - // stage agree at every tile boundary. Keeping this expression - // loop-local lets the unroll pass fold each stage match. - #[unroll(4)] - while k_idx < k_iters { - let global_k = tile_k_base + k_idx; - let stage = k_idx & 3; - let tma_parity = (global_k >> 2) & 1; - - let (smem_a_base, smem_b_base, tma_bar_const, mma_bar_mut): ( - u64, - u64, - *const Barrier, - *mut Barrier, - ) = match stage { - 0 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A0 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B0 as *const u8), - &raw const TMA_BAR0 as *const Barrier, - &raw mut MMA_BAR0 as *mut Barrier, - ), - 1 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A1 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B1 as *const u8), - &raw const TMA_BAR1 as *const Barrier, - &raw mut MMA_BAR1 as *mut Barrier, - ), - 2 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A2 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B2 as *const u8), - &raw const TMA_BAR2 as *const Barrier, - &raw mut MMA_BAR2 as *mut Barrier, - ), - _ => ( - cvta_generic_to_shared_offset(&raw const SMEM_A3 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B3 as *const u8), - &raw const TMA_BAR3 as *const Barrier, - &raw mut MMA_BAR3 as *mut Barrier, - ), - }; + let mut k_idx: u32 = 0; + #[unroll(4)] + while k_idx < k_iters { + let stage_desc_offset = (stage as u64) << 10; + let a_desc_s = a_desc0 + stage_desc_offset; + let b_desc_s = b_desc0 + stage_desc_offset; + let tma_bar_const = + (&raw const TMA_BAR as *const Barrier).add(stage); + let mma_bar_mut = + (&raw mut MMA_BAR as *mut Barrier).add(stage); - // LEADER-ONLY TMA WAIT + MMA: - // Because TMA completions are aliased to the leader's barrier, only - // the leader can (and should) wait on tma_bar_const. The follower's - // TMA barrier is never signaled. The follower's MMA warp simply loops - // through the K iterations without doing work — pair-UMMA is issued - // by the leader and operates on both CTAs' SMEM simultaneously. - if elect_one_cta { while !mbarrier_try_wait_parity(tma_bar_const, tma_parity) {} - if is_lane0 { - let mut j: u32 = 0; - while j < 4 { - let byte_offset = (j * 32) as u64; - let a_desc = build_smem_descriptor( - smem_a_base + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); - let b_desc = build_smem_descriptor( - smem_b_base + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); + let tmem_dst = tmem_addr + tmem_stage_offset; + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s, + b_desc_s, + idesc, + k_idx > 0, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 2, + b_desc_s + 2, + idesc, + true, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 4, + b_desc_s + 4, + idesc, + true, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 6, + b_desc_s + 6, + idesc, + true, + ); - let accumulate = k_idx > 0 || j > 0; - tcgen05_mma_f16_cg2( - tmem_addr + tmem_stage_offset, - a_desc, - b_desc, - idesc, - accumulate, - ); - j += 1; - } + tcgen05_commit_multicast_cg2( + mma_bar_mut as *mut u64, + CTA_MASK_PAIR, + ); - tcgen05_commit_multicast_cg2( - mma_bar_mut as *mut u64, - CTA_MASK_PAIR, - ); + stage += 1; + if stage == num_stages { + stage = 0; + tma_parity ^= 1; } + k_idx += 1; } - k_idx += 1; - } - - if elect_one_cta && is_lane0 { if accum_stage == 0 { tcgen05_commit_multicast_cg2( &raw mut ACCUM_FULL0 as *mut u64, @@ -457,9 +467,9 @@ mod kernels { CTA_MASK_PAIR, ); } - } - tile_iter += 1; + tile_iter += 1; + } } tcgen05_relinquish_alloc_permit_cg2(); @@ -467,31 +477,35 @@ mod kernels { if warp_id < 4 { let mut epi_tile_iter: u32 = 0; - let mut tile_parity: u32 = 0; let leader_accum_empty0_addr = cvta_generic_to_shared_offset(cluster::map_shared_rank(&raw const ACCUM_EMPTY0, 0) as *const u8); let leader_accum_empty1_addr = cvta_generic_to_shared_offset(cluster::map_shared_rank(&raw const ACCUM_EMPTY1, 0) as *const u8); - const TILE_N: usize = 256; - let warp_row_base = (warp_id * 32) as usize; - let row_stride_bytes = TILE_N * 2; - let row_within_8 = (lane_id % 8) as usize; - let is_second_matrix = (8..16).contains(&lane_id); - let col_offset_for_matrix2 = if is_second_matrix { 16usize } else { 0usize }; + let is_odd_lane = (lane_id & 1) != 0; + let is_bit1_lane = (lane_id & 2) != 0; + let lane_row_in_16 = ((lane_id / 4) + (lane_id & 1) * 8) as usize; + let lane_col_u32 = ((lane_id & 2) * 2) as usize; loop { - while !mbarrier_try_wait_parity(&raw const TILE_READY, tile_parity) {} - tile_parity ^= 1; - - let has_work = *(&raw const TILE_INFO as *const u32).add(2); + let t_slot = (epi_tile_iter & 3) as usize; + let t_parity = (epi_tile_iter >> 2) & 1; + while !mbarrier_try_wait_parity( + (&raw const TILE_READY as *const Barrier).add(t_slot), + t_parity, + ) {} + + let info_base = (&raw const TILE_INFO as *const u32).add(t_slot * 4); + let has_work = *info_base.add(2); if has_work == 0 { break; } - let tile_m = *(&raw const TILE_INFO as *const u32).add(0); - let tile_n = *(&raw const TILE_INFO as *const u32).add(1); + let tile_m = *info_base.add(0); + let packed_col = *info_base.add(1); + let tile_kind = packed_col >> 30; + let col_u32_start = ((packed_col & 0x3FFF_FFFF) * 16) as usize; let accum_stage = epi_tile_iter % NUM_ACCUM_STAGES; let tmem_stage_offset = accum_stage * ACCUM_STAGE_COLS; @@ -503,160 +517,102 @@ mod kernels { while !mbarrier_try_wait_parity(&raw const ACCUM_FULL1, full_parity) {} } + let n_u32 = (n / 2) as usize; + let tile_row_base = (tile_m * 256 + my_rank * 128 + warp_id * 32) as usize; + let tile_col_base = col_u32_start + lane_col_u32; + let col_chunks: u32 = if tile_kind == 1 { 3 } else { 4 }; + let mut tmem_row_block = 0u32; while tmem_row_block < 2 { let tmem_row = warp_id * 32 + tmem_row_block * 16; + let global_row = + tile_row_base + (tmem_row_block as usize * 16) + lane_row_in_16; + let row_out_base = global_row * n_u32 + tile_col_base; - let mut col_block = 0u32; - while col_block < 16 { - let col_offset = (col_block * 16) as usize; + let mut col_chunk = 0u32; + #[unroll(never)] + while col_chunk < col_chunks { + let col_offset = col_chunk * 64; - let regs_a = tcgen05_ld_16x256b_pure( - tmem_addr - + tmem_stage_offset - + (tmem_row << 16) - + col_offset as u32, + let regs = tcgen05_ld_16x256b_x8_pure( + tmem_addr + tmem_stage_offset + (tmem_row << 16) + col_offset, ); tcgen05_load_wait(); - let regs_b = tcgen05_ld_16x256b_pure( - tmem_addr - + tmem_stage_offset - + (tmem_row << 16) - + col_offset as u32 - + 8, - ); - tcgen05_load_wait(); + if tmem_row_block == 1 && col_chunk + 1 == col_chunks && lane_id == 0 { + if elect_one_cta { + if accum_stage == 0 { + mbarrier_arrive(&raw const ACCUM_EMPTY0); + } else { + mbarrier_arrive(&raw const ACCUM_EMPTY1); + } + } else if accum_stage == 0 { + mbarrier_arrive_cluster(leader_accum_empty0_addr); + } else { + mbarrier_arrive_cluster(leader_accum_empty1_addr); + } + } - let p0_lo = cvt_bf16x2_f32(regs_a[0], regs_a[1]); - let p1_lo = cvt_bf16x2_f32(regs_b[0], regs_b[1]); - let out_row_lo = - warp_row_base + (tmem_row_block as usize * 16) + row_within_8; - let smem_addr_lo = (&raw mut SMEM_OUT as *mut u8).add( - out_row_lo * row_stride_bytes - + col_offset * 2 - + col_offset_for_matrix2, - ); - stmatrix_m8n8_x2(smem_addr_lo, p0_lo, p1_lo); - - let p0_hi = cvt_bf16x2_f32(regs_a[2], regs_a[3]); - let p1_hi = cvt_bf16x2_f32(regs_b[2], regs_b[3]); - let out_row_hi = - warp_row_base + (tmem_row_block as usize * 16) + 8 + row_within_8; - let smem_addr_hi = (&raw mut SMEM_OUT as *mut u8).add( - out_row_hi * row_stride_bytes - + col_offset * 2 - + col_offset_for_matrix2, - ); - stmatrix_m8n8_x2(smem_addr_hi, p0_hi, p1_hi); + let col_u32_base = (col_chunk as usize) * 32; + let max_pairs: usize = + if tile_kind == 2 && col_chunk == 3 { 2 } else { 4 }; + let mut pair = 0usize; + #[unroll] + while pair < 4 { + if pair < max_pairs { + let base = pair * 8; + let p0_lo = cvt_bf16x2_f32(regs[base], regs[base + 1]); + let p0_hi = cvt_bf16x2_f32(regs[base + 2], regs[base + 3]); + let send0 = if is_odd_lane { p0_lo } else { p0_hi }; + let recv0 = warp::shuffle_xor_sync(u32::MAX, send0, 1); + let lo0 = if is_odd_lane { recv0 } else { p0_lo }; + let hi0 = if is_odd_lane { p0_hi } else { recv0 }; + + let p1_lo = cvt_bf16x2_f32(regs[base + 4], regs[base + 5]); + let p1_hi = cvt_bf16x2_f32(regs[base + 6], regs[base + 7]); + let send1 = if is_odd_lane { p1_lo } else { p1_hi }; + let recv1 = warp::shuffle_xor_sync(u32::MAX, send1, 1); + let lo1 = if is_odd_lane { recv1 } else { p1_lo }; + let hi1 = if is_odd_lane { p1_hi } else { recv1 }; + + let send_a = if is_bit1_lane { lo0 } else { lo1 }; + let send_b = if is_bit1_lane { hi0 } else { hi1 }; + let recv_a = warp::shuffle_xor_sync(u32::MAX, send_a, 2); + let recv_b = warp::shuffle_xor_sync(u32::MAX, send_b, 2); + + let w0 = if is_bit1_lane { recv_a } else { lo0 }; + let w1 = if is_bit1_lane { recv_b } else { hi0 }; + let w2 = if is_bit1_lane { lo1 } else { recv_a }; + let w3 = if is_bit1_lane { hi1 } else { recv_b }; + + #[repr(C, align(16))] + struct U32x4(u32, u32, u32, u32); + let ptr = out.get_unchecked_mut( + row_out_base + col_u32_base + pair * 8, + ) as *mut u32 + as *mut U32x4; + *ptr = U32x4(w0, w1, w2, w3); + } + pair += 1; + } - col_block += 1; + col_chunk += 1; } tmem_row_block += 1; } - let n_u32 = (n / 2) as usize; - let tile_row_base = (tile_m * 256 + my_rank * 128) as usize; - let tile_col_base = (tile_n * 128) as usize; - let base_row = warp_id as usize * 32; - - let mut elem = lane_id as usize * 2; - while elem < 4096 { - let local_row = elem / 128; - let local_col = elem % 128; - let smem_idx = (base_row + local_row) * 128 + local_col; - let global_row = tile_row_base + base_row + local_row; - let global_col = tile_col_base + local_col; - let global_idx = global_row * n_u32 + global_col; - let packed = (SMEM_OUT[smem_idx] as u64) - | ((SMEM_OUT[smem_idx + 1] as u64) << 32); - let global_ptr = - out.get_unchecked_mut(global_idx) as *mut u32 as *mut u64; - - *global_ptr = packed; - elem += 64; - } - - if elect_one_cta { - if accum_stage == 0 { - mbarrier_arrive(&raw const ACCUM_EMPTY0); - } else { - mbarrier_arrive(&raw const ACCUM_EMPTY1); - } - } else { - if accum_stage == 0 { - mbarrier_arrive_cluster(leader_accum_empty0_addr); - } else { - mbarrier_arrive_cluster(leader_accum_empty1_addr); - } - } - epi_tile_iter += 1; } } - // BUG FIX: cluster_sync before exit prevents "Cluster target block not present" - // (CUDA_EXCEPTION_17). Without this, a fast CTA can exit while its partner is still - // executing cross-CTA operations (e.g., mbarrier_arrive_cluster in the epilogue). cluster::cluster_sync(); if warp_id == 0 { tcgen05_dealloc_cg2(tmem_addr, 512); } - if tid == 0 { - mbarrier_inval(&raw mut TMA_BAR0); - mbarrier_inval(&raw mut TMA_BAR1); - mbarrier_inval(&raw mut TMA_BAR2); - mbarrier_inval(&raw mut TMA_BAR3); - mbarrier_inval(&raw mut MMA_BAR0); - mbarrier_inval(&raw mut MMA_BAR1); - mbarrier_inval(&raw mut MMA_BAR2); - mbarrier_inval(&raw mut MMA_BAR3); - mbarrier_inval(&raw mut ACCUM_FULL0); - mbarrier_inval(&raw mut ACCUM_FULL1); - mbarrier_inval(&raw mut ACCUM_EMPTY0); - mbarrier_inval(&raw mut ACCUM_EMPTY1); - mbarrier_inval(&raw mut TILE_READY); - mbarrier_inval(&raw mut CLC_BAR); } - } } - /// Target kernel: CLC + TMA multicast with 4 SMEM stages and no MCAST_BAR. - /// - /// Blackwell GEMM kernel: one cta_group::2 worker computes an M512xN256 macro tile. - /// - /// Architecture: 6 warps per CTA, cluster size 2 (CTA pairs). - /// - Warp 4 (TMA): loads two A panels and one shared B tile per K stage. - /// - Warp 5 (MMA): computes two adjacent M256xN256 outputs that reuse B. - /// - Warps 0-3 (Epilogue): stream the two fixed TMEM output halves to global memory. - /// - /// Pipeline stages: 4 SMEM stages (TMA_BAR0..3 / MMA_BAR0..3), two fixed TMEM - /// output halves (FULL0/1, EMPTY0/1), plus TILE_READY for macro-tile handoff. - /// - /// Key synchronization protocol (cta_group::2 barrier aliasing): - /// Both CTAs in a pair issue TMA loads, but the barrier pointer is masked with - /// PEER_BIT_MASK (0xFEFFFFF8) before being passed to the TMA instruction. This - /// clears bit 24, redirecting both CTAs' completion signals to the leader CTA's - /// (rank 0) barrier. Consequently: - /// - Only the leader sets expect_tx (doubled: both CTAs' bytes land on one barrier). - /// - Only the leader waits on TMA barriers in the MMA warp. - /// - The follower's MMA warp skips TMA waits since pair-UMMA is leader-issued. - /// MMA barriers still use normal multicast (tcgen05_commit_multicast_cg2 with - /// CTA_MASK_PAIR=0b11), so both CTAs receive MMA completion signals. - /// - /// CLC work-stealing: rank 0 issues clc_try_cancel_multicast; both CTAs receive the - /// response via CLC_BAR. Tile indices are derived by dividing the CLC first_ctaid_x - /// by the cluster size (2), NOT using raw CTA IDs. - /// - /// # Shape contract - /// - /// `n` must be a multiple of 256 and `tiles_m` must be even. The fixed host grid - /// remains expressed in M256xN128 work IDs; only its first quarter is reinterpreted - /// as M512xN256 macro tiles and the remaining IDs are drained through CLC. - /// - /// `k` must be a multiple of 256. With a K tile size of 64, this makes `k_iters` - /// a multiple of four, so every macro tile starts on pipeline stage 0. #[kernel] #[cluster_launch(2, 1, 1)] pub unsafe fn gemm_sol_clc_multicast_4_stage_pipeline_large( @@ -668,58 +624,41 @@ mod kernels { tiles_m: u32, _tiles_n: u32, ) { - unsafe { - static mut SMEM_A0: SharedArray = SharedArray::UNINIT; - static mut SMEM_A1: SharedArray = SharedArray::UNINIT; - static mut SMEM_A2: SharedArray = SharedArray::UNINIT; - static mut SMEM_A3: SharedArray = SharedArray::UNINIT; - static mut SMEM_B0: SharedArray = SharedArray::UNINIT; - static mut SMEM_B1: SharedArray = SharedArray::UNINIT; - static mut SMEM_B2: SharedArray = SharedArray::UNINIT; - static mut SMEM_B3: SharedArray = SharedArray::UNINIT; - // Rolling 32 KiB epilogue scratch: each warp owns 32x128 BF16 values. - static mut SMEM_OUT: SharedArray = SharedArray::UNINIT; + unsafe { + static mut SMEM_BUF: SharedArray = SharedArray::UNINIT; static mut TMEM_ADDR: SharedArray = SharedArray::UNINIT; - static mut TILE_INFO: SharedArray = SharedArray::UNINIT; - - // 4-stage TMA <-> MMA pipeline - static mut TMA_BAR0: Barrier = Barrier::UNINIT; - static mut TMA_BAR1: Barrier = Barrier::UNINIT; - static mut TMA_BAR2: Barrier = Barrier::UNINIT; - static mut TMA_BAR3: Barrier = Barrier::UNINIT; - static mut MMA_BAR0: Barrier = Barrier::UNINIT; - static mut MMA_BAR1: Barrier = Barrier::UNINIT; - static mut MMA_BAR2: Barrier = Barrier::UNINIT; - static mut MMA_BAR3: Barrier = Barrier::UNINIT; + static mut TILE_INFO: SharedArray = SharedArray::UNINIT; + + // Up to 8-stage TMA <-> MMA pipeline + static mut TMA_BAR: SharedArray = SharedArray::UNINIT; + static mut MMA_BAR: SharedArray = SharedArray::UNINIT; static mut ACCUM_FULL0: Barrier = Barrier::UNINIT; static mut ACCUM_FULL1: Barrier = Barrier::UNINIT; static mut ACCUM_EMPTY0: Barrier = Barrier::UNINIT; static mut ACCUM_EMPTY1: Barrier = Barrier::UNINIT; - static mut TILE_READY: Barrier = Barrier::UNINIT; + static mut TILE_READY: SharedArray = SharedArray::UNINIT; static mut CLC_RESPONSE: SharedArray = SharedArray::UNINIT; static mut CLC_BAR: Barrier = Barrier::UNINIT; const A_TILE_BYTES: u32 = 128 * 64 * 2; - const B_PANEL_BYTES: u32 = 64 * 64 * 2; - const B_TILE_BYTES: u32 = 128 * 64 * 2; // 128 B rows per CTA (split by rank) + const B_PANEL32_BYTES: u32 = 32 * 64 * 2; + const B_TILE_BYTES_256: u32 = 128 * 64 * 2; + const B_TILE_BYTES_192: u32 = 96 * 64 * 2; const SBO_BYTES: u32 = 1024; const LBO_BYTES: u32 = 16; const SWIZZLE_128B: u8 = 2; const TMA_WARP: u32 = 4; const MMA_WARP: u32 = 5; - const ACCUM_HALF_COLS: u32 = 256; + const NUM_ACCUM_STAGES: u32 = 2; + const ACCUM_STAGE_COLS: u32 = 256; const CTA_MASK_PAIR: u16 = 0b11; - // Clears bit 24 (CTA rank within pair) + alignment bits 2:0 of a shared - // memory barrier address, redirecting TMA completions to the leader CTA's - // barrier. const PEER_BIT_MASK: u32 = 0xFEFFFFF8; - // L2 cache-blocking: the 256-column output tile doubles the physical - // width of each N tile, so preserve the accepted cache-band widths by - // halving the prior logical group sizes. - let swizzle_g: u32 = if tiles_m <= 16 { 2 } else { 8 }; + let num_stages: usize = 7; + let b_smem_offset: usize = 7 * 16384; + let swizzle_g: u32 = if tiles_m == 16 { 2 } else { 8 }; let n = n as u32; let k = k as u32; @@ -731,177 +670,189 @@ mod kernels { let my_rank = cluster::cluster_ctaidX(); let self_mask: u16 = 1u16 << (my_rank as u16); if tid == 0 { - mbarrier_init(&raw mut TMA_BAR0, 1); - mbarrier_init(&raw mut TMA_BAR1, 1); - mbarrier_init(&raw mut TMA_BAR2, 1); - mbarrier_init(&raw mut TMA_BAR3, 1); - mbarrier_init(&raw mut MMA_BAR0, 2); - mbarrier_init(&raw mut MMA_BAR1, 2); - mbarrier_init(&raw mut MMA_BAR2, 2); - mbarrier_init(&raw mut MMA_BAR3, 2); + let mut s = 0usize; + while s < 7 { + mbarrier_init((&raw mut TMA_BAR as *mut Barrier).add(s), 1); + mbarrier_init((&raw mut MMA_BAR as *mut Barrier).add(s), 1); + s += 1; + } mbarrier_init(&raw mut ACCUM_FULL0, 1); mbarrier_init(&raw mut ACCUM_FULL1, 1); - mbarrier_init(&raw mut ACCUM_EMPTY0, 256); - mbarrier_init(&raw mut ACCUM_EMPTY1, 256); - mbarrier_init(&raw mut TILE_READY, 1); + mbarrier_init(&raw mut ACCUM_EMPTY0, 8); + mbarrier_init(&raw mut ACCUM_EMPTY1, 8); + let mut ts = 0usize; + while ts < 4 { + mbarrier_init((&raw mut TILE_READY as *mut Barrier).add(ts), 1); + ts += 1; + } mbarrier_init(&raw mut CLC_BAR, 1); fence_proxy_async_shared_cta(); } - thread::sync_threads(); - - // Pre-arrive all MMA stage barriers so producer can start immediately. - if tid == 0 { - mbarrier_arrive(&raw const MMA_BAR0); - mbarrier_arrive(&raw const MMA_BAR1); - mbarrier_arrive(&raw const MMA_BAR2); - mbarrier_arrive(&raw const MMA_BAR3); - mbarrier_arrive(&raw const MMA_BAR0); - mbarrier_arrive(&raw const MMA_BAR1); - mbarrier_arrive(&raw const MMA_BAR2); - mbarrier_arrive(&raw const MMA_BAR3); - } - thread::sync_threads(); - if warp_id == 0 { tcgen05_alloc_cg2(&raw mut TMEM_ADDR as *mut u32, 512); } - thread::sync_threads(); + cluster::cluster_sync(); let tmem_addr = *(&raw const TMEM_ADDR as *const u32); let elect_one_cta = my_rank == 0; - let idesc = Tcgen05InstructionDescriptor::builder() + let idesc_256 = Tcgen05InstructionDescriptor::builder() .shape(Tcgen05MmaShape::M256_N256) .element_type(Tcgen05ElementType::F16) .accumulator_type(Tcgen05AccumulatorType::F32) .build() .raw(); + let idesc_192 = Tcgen05InstructionDescriptor::builder() + .shape(Tcgen05MmaShape { m: 256, n: 192 }) + .element_type(Tcgen05ElementType::F16) + .accumulator_type(Tcgen05AccumulatorType::F32) + .build() + .raw(); + let idesc_224 = Tcgen05InstructionDescriptor::builder() + .shape(Tcgen05MmaShape { m: 256, n: 224 }) + .element_type(Tcgen05ElementType::F16) + .accumulator_type(Tcgen05AccumulatorType::F32) + .build() + .raw(); let k_iters = k / 64; - cluster::cluster_sync(); if warp_id == TMA_WARP { - let is_lane0 = lane_id == 0; - let mut macro_seq: u32 = 0; - let mut clc_iter: u32 = 0; - - // The fixed host grid contains one raw CLC work ID per M256xN128 - // tile. Reinterpret its first quarter as M512xN256 macro tiles and - // drain the remaining raw IDs without publishing consumer work. - let cluster_base_id = thread::blockIdx_x() - my_rank; - let mut raw_tile_idx = cluster_base_id / 2; - let macro_tiles_m = tiles_m / 2; - let wide_tiles_n = _tiles_n / 2; - let macro_total = macro_tiles_m * wide_tiles_n; - let resp_ptr = &raw mut CLC_RESPONSE as *mut u64; + if lane_id == 0 { + let mut stage: usize = 0; + let mut mma_parity: u32 = 1; + let mut tile_seq: usize = 0; + let mut clc_iter: u32 = 0; + + let cluster_base_id = thread::blockIdx_x() - my_rank; + let mut raw_tile_idx = cluster_base_id / 2; + let wide_tiles_n = _tiles_n / 2; + let wide_total = tiles_m * wide_tiles_n; + let num_waves = (wide_total + 73) / 74; + let work_total = num_waves * 74; + let num_convert = work_total - wide_total; + let use_n224 = tiles_m >= 32; + let full_tiles = if use_n224 { + wide_total - 7 * num_convert + } else { + wide_total - 3 * num_convert + }; + let head_cols = if use_n224 { + wide_tiles_n - 7 + } else if tiles_m == 16 { + wide_tiles_n - 9 + } else { + wide_tiles_n - 3 + }; + let head_tiles = tiles_m * head_cols; + let resp_ptr = &raw mut CLC_RESPONSE as *mut u64; + let aliased_tma_bar0 = (((&raw mut TMA_BAR as *mut Barrier) as usize) + & (PEER_BIT_MASK as usize)) + as *mut Barrier; + + loop { + let clc_parity = clc_iter & 1; + if raw_tile_idx < work_total { + let (tile_m, col_bf16_base, tile_kind) = if raw_tile_idx < head_tiles { + let group_tiles = swizzle_g * tiles_m; + let group = raw_tile_idx / group_tiles; + let in_group = raw_tile_idx % group_tiles; + let n_start = group * swizzle_g; + let band_w = if swizzle_g < head_cols - n_start { + swizzle_g + } else { + head_cols - n_start + }; + let raw_m = in_group / band_w; + let tm = if (group & 1) != 0 { + tiles_m - 1 - raw_m + } else { + raw_m + }; + let tn = n_start + in_group % band_w; + (tm, tn * 256, 0u32) + } else if use_n224 { + if raw_tile_idx < full_tiles { + let rem = raw_tile_idx - head_tiles; + let t = rem / 7; + let sub = rem % 7; + (t, (head_cols + sub) * 256, 0u32) + } else { + let tail = raw_tile_idx - full_tiles; + let t = (full_tiles - head_tiles) / 7 + (tail / 8); + let sub = tail & 7; + (t, head_cols * 256 + sub * 224, 2u32) + } + } else if raw_tile_idx < full_tiles { + let rem = raw_tile_idx - head_tiles; + let t = rem / 3; + let sub = rem % 3; + let band = t / tiles_m; + let tm = t % tiles_m; + let tn = head_cols + band * 3 + sub; + (tm, tn * 256, 0u32) + } else { + let tail = raw_tile_idx - full_tiles; + let t = (full_tiles - head_tiles) / 3 + (tail / 4); + let sub = tail & 3; + let band = t / tiles_m; + let tm = t % tiles_m; + let tn_base = head_cols + band * 3; + (tm, tn_base * 256 + sub * 192, 1u32) + }; - loop { - if raw_tile_idx < macro_total { - // L2 cache-blocking swizzle over M512xN256 macro tiles. - let group_tiles = swizzle_g * macro_tiles_m; - let group = raw_tile_idx / group_tiles; - let in_group = raw_tile_idx % group_tiles; - let n_start = group * swizzle_g; - let band_w = if swizzle_g < wide_tiles_n - n_start { - swizzle_g - } else { - wide_tiles_n - n_start - }; - let macro_m = in_group / band_w; - let tile_n = n_start + in_group % band_w; - - if is_lane0 { - *(&raw mut TILE_INFO as *mut u32).add(0) = macro_m; - *(&raw mut TILE_INFO as *mut u32).add(1) = tile_n; - *(&raw mut TILE_INFO as *mut u32).add(2) = 1; - mbarrier_arrive(&raw const TILE_READY); - } + let packed_col = (col_bf16_base >> 5) | (tile_kind << 30); - let m_offset0 = (macro_m * 512 + my_rank * 128) as i32; - let m_offset1 = m_offset0 + 256; - let b_n_offset = (tile_n * 256 + my_rank * 128) as i32; + let t_slot = tile_seq & 3; + let info_base = (&raw mut TILE_INFO as *mut u32).add(t_slot * 4); + *info_base.add(0) = tile_m; + *info_base.add(1) = packed_col; + *info_base.add(2) = 1; + mbarrier_arrive( + (&raw const TILE_READY as *const Barrier).add(t_slot), + ); + tile_seq += 1; - let mut k_idx: u32 = 0; - while k_idx < k_iters { - let global_k = macro_seq * k_iters + k_idx; - let stage = global_k & 3; - let mma_parity = (global_k >> 2) & 1; - - let ( - smem_a_ptr, - smem_b_ptr, - tma_bar_const, - tma_bar_mut, - mma_bar_const, - ): ( - *mut u8, - *mut u8, - *const Barrier, - *mut Barrier, - *const Barrier, - ) = match stage { - 0 => ( - &raw mut SMEM_A0 as *mut u8, - &raw mut SMEM_B0 as *mut u8, - &raw const TMA_BAR0 as *const Barrier, - &raw mut TMA_BAR0 as *mut Barrier, - &raw const MMA_BAR0 as *const Barrier, - ), - 1 => ( - &raw mut SMEM_A1 as *mut u8, - &raw mut SMEM_B1 as *mut u8, - &raw const TMA_BAR1 as *const Barrier, - &raw mut TMA_BAR1 as *mut Barrier, - &raw const MMA_BAR1 as *const Barrier, - ), - 2 => ( - &raw mut SMEM_A2 as *mut u8, - &raw mut SMEM_B2 as *mut u8, - &raw const TMA_BAR2 as *const Barrier, - &raw mut TMA_BAR2 as *mut Barrier, - &raw const MMA_BAR2 as *const Barrier, - ), - _ => ( - &raw mut SMEM_A3 as *mut u8, - &raw mut SMEM_B3 as *mut u8, - &raw const TMA_BAR3 as *const Barrier, - &raw mut TMA_BAR3 as *mut Barrier, - &raw const MMA_BAR3 as *const Barrier, - ), + let m_offset = (tile_m * 256 + my_rank * 128) as i32; + let b_n_offset = if tile_kind == 1 { + (col_bf16_base + my_rank * 96) as i32 + } else if tile_kind == 2 { + (col_bf16_base + my_rank * 112) as i32 + } else { + (col_bf16_base + my_rank * 128) as i32 + }; + let tx_bytes = if tile_kind == 1 { + (A_TILE_BYTES + B_TILE_BYTES_192) * 2 + } else { + (A_TILE_BYTES + B_TILE_BYTES_256) * 2 }; - // Both CTAs receive the MMA completion, so both producer - // warps can safely wait before reusing this SMEM stage. - while !mbarrier_try_wait_parity(mma_bar_const, mma_parity) {} + let mut k_idx: u32 = 0; + #[unroll(4)] + while k_idx < k_iters { + let smem_a_ptr = + (&raw mut SMEM_BUF as *mut u8).add(stage * 16384); + let smem_b_ptr = (&raw mut SMEM_BUF as *mut u8) + .add(b_smem_offset + stage * 16384); + let tma_bar_const = + (&raw const TMA_BAR as *const Barrier).add(stage); + let aliased_bar = aliased_tma_bar0.add(stage); + let mma_bar_const = + (&raw const MMA_BAR as *const Barrier).add(stage); + + while !mbarrier_try_wait_parity(mma_bar_const, mma_parity) {} - if is_lane0 { if elect_one_cta { - mbarrier_arrive_expect_tx( - tma_bar_const, - 1, - (2 * A_TILE_BYTES + B_TILE_BYTES) * 2, - ); + mbarrier_arrive_expect_tx(tma_bar_const, 1, tx_bytes); } - let aliased_bar = - ((tma_bar_mut as u32) & PEER_BIT_MASK) as *mut Barrier; let k_base = (k_idx * 64) as i32; cp_async_bulk_tensor_2d_g2s_multicast_cg2( smem_a_ptr, a_tma, k_base, - m_offset0, - aliased_bar, - self_mask, - ); - cp_async_bulk_tensor_2d_g2s_multicast_cg2( - smem_a_ptr.add(A_TILE_BYTES as usize), - a_tma, - k_base, - m_offset1, + m_offset, aliased_bar, self_mask, ); - // The fixed host descriptor exposes a 64x64 B panel. - // Concatenate two panels to form this CTA's 128 columns. cp_async_bulk_tensor_2d_g2s_multicast_cg2( smem_b_ptr, b_tma, @@ -910,395 +861,363 @@ mod kernels { aliased_bar, self_mask, ); - cp_async_bulk_tensor_2d_g2s_multicast_cg2( - smem_b_ptr.add(B_PANEL_BYTES as usize), - b_tma, - k_base, - b_n_offset + 64, - aliased_bar, - self_mask, + if tile_kind == 1 { + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add(B_PANEL32_BYTES as usize), + b_tma, + k_base, + b_n_offset + 32, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 2) as usize), + b_tma, + k_base, + b_n_offset + 64, + aliased_bar, + self_mask, + ); + } else if tiles_m == 16 { + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add(B_PANEL32_BYTES as usize), + b_tma, + k_base, + b_n_offset + 32, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 2) as usize), + b_tma, + k_base, + b_n_offset + 64, + aliased_bar, + self_mask, + ); + cp_async_bulk_tensor_2d_g2s_multicast_cg2( + smem_b_ptr.add((B_PANEL32_BYTES * 3) as usize), + b_tma, + k_base, + b_n_offset + 96, + aliased_bar, + self_mask, + ); + } + if k_idx + 4 == k_iters { + mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); + if elect_one_cta { + clc_try_cancel_multicast( + resp_ptr as *mut u8, + &raw mut CLC_BAR, + ); + } + } + + stage += 1; + if stage == num_stages { + stage = 0; + mma_parity ^= 1; + } + k_idx += 1; + } + } else { + mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); + if elect_one_cta { + clc_try_cancel_multicast( + resp_ptr as *mut u8, + &raw mut CLC_BAR, ); } - - k_idx += 1; } - macro_seq += 1; - } - // One cancellation query consumes one CLC barrier phase even - // when its raw work ID is outside the reduced wide-tile range. - let clc_parity = clc_iter & 1; - if is_lane0 { - mbarrier_arrive_expect_tx(&raw const CLC_BAR, 1, 16); - if elect_one_cta { - clc_try_cancel_multicast(resp_ptr as *mut u8, &raw mut CLC_BAR); - } - } - while !mbarrier_try_wait_parity(&raw const CLC_BAR, clc_parity) {} - clc_iter += 1; - - let resp_lo = *resp_ptr; - let resp_hi = *resp_ptr.add(1); - let is_canceled = clc_query_is_canceled(resp_lo, resp_hi); - if is_canceled == 0 { - if is_lane0 { - *(&raw mut TILE_INFO as *mut u32).add(2) = 0; - mbarrier_arrive(&raw const TILE_READY); + while !mbarrier_try_wait_parity(&raw const CLC_BAR, clc_parity) {} + clc_iter += 1; + + let resp_lo = *resp_ptr; + let resp_hi = *resp_ptr.add(1); + let is_canceled = clc_query_is_canceled(resp_lo, resp_hi); + if is_canceled == 0 { + let t_slot = tile_seq & 3; + let info_base = (&raw mut TILE_INFO as *mut u32).add(t_slot * 4); + *info_base.add(2) = 0; + mbarrier_arrive( + (&raw const TILE_READY as *const Barrier).add(t_slot), + ); + break; } - break; - } - let first_stolen = clc_query_get_first_ctaid_x(resp_lo, resp_hi); - raw_tile_idx = first_stolen / 2; + let first_stolen = clc_query_get_first_ctaid_x(resp_lo, resp_hi); + raw_tile_idx = first_stolen / 2; + } } } if warp_id == MMA_WARP { - let is_lane0 = lane_id == 0; - let mut macro_iter: u32 = 0; - let mut tile_parity: u32 = 0; - - loop { - while !mbarrier_try_wait_parity(&raw const TILE_READY, tile_parity) {} - tile_parity ^= 1; - - let has_work = *(&raw const TILE_INFO as *const u32).add(2); - if has_work == 0 { - break; - } - - if elect_one_cta && macro_iter > 0 { - let empty_parity = (macro_iter - 1) & 1; - while !mbarrier_try_wait_parity(&raw const ACCUM_EMPTY0, empty_parity) {} - } - - let macro_k_base = macro_iter * k_iters; - let mut k_idx: u32 = 0; - // Four K64 stages form one complete SMEM ring cycle. Each stage - // commits its half-0 and half-1 MMA groups separately. - #[unroll(4)] - while k_idx < k_iters { - let global_k = macro_k_base + k_idx; - let stage = k_idx & 3; - let tma_parity = (global_k >> 2) & 1; - - let (smem_a_base, smem_b_base, tma_bar_const, mma_bar_mut): ( - u64, - u64, - *const Barrier, - *mut Barrier, - ) = match stage { - 0 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A0 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B0 as *const u8), - &raw const TMA_BAR0 as *const Barrier, - &raw mut MMA_BAR0 as *mut Barrier, - ), - 1 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A1 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B1 as *const u8), - &raw const TMA_BAR1 as *const Barrier, - &raw mut MMA_BAR1 as *mut Barrier, - ), - 2 => ( - cvta_generic_to_shared_offset(&raw const SMEM_A2 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B2 as *const u8), - &raw const TMA_BAR2 as *const Barrier, - &raw mut MMA_BAR2 as *mut Barrier, - ), - _ => ( - cvta_generic_to_shared_offset(&raw const SMEM_A3 as *const u8), - cvta_generic_to_shared_offset(&raw const SMEM_B3 as *const u8), - &raw const TMA_BAR3 as *const Barrier, - &raw mut MMA_BAR3 as *mut Barrier, - ), + if elect_one_cta && lane_id == 0 { + let mut tile_iter: u32 = 0; + let mut stage: usize = 0; + let mut tma_parity: u32 = 0; + let smem_a_base0 = + cvta_generic_to_shared_offset(&raw const SMEM_BUF as *const u8); + let smem_b_base0 = smem_a_base0 + b_smem_offset as u64; + let a_desc0 = + build_smem_descriptor(smem_a_base0, LBO_BYTES, SBO_BYTES, SWIZZLE_128B); + let b_desc0 = + build_smem_descriptor(smem_b_base0, LBO_BYTES, SBO_BYTES, SWIZZLE_128B); + + loop { + let t_slot = (tile_iter & 3) as usize; + let t_parity = (tile_iter >> 2) & 1; + while !mbarrier_try_wait_parity( + (&raw const TILE_READY as *const Barrier).add(t_slot), + t_parity, + ) {} + + let info_base = (&raw const TILE_INFO as *const u32).add(t_slot * 4); + let has_work = *info_base.add(2); + if has_work == 0 { + break; + } + let packed_col = *info_base.add(1); + let tile_kind = packed_col >> 30; + let idesc = if tile_kind == 1 { + idesc_192 + } else if tile_kind == 2 { + idesc_224 + } else { + idesc_256 }; - // TMA completions are aliased to the leader CTA. The leader - // issues both pair-UMMA outputs while the follower advances - // the same K-stage loop without touching the barriers. - if elect_one_cta { - while !mbarrier_try_wait_parity(tma_bar_const, tma_parity) {} - - if is_lane0 { - let mut j: u32 = 0; - while j < 4 { - let byte_offset = (j * 32) as u64; - let a_desc = build_smem_descriptor( - smem_a_base + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); - let b_desc = build_smem_descriptor( - smem_b_base + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); - let accumulate = k_idx > 0 || j > 0; - tcgen05_mma_f16_cg2( - tmem_addr, - a_desc, - b_desc, - idesc, - accumulate, - ); - j += 1; - } - tcgen05_commit_multicast_cg2( - mma_bar_mut as *mut u64, - CTA_MASK_PAIR, - ); - if k_idx + 1 == k_iters { - tcgen05_commit_multicast_cg2( - &raw mut ACCUM_FULL0 as *mut u64, - CTA_MASK_PAIR, - ); - } - } + let accum_stage = tile_iter % NUM_ACCUM_STAGES; + let tmem_stage_offset = accum_stage * ACCUM_STAGE_COLS; - // Once half 0 of the new macro has started, wait until the - // prior macro's half 1 has drained before overwriting it. - if macro_iter > 0 && k_idx == 0 { - let empty_parity = (macro_iter - 1) & 1; + if tile_iter >= NUM_ACCUM_STAGES { + let empty_parity = + ((tile_iter - NUM_ACCUM_STAGES) / NUM_ACCUM_STAGES) & 1; + if accum_stage == 0 { + while !mbarrier_try_wait_parity( + &raw const ACCUM_EMPTY0, + empty_parity, + ) {} + } else { while !mbarrier_try_wait_parity( &raw const ACCUM_EMPTY1, empty_parity, ) {} } + } - if is_lane0 { - let mut j: u32 = 0; - while j < 4 { - let byte_offset = (j * 32) as u64; - let a_desc = build_smem_descriptor( - smem_a_base + A_TILE_BYTES as u64 + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); - let b_desc = build_smem_descriptor( - smem_b_base + byte_offset, - LBO_BYTES, - SBO_BYTES, - SWIZZLE_128B, - ); - let accumulate = k_idx > 0 || j > 0; - tcgen05_mma_f16_cg2( - tmem_addr + ACCUM_HALF_COLS, - a_desc, - b_desc, - idesc, - accumulate, - ); - j += 1; - } - tcgen05_commit_multicast_cg2( - mma_bar_mut as *mut u64, - CTA_MASK_PAIR, - ); - if k_idx + 1 == k_iters { - tcgen05_commit_multicast_cg2( - &raw mut ACCUM_FULL1 as *mut u64, - CTA_MASK_PAIR, - ); - } + let mut k_idx: u32 = 0; + #[unroll(4)] + while k_idx < k_iters { + let stage_desc_offset = (stage as u64) << 10; + let a_desc_s = a_desc0 + stage_desc_offset; + let b_desc_s = b_desc0 + stage_desc_offset; + let tma_bar_const = + (&raw const TMA_BAR as *const Barrier).add(stage); + let mma_bar_mut = + (&raw mut MMA_BAR as *mut Barrier).add(stage); + + while !mbarrier_try_wait_parity(tma_bar_const, tma_parity) {} + + let tmem_dst = tmem_addr + tmem_stage_offset; + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s, + b_desc_s, + idesc, + k_idx > 0, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 2, + b_desc_s + 2, + idesc, + true, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 4, + b_desc_s + 4, + idesc, + true, + ); + tcgen05_mma_f16_cg2( + tmem_dst, + a_desc_s + 6, + b_desc_s + 6, + idesc, + true, + ); + + tcgen05_commit_multicast_cg2( + mma_bar_mut as *mut u64, + CTA_MASK_PAIR, + ); + + stage += 1; + if stage == num_stages { + stage = 0; + tma_parity ^= 1; } + k_idx += 1; } - k_idx += 1; - } + if accum_stage == 0 { + tcgen05_commit_multicast_cg2( + &raw mut ACCUM_FULL0 as *mut u64, + CTA_MASK_PAIR, + ); + } else { + tcgen05_commit_multicast_cg2( + &raw mut ACCUM_FULL1 as *mut u64, + CTA_MASK_PAIR, + ); + } - macro_iter += 1; + tile_iter += 1; + } } tcgen05_relinquish_alloc_permit_cg2(); } if warp_id < 4 { - let mut epi_macro_iter: u32 = 0; - let mut tile_parity: u32 = 0; + let mut epi_tile_iter: u32 = 0; let leader_accum_empty0_addr = cvta_generic_to_shared_offset(cluster::map_shared_rank(&raw const ACCUM_EMPTY0, 0) as *const u8); let leader_accum_empty1_addr = cvta_generic_to_shared_offset(cluster::map_shared_rank(&raw const ACCUM_EMPTY1, 0) as *const u8); - const SCRATCH_BF16_COLS: usize = 128; - const SCRATCH_U32_COLS: usize = SCRATCH_BF16_COLS / 2; - const SCRATCH_U32_PER_WARP: usize = 32 * SCRATCH_U32_COLS; - const SCRATCH_ROW_BYTES: usize = SCRATCH_BF16_COLS * 2; - let warp_scratch_u32 = warp_id as usize * SCRATCH_U32_PER_WARP; - let warp_scratch_u8 = - (&raw mut SMEM_OUT as *mut u8).add(warp_scratch_u32 * 4); - let row_within_8 = (lane_id % 8) as usize; - let is_second_matrix = (8..16).contains(&lane_id); - let col_offset_for_matrix2 = if is_second_matrix { 16usize } else { 0usize }; + let is_odd_lane = (lane_id & 1) != 0; + let is_bit1_lane = (lane_id & 2) != 0; + let lane_row_in_16 = ((lane_id / 4) + (lane_id & 1) * 8) as usize; + let lane_col_u32 = ((lane_id & 2) * 2) as usize; loop { - while !mbarrier_try_wait_parity(&raw const TILE_READY, tile_parity) {} - tile_parity ^= 1; - - let has_work = *(&raw const TILE_INFO as *const u32).add(2); + let t_slot = (epi_tile_iter & 3) as usize; + let t_parity = (epi_tile_iter >> 2) & 1; + while !mbarrier_try_wait_parity( + (&raw const TILE_READY as *const Barrier).add(t_slot), + t_parity, + ) {} + + let info_base = (&raw const TILE_INFO as *const u32).add(t_slot * 4); + let has_work = *info_base.add(2); if has_work == 0 { break; } - let macro_m = *(&raw const TILE_INFO as *const u32).add(0); - let tile_n = *(&raw const TILE_INFO as *const u32).add(1); - let full_parity = epi_macro_iter & 1; - let n_u32 = (n / 2) as usize; - let tile_col_base = (tile_n * 128) as usize; - - let mut macro_half = 0u32; - while macro_half < 2 { - if macro_half == 0 { - while !mbarrier_try_wait_parity( - &raw const ACCUM_FULL0, - full_parity, - ) {} - } else { - while !mbarrier_try_wait_parity( - &raw const ACCUM_FULL1, - full_parity, - ) {} - } + let tile_m = *info_base.add(0); + let packed_col = *info_base.add(1); + let tile_kind = packed_col >> 30; + let col_u32_start = ((packed_col & 0x3FFF_FFFF) * 16) as usize; - let tmem_half_offset = macro_half * ACCUM_HALF_COLS; - let macro_row_base = - (macro_m * 512 + macro_half * 256 + my_rank * 128) as usize; + let accum_stage = epi_tile_iter % NUM_ACCUM_STAGES; + let tmem_stage_offset = accum_stage * ACCUM_STAGE_COLS; - let mut n_chunk = 0u32; - while n_chunk < 2 { - let tmem_chunk_offset = tmem_half_offset + n_chunk * 128; + let full_parity = (epi_tile_iter / NUM_ACCUM_STAGES) & 1; + if accum_stage == 0 { + while !mbarrier_try_wait_parity(&raw const ACCUM_FULL0, full_parity) {} + } else { + while !mbarrier_try_wait_parity(&raw const ACCUM_FULL1, full_parity) {} + } - let mut tmem_row_block = 0u32; - while tmem_row_block < 2 { - let tmem_row = warp_id * 32 + tmem_row_block * 16; - let scratch_row_base = tmem_row_block as usize * 16; + let n_u32 = (n / 2) as usize; + let tile_row_base = (tile_m * 256 + my_rank * 128 + warp_id * 32) as usize; + let tile_col_base = col_u32_start + lane_col_u32; + let col_chunks: u32 = if tile_kind == 1 { 3 } else { 4 }; - let mut col_block = 0u32; - while col_block < 8 { - let col_offset = (col_block * 16) as usize; + let mut tmem_row_block = 0u32; + while tmem_row_block < 2 { + let tmem_row = warp_id * 32 + tmem_row_block * 16; + let global_row = + tile_row_base + (tmem_row_block as usize * 16) + lane_row_in_16; + let row_out_base = global_row * n_u32 + tile_col_base; - let regs_a = tcgen05_ld_16x256b_pure( - tmem_addr - + (tmem_row << 16) - + tmem_chunk_offset - + col_offset as u32, - ); - tcgen05_load_wait(); - - let regs_b = tcgen05_ld_16x256b_pure( - tmem_addr - + (tmem_row << 16) - + tmem_chunk_offset - + col_offset as u32 - + 8, - ); - tcgen05_load_wait(); - - let p0_lo = cvt_bf16x2_f32(regs_a[0], regs_a[1]); - let p1_lo = cvt_bf16x2_f32(regs_b[0], regs_b[1]); - let scratch_row_lo = scratch_row_base + row_within_8; - let smem_addr_lo = warp_scratch_u8.add( - scratch_row_lo * SCRATCH_ROW_BYTES - + col_offset * 2 - + col_offset_for_matrix2, - ); - stmatrix_m8n8_x2(smem_addr_lo, p0_lo, p1_lo); - - let p0_hi = cvt_bf16x2_f32(regs_a[2], regs_a[3]); - let p1_hi = cvt_bf16x2_f32(regs_b[2], regs_b[3]); - let scratch_row_hi = - scratch_row_base + 8 + row_within_8; - let smem_addr_hi = warp_scratch_u8.add( - scratch_row_hi * SCRATCH_ROW_BYTES - + col_offset * 2 - + col_offset_for_matrix2, - ); - stmatrix_m8n8_x2(smem_addr_hi, p0_hi, p1_hi); + let mut col_chunk = 0u32; + #[unroll(never)] + while col_chunk < col_chunks { + let col_offset = col_chunk * 64; - col_block += 1; - } + let regs = tcgen05_ld_16x256b_x8_pure( + tmem_addr + tmem_stage_offset + (tmem_row << 16) + col_offset, + ); + tcgen05_load_wait(); - tmem_row_block += 1; + if tmem_row_block == 1 && col_chunk + 1 == col_chunks && lane_id == 0 { + if elect_one_cta { + if accum_stage == 0 { + mbarrier_arrive(&raw const ACCUM_EMPTY0); + } else { + mbarrier_arrive(&raw const ACCUM_EMPTY1); + } + } else if accum_stage == 0 { + mbarrier_arrive_cluster(leader_accum_empty0_addr); + } else { + mbarrier_arrive_cluster(leader_accum_empty1_addr); + } } - warp::sync_mask(u32::MAX); - - let global_row_base = - macro_row_base + warp_id as usize * 32; - let global_col_base = - tile_col_base + n_chunk as usize * SCRATCH_U32_COLS; - - let mut elem = lane_id as usize * 2; - while elem < SCRATCH_U32_PER_WARP { - let local_row = elem / SCRATCH_U32_COLS; - let local_col = elem % SCRATCH_U32_COLS; - let smem_idx = warp_scratch_u32 - + local_row * SCRATCH_U32_COLS - + local_col; - let global_idx = (global_row_base + local_row) * n_u32 - + global_col_base - + local_col; - let packed = (SMEM_OUT[smem_idx] as u64) - | ((SMEM_OUT[smem_idx + 1] as u64) << 32); - let global_ptr = - out.get_unchecked_mut(global_idx) as *mut u32 as *mut u64; - - *global_ptr = packed; - elem += 64; + let col_u32_base = (col_chunk as usize) * 32; + let max_pairs: usize = + if tile_kind == 2 && col_chunk == 3 { 2 } else { 4 }; + let mut pair = 0usize; + #[unroll] + while pair < 4 { + if pair < max_pairs { + let base = pair * 8; + let p0_lo = cvt_bf16x2_f32(regs[base], regs[base + 1]); + let p0_hi = cvt_bf16x2_f32(regs[base + 2], regs[base + 3]); + let send0 = if is_odd_lane { p0_lo } else { p0_hi }; + let recv0 = warp::shuffle_xor_sync(u32::MAX, send0, 1); + let lo0 = if is_odd_lane { recv0 } else { p0_lo }; + let hi0 = if is_odd_lane { p0_hi } else { recv0 }; + + let p1_lo = cvt_bf16x2_f32(regs[base + 4], regs[base + 5]); + let p1_hi = cvt_bf16x2_f32(regs[base + 6], regs[base + 7]); + let send1 = if is_odd_lane { p1_lo } else { p1_hi }; + let recv1 = warp::shuffle_xor_sync(u32::MAX, send1, 1); + let lo1 = if is_odd_lane { recv1 } else { p1_lo }; + let hi1 = if is_odd_lane { p1_hi } else { recv1 }; + + let send_a = if is_bit1_lane { lo0 } else { lo1 }; + let send_b = if is_bit1_lane { hi0 } else { hi1 }; + let recv_a = warp::shuffle_xor_sync(u32::MAX, send_a, 2); + let recv_b = warp::shuffle_xor_sync(u32::MAX, send_b, 2); + + let w0 = if is_bit1_lane { recv_a } else { lo0 }; + let w1 = if is_bit1_lane { recv_b } else { hi0 }; + let w2 = if is_bit1_lane { lo1 } else { recv_a }; + let w3 = if is_bit1_lane { hi1 } else { recv_b }; + + #[repr(C, align(16))] + struct U32x4(u32, u32, u32, u32); + let ptr = out.get_unchecked_mut( + row_out_base + col_u32_base + pair * 8, + ) as *mut u32 + as *mut U32x4; + *ptr = U32x4(w0, w1, w2, w3); + } + pair += 1; } - warp::sync_mask(u32::MAX); - n_chunk += 1; + col_chunk += 1; } - - if elect_one_cta { - if macro_half == 0 { - mbarrier_arrive(&raw const ACCUM_EMPTY0); - } else { - mbarrier_arrive(&raw const ACCUM_EMPTY1); - } - } else if macro_half == 0 { - mbarrier_arrive_cluster(leader_accum_empty0_addr); - } else { - mbarrier_arrive_cluster(leader_accum_empty1_addr); - } - - macro_half += 1; + tmem_row_block += 1; } - epi_macro_iter += 1; + epi_tile_iter += 1; } } - // BUG FIX: cluster_sync before exit prevents "Cluster target block not present" - // (CUDA_EXCEPTION_17). Without this, a fast CTA can exit while its partner is still - // executing cross-CTA operations (e.g., mbarrier_arrive_cluster in the epilogue). cluster::cluster_sync(); if warp_id == 0 { tcgen05_dealloc_cg2(tmem_addr, 512); } - if tid == 0 { - mbarrier_inval(&raw mut TMA_BAR0); - mbarrier_inval(&raw mut TMA_BAR1); - mbarrier_inval(&raw mut TMA_BAR2); - mbarrier_inval(&raw mut TMA_BAR3); - mbarrier_inval(&raw mut MMA_BAR0); - mbarrier_inval(&raw mut MMA_BAR1); - mbarrier_inval(&raw mut MMA_BAR2); - mbarrier_inval(&raw mut MMA_BAR3); - mbarrier_inval(&raw mut ACCUM_FULL0); - mbarrier_inval(&raw mut ACCUM_FULL1); - mbarrier_inval(&raw mut ACCUM_EMPTY0); - mbarrier_inval(&raw mut ACCUM_EMPTY1); - mbarrier_inval(&raw mut TILE_READY); - mbarrier_inval(&raw mut CLC_BAR); } - } } } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/main.rs index 4f1e2fd470..1734484d07 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/main.rs @@ -7,7 +7,9 @@ clippy::not_unsafe_ptr_arg_deref, clippy::needless_range_loop, clippy::unnecessary_cast, - clippy::too_many_arguments + clippy::too_many_arguments, + clippy::collapsible_if, + clippy::manual_div_ceil )] //! Canonical cuda-oxide Blackwell GEMM speed-of-light example. @@ -22,11 +24,11 @@ //! - B: N×K f16, row-major (transposed storage, K contiguous) //! - C: M×N bf16 output, row-major (packed as u32 pairs) -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::barrier::{ Barrier, fence_proxy_async_shared_cta, mbarrier_arrive, mbarrier_arrive_cluster, - mbarrier_arrive_expect_tx, mbarrier_init, mbarrier_inval, mbarrier_try_wait_parity, + mbarrier_arrive_expect_tx, mbarrier_init, mbarrier_try_wait_parity, }; use cuda_device::clc::{ clc_query_get_first_ctaid_x, clc_query_is_canceled, clc_try_cancel_multicast, @@ -36,8 +38,8 @@ use cuda_device::convert::{bf16_to_f32, cvt_bf16x2_f32}; use cuda_device::shared::{SharedArray, cvta_generic_to_shared_offset}; use cuda_device::tcgen05::{ Tcgen05AccumulatorType, Tcgen05ElementType, Tcgen05InstructionDescriptor, Tcgen05MmaShape, - stmatrix_m8n8_x2, tcgen05_alloc_cg2, tcgen05_commit_multicast_cg2, tcgen05_dealloc_cg2, - tcgen05_ld_16x256b_pure, tcgen05_load_wait, tcgen05_mma_f16_cg2, + tcgen05_alloc_cg2, tcgen05_commit_multicast_cg2, tcgen05_dealloc_cg2, + tcgen05_ld_16x256b_x8_pure, tcgen05_load_wait, tcgen05_mma_f16_cg2, tcgen05_relinquish_alloc_permit_cg2, }; use cuda_device::tma::{TmaDescriptor, cp_async_bulk_tensor_2d_g2s_multicast_cg2}; @@ -88,6 +90,9 @@ mod cublas_baseline { } fn load() -> Option> { + if std::env::var("GEMM_CUBLAS_SKIP").as_deref() == Ok("1") { + return None; + } let bin = bench_binary(); if !bin.exists() { eprintln!( @@ -101,8 +106,20 @@ mod cublas_baseline { return None; } - eprintln!("ℹ️ Measuring live cuBLASLt FP16 reference on host GPU..."); - let out = match std::process::Command::new(&bin).output() { + eprintln!("ℹ️ Measuring live cuBLASLt tuned FP16 reference on host GPU..."); + let mut command = std::process::Command::new(&bin); + // Forward any timing overrides so the live baseline and Rust batches + // use the same protocol. Defaults are identical in the two programs. + for (name, flag) in [ + ("GEMM_SOL_WARMUP", "--warmup"), + ("GEMM_SOL_ITERS", "--iters"), + ("GEMM_SOL_REPEATS", "--repeats"), + ] { + if let Ok(value) = std::env::var(name) { + command.args([flag, value.as_str()]); + } + } + let out = match command.output() { Ok(o) => o, Err(e) => { eprintln!("⚠️ Failed to run {}: {e}", bin.display()); @@ -182,7 +199,7 @@ fn print_cublas_comparison(tflops: f64, m: usize) { Some(sol) => { let pct = (tflops / sol) * 100.0; println!( - " vs cuBLAS: {:.2}% of live cublasLt FP16 reference ({:.0} TFLOPS)", + " vs cuBLAS: {:.2}% of live cublasLt tuned FP16 reference ({:.0} TFLOPS)", pct, sol ); } @@ -193,35 +210,27 @@ fn print_cublas_comparison(tflops: f64, m: usize) { } fn print_benchmark_summary(measurements: &[(usize, f64)]) { - let count = measurements.len() as f64; - let kernel_geomean = (measurements - .iter() - .map(|(_, value)| value.ln()) - .sum::() - / count) - .exp(); - - println!("── Fixed-size geomean summary ───────────────────────"); - println!(" Kernel: {:.3} TFLOPS", kernel_geomean); - - let reference_values: Option> = measurements - .iter() - .map(|(m, _)| cublas_baseline::reference_tflops(*m)) - .collect(); - if let Some(reference_values) = reference_values { - let reference_geomean = - (reference_values.iter().map(|value| value.ln()).sum::() / count).exp(); - println!(" cuBLAS ref: {:.3} TFLOPS", reference_geomean); - println!( - " Ratio: {:.2}% of live FP16 reference", - kernel_geomean / reference_geomean * 100.0 - ); - } else { - println!(" cuBLAS ref: unavailable"); + println!("── Individual fixed-size results ────────────────────"); + for (m, tflops) in measurements { + println!(" {m}³: {tflops:.3} TFLOPS"); } println!("─────────────────────────────────────────────────────"); } +fn benchmark_count(name: &str, default: usize) -> Result> { + match std::env::var(name) { + Ok(value) => { + let count = value.parse::()?; + if count == 0 || count > 10_000_000 { + return Err(format!("{name} must be between 1 and 10000000").into()); + } + Ok(count) + } + Err(std::env::VarError::NotPresent) => Ok(default), + Err(error) => Err(error.into()), + } +} + // ============================================================================= // KERNEL // ============================================================================= @@ -757,13 +766,12 @@ fn run_correctness_test_clc_multicast_4_stage_pipeline( let a_tma = create_tma_descriptor_f16_swizzled(a_ptr as *mut std::ffi::c_void, k as u64, m as u64)?; - // B TMA: 64 rows per CTA (each CTA loads half the N tile, split by rank) let b_tma = create_tma_descriptor_f16_swizzled_box( b_ptr as *mut std::ffi::c_void, k as u64, n as u64, 64, - 64, + if m == 4096 { 32 } else { 128 }, )?; let dev_a_tma = DeviceBuffer::from_host(stream, &a_tma.opaque)?; @@ -771,7 +779,8 @@ fn run_correctness_test_clc_multicast_4_stage_pipeline( let tiles_m = (m / 256) as u32; let tiles_n = (n / 128) as u32; - let total_tiles = tiles_m * tiles_n; + let wide_total = tiles_m * (tiles_n / 2); + let total_tiles = wide_total.div_ceil(74) * 74; let cfg = LaunchConfig { grid_dim: (total_tiles * 2, 1, 1), @@ -905,8 +914,12 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( n: usize, k: usize, ) -> Result> { - const WARMUP: usize = 10; - const ITERS: usize = 100; + let warmup = benchmark_count("GEMM_SOL_WARMUP", 2000)?; + let iters = benchmark_count("GEMM_SOL_ITERS", 301)?; + let repeats = benchmark_count("GEMM_SOL_REPEATS", 5)?; + if repeats % 2 == 0 { + return Err("GEMM_SOL_REPEATS must be odd for the median".into()); + } let problem = GemmProblem::new(m, n, k); let selected = select_gemm_variant(family, &problem, selection_mode)?; @@ -928,7 +941,7 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( k as u64, n as u64, 64, - 64, + if m == 4096 { 32 } else { 128 }, )?; let dev_a_tma = DeviceBuffer::from_host(stream, &a_tma.opaque)?; @@ -942,7 +955,8 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( let tiles_m = (m / 256) as u32; let tiles_n = (n / 128) as u32; - let total_tiles = tiles_m * tiles_n; + let wide_total = tiles_m * (tiles_n / 2); + let total_tiles = wide_total.div_ceil(74) * 74; let cfg = LaunchConfig { grid_dim: (total_tiles * 2, 1, 1), @@ -953,7 +967,7 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( let output_u32_count = m * n / 2; let mut dev_output = DeviceBuffer::::zeroed(stream, output_u32_count)?; - for _ in 0..WARMUP { + for _ in 0..warmup { unsafe { launch_kernel( module, @@ -971,31 +985,34 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( } stream.synchronize()?; - let start_event = - stream.record_event(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?; - - for _ in 0..ITERS { - unsafe { - launch_kernel( - module, - stream.as_ref(), - cfg, - a_tma_ptr, - b_tma_ptr, - &mut dev_output, - n_arg, - k_arg, - tiles_m, - tiles_n, - ) - }?; + let mut samples_us = Vec::with_capacity(repeats); + for _ in 0..repeats { + let start_event = + stream.record_event(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?; + for _ in 0..iters { + unsafe { + launch_kernel( + module, + stream.as_ref(), + cfg, + a_tma_ptr, + b_tma_ptr, + &mut dev_output, + n_arg, + k_arg, + tiles_m, + tiles_n, + ) + }?; + } + let end_event = + stream.record_event(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?; + let elapsed_ms = start_event.elapsed_ms(&end_event)?; + samples_us.push(elapsed_ms as f64 * 1000.0 / iters as f64); } - - let end_event = - stream.record_event(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?; - let elapsed_ms = start_event.elapsed_ms(&end_event)?; - - let avg_ms = elapsed_ms as f64 / ITERS as f64; + let mut ordered = samples_us.clone(); + ordered.sort_by(f64::total_cmp); + let avg_ms = ordered[repeats / 2] / 1000.0; let avg_us = avg_ms * 1000.0; let flops = 2.0 * m as f64 * n as f64 * k as f64; let tflops = (flops / (avg_ms / 1000.0)) / 1e12; @@ -1014,11 +1031,18 @@ fn run_benchmark_clc_multicast_4_stage_pipeline( ); println!(" K-loop: {} outer iters (BK=64, 4 MMAs each)", k / 64); println!(" Pipeline: CLC + cta_group::2 + 4-stage SMEM + unroll(4)"); - println!(" Iterations: {} (after {} warmup)", ITERS, WARMUP); - println!(" Total time: {:.3} ms", elapsed_ms); + println!( + " Iterations: {} × {} batches (after {} warmup)", + iters, repeats, warmup + ); + println!(" Statistic: median batch mean"); + println!(" Samples us: {:?}", samples_us); println!(" Average: {:.3} us / kernel", avg_us); println!(" FLOPS/kern: {:.3e}", flops); println!(" Throughput: {:.3} TFLOPS", tflops); + println!( + "RESULT {{\"implementation\":\"cuda-oxide\",\"n\":{m},\"warmup\":{warmup},\"iters\":{iters},\"repeats\":{repeats},\"samples_us\":{samples_us:?},\"median_us\":{avg_us:.9},\"tflops\":{tflops:.9},\"variant\":\"{output_tile}\",\"input\":\"fp16\",\"output\":\"bf16\",\"compute\":\"fp32\"}}" + ); print_cublas_comparison(tflops, m); println!("═══════════════════════════════════════════════════════\n"); @@ -1057,7 +1081,7 @@ fn create_tma_descriptor_f16_swizzled( CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16, CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, - CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE, + CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, cuTensorMapEncodeTiled, cudaError_enum_CUDA_SUCCESS, }; @@ -1080,7 +1104,7 @@ fn create_tma_descriptor_f16_swizzled( element_strides.as_ptr(), CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE, + CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, ) }; @@ -1103,7 +1127,7 @@ fn create_tma_descriptor_f16_swizzled_box( CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16, CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, - CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE, + CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, cuTensorMapEncodeTiled, cudaError_enum_CUDA_SUCCESS, }; @@ -1126,7 +1150,7 @@ fn create_tma_descriptor_f16_swizzled_box( element_strides.as_ptr(), CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE, + CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, ) }; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.lock index 9a3324ea0b..021a9ba09e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -265,6 +294,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.lock index cd9f2ba0b2..c66cccf280 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -269,6 +298,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/src/main.rs index 8c659d964d..f772e29912 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/src/main.rs @@ -8,7 +8,7 @@ //! The kernel calls generated coordinate, lane-mask, vote, and shuffle intrinsics //! directly. This covers the raw path instead of only `cuda-device` wrappers. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel}; use cuda_device::{ diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.lock index 04de2c5f8b..86beb77dfa 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -269,6 +298,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/src/main.rs index 0d29f842d9..25dd4dcc60 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/src/main.rs @@ -17,7 +17,7 @@ //! //! cargo oxide run generated_intrinsics_blackwell -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ DisjointSlice, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.lock index bb356b611b..bd580cdfda 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +220,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -269,6 +298,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/src/main.rs index 9bd7a2f03f..3267e9e6a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/src/main.rs @@ -18,7 +18,7 @@ //! Build and run with: //! cargo oxide run generated_ldmatrix -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::shared::cvta_generic_to_shared_u32; use cuda_device::wmma::{ diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock index 6a78ca4f12..9382a82e6d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -264,6 +293,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/src/main.rs index 013cdd4baa..543133fead 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/src/main.rs @@ -22,7 +22,7 @@ //! //! We should see a PTX entry point for `scale` (or `scale_f32` if we add type info). -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; use std::ops::{Add, Mul}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock index 86b24f7d2a..621806ac5c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/src/main.rs index e4984f32f9..cdaff2722c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/src/main.rs @@ -27,7 +27,7 @@ use std::sync::Arc; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU64}; use cuda_device::{DisjointSlice, kernel, thread}; @@ -273,7 +273,7 @@ impl GpuHashMap { let slots = DeviceBuffer::::zeroed(stream, capacity)?; unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( slots.cu_deviceptr(), 0xFF, slots.num_bytes(), @@ -285,7 +285,7 @@ impl GpuHashMap { // through a shared `&[DeviceAtomicU64]` (issue #151). `DeviceAtomicU64` // is `#[repr(transparent)]` over `u64`, so the pointer, length, and the // `0xFF`-filled bytes are all unchanged. - let slots = slots.cast_elem::(); + let slots = slots.wrap_elem::(); Ok(Self { slots, capacity }) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.lock index 0fb7f0e99e..e49e81af39 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.lock @@ -61,6 +61,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -159,7 +179,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -171,18 +191,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -193,7 +215,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -202,6 +226,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -222,6 +247,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -286,6 +315,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/bin/bench.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/bin/bench.rs index ad66f6359a..77e5fe3c2c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/bin/bench.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/bin/bench.rs @@ -36,7 +36,7 @@ use std::sync::Arc; use std::time::Instant; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use hashbrown::HashMap as HbMap; use hashmap_v2::*; @@ -96,13 +96,13 @@ unsafe fn reset_table_async( stream: &Arc, ) -> Result<(), Box> { unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( map.ctrl.cu_deviceptr(), 0xFF, map.ctrl.num_bytes(), stream.cu_stream(), )?; - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( map.slots.cu_deviceptr(), 0xFF, map.slots.num_bytes(), diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/lib.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/lib.rs index 89e171695e..678581d375 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/lib.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/src/lib.rs @@ -61,7 +61,7 @@ use std::sync::Arc; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaStream, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32, DeviceAtomicU64}; use cuda_device::cooperative_groups::{ThreadGroup, WarpCollective, this_thread_block}; @@ -1142,13 +1142,13 @@ impl GpuSwissMap { let ctrl = DeviceBuffer::::zeroed(stream, capacity / GROUP)?; let slots = DeviceBuffer::::zeroed(stream, capacity)?; unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( ctrl.cu_deviceptr(), 0xFF, ctrl.num_bytes(), stream.cu_stream(), )?; - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( slots.cu_deviceptr(), 0xFF, slots.num_bytes(), @@ -1162,8 +1162,8 @@ impl GpuSwissMap { // `&[u32]` / `&[u64]` would be readonly+noalias and writing // through it is UB (issue #151). The memset above is byte-level so // it stays valid across the cast. - let ctrl = ctrl.cast_elem::(); - let slots = slots.cast_elem::(); + let ctrl = ctrl.wrap_elem::(); + let slots = slots.wrap_elem::(); Ok(Self { ctrl, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.lock index b781c994c5..9afd413785 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.lock @@ -61,6 +61,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -159,7 +179,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -171,18 +191,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -193,7 +215,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -202,6 +226,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -222,6 +247,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -286,6 +315,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/bin/bench.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/bin/bench.rs index 7a8e5945de..2c65bd8b58 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/bin/bench.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/bin/bench.rs @@ -41,7 +41,7 @@ use std::sync::Arc; use std::time::Instant; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use hashbrown::HashMap as HbMap; use hashmap_v3::*; @@ -101,13 +101,13 @@ unsafe fn reset_table_async( stream: &Arc, ) -> Result<(), Box> { unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( map.ctrl.cu_deviceptr(), 0xFF, map.ctrl.num_bytes(), stream.cu_stream(), )?; - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( map.slots.cu_deviceptr(), 0xFF, map.slots.num_bytes(), diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/lib.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/lib.rs index bfbb8a00e4..6298d91977 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/lib.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/src/lib.rs @@ -57,7 +57,7 @@ use std::sync::Arc; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaStream, DeviceBuffer}; use cuda_device::atomic::{AtomicOrdering, DeviceAtomicU32, DeviceAtomicU64}; use cuda_device::cooperative_groups::{ThreadGroup, WarpCollective, this_thread_block}; @@ -1069,13 +1069,13 @@ impl GpuSwissMap { let ctrl = DeviceBuffer::::zeroed(stream, capacity / GROUP)?; let slots = DeviceBuffer::::zeroed(stream, capacity)?; unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( ctrl.cu_deviceptr(), 0xFF, ctrl.num_bytes(), stream.cu_stream(), )?; - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( slots.cu_deviceptr(), 0xFF, slots.num_bytes(), @@ -1085,8 +1085,8 @@ impl GpuSwissMap { // Reinterpret the allocations as slices of the atomic wrapper // (`#[repr(transparent)]`, identical layout). The memset above // is byte-level so it is unaffected by the element type. - let ctrl = ctrl.cast_elem::(); - let slots = slots.cast_elem::(); + let ctrl = ctrl.wrap_elem::(); + let slots = slots.wrap_elem::(); Ok(Self { ctrl, @@ -1379,21 +1379,21 @@ impl GpuSwissMap { let new_ctrl = DeviceBuffer::::zeroed(stream, new_capacity / GROUP)?; let new_slots = DeviceBuffer::::zeroed(stream, new_capacity)?; unsafe { - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( new_ctrl.cu_deviceptr(), 0xFF, new_ctrl.num_bytes(), stream.cu_stream(), )?; - cuda_core::simt::memory::memset_d8_async( + cuda_core::memory::memset_d8_async( new_slots.cu_deviceptr(), 0xFF, new_slots.num_bytes(), stream.cu_stream(), )?; } - let new_ctrl = new_ctrl.cast_elem::(); - let new_slots = new_slots.cast_elem::(); + let new_ctrl = new_ctrl.wrap_elem::(); + let new_slots = new_slots.wrap_elem::(); let rehash_threads = self.capacity.min(REHASH_MAX_THREADS) as u32; let cfg = LaunchConfig::for_num_elems(rehash_threads); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.lock index f3fe99746f..4e5e83344c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/src/main.rs index 4fbb97cfb8..fcb76b45c5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/src/main.rs @@ -11,7 +11,7 @@ //! Build and run with: //! cargo oxide run hello_constant -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::kernel; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.lock index 7de74fdb2b..9fc9e597f3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/src/main.rs index 8dafaf4c98..75b82d0ad2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/src/main.rs @@ -12,7 +12,7 @@ //! Build and run with: //! cargo oxide run helper_fn -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.lock index 05cdf65074..d3c7da91ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -134,7 +154,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -149,7 +169,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -161,18 +181,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -183,7 +205,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -193,6 +217,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -213,6 +238,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -379,6 +408,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/src/main.rs index 1e267a6268..f7162ec01d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/src/main.rs @@ -30,8 +30,8 @@ //! 8. Both reordered (default Rust repr, mixed-size fields) and identity-order //! (`#[repr(C)]`) struct captures. -use cuda_async::simt::device_operation::DeviceOperation; -use cuda_core::simt::LaunchConfig; +use cuda_async::device_operation::DeviceOperation; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_launch, cuda_launch_async, cuda_module, load_kernel_module}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.lock index 6c77a5ff67..c2c4eab309 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/src/main.rs index 5d5989b2f4..8c1c5576a3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/src/main.rs @@ -23,7 +23,7 @@ use std::sync::Arc; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, load_kernel_module}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.lock index 86c75b8ff3..8218065906 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/src/main.rs index 3106b31652..e5d51c2ff9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/src/main.rs @@ -11,7 +11,7 @@ //! cargo oxide build iket_trace --arch sm_120 //! ``` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, iket, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock index d5aa2ae65e..ba70558bbb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock @@ -67,6 +67,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -140,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -152,18 +172,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -174,7 +196,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -183,6 +207,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -203,6 +228,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -287,6 +316,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.lock index 3836acca0b..d215dbecc5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/src/main.rs index 5888f2b09b..aed40c7828 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/src/main.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock index 1a1e93d7b5..10e2ffbac9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/src/main.rs index 1557e57f06..38c794f608 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/src/main.rs @@ -31,7 +31,7 @@ //! Usage: //! cargo oxide run index_field_assign -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock index cb8ef2d6db..e650f1857e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/src/main.rs index 037b086542..c9b22840d9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/src/main.rs @@ -24,7 +24,7 @@ //! reaches LLVM as its own definition. Otherwise rustc folds the small ones //! into the kernel before any attribute could be observed. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, device, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.lock index 53bc623360..6c265f6b2e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/src/main.rs index 18238fd1f7..48795d93f7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/src/main.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, ptx_asm}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock index fe88f463bb..0459e2e67c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/src/main.rs index c8e331e578..abc9883fb9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/src/main.rs @@ -5,7 +5,7 @@ //! End-to-end example for scalar and packed integer min/max intrinsics (sm_90+). -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::i16x2::{ max_relu_s16x2, max_s16x2, max_u16x2, min_relu_s16x2, min_s16x2, min_u16x2, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock index a6f716be40..e5b12c1c46 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/src/main.rs index adf0c23f65..140955dcd3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/src/main.rs @@ -13,7 +13,7 @@ //! cargo oxide run interior_mutability //! CUDA_OXIDE_NO_OPT=1 cargo oxide run interior_mutability -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock index 8746a26f4b..3d2813e85c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -117,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -129,24 +149,30 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "digest" version = "0.10.7" @@ -207,6 +233,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.lock index f22694c1a4..f8feb1a8bd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.lock index 6100df50d1..80c93260b7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -134,7 +154,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -149,7 +169,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -161,18 +181,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -183,7 +205,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -193,6 +217,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -213,6 +238,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -379,6 +408,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/src/main.rs index 5152441dd3..b57c29a532 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/src/main.rs @@ -186,7 +186,7 @@ mod execution { fn verify_launch_packets() -> Result<(), Box> { use cuda_core::{DeviceBuffer, LaunchConfig1D}; - use cuda_host::cuda_async::simt::{ + use cuda_host::cuda_async::{ device_context::{init_device_contexts, with_cuda_context}, device_operation::DeviceOperation, }; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.lock index 394622c351..eda8dcfb02 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/src/main.rs index 3c0fd45658..d77b163752 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/src/main.rs @@ -114,7 +114,7 @@ mod kernels { // ============================================================================= fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== Lane-Position Masks & Warp Prefix Sum ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.lock index 1e2ee59815..f4b234f9ca 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/src/main.rs index 071f388f15..e65167feb8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/src/main.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ DisjointSlice, @@ -39,8 +39,8 @@ fn main() -> Result<(), Box> { let context = CudaContext::new(0)?; let stream = context.default_stream(); let module = kernels::load(&context)?; - let counter_f32 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); - let counter_f64 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); + let counter_f32 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); + let counter_f64 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); let mut old_values = DeviceBuffer::<(f32, f64)>::zeroed(&stream, N)?; // SAFETY: the launch covers exactly N unique 1D indices, `old_values` has @@ -57,8 +57,8 @@ fn main() -> Result<(), Box> { } stream.synchronize()?; - let got_f32 = counter_f32.cast_elem::().to_host_vec(&stream)?[0]; - let got_f64 = counter_f64.cast_elem::().to_host_vec(&stream)?[0]; + let got_f32 = counter_f32.peel_elem::().to_host_vec(&stream)?[0]; + let got_f64 = counter_f64.peel_elem::().to_host_vec(&stream)?[0]; let old_values = old_values.to_host_vec(&stream)?; let mut old_f32 = old_values .iter() diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.lock index ecfb175d50..157cf73b72 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/src/main.rs index 8ddb77f31a..74eb386556 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/src/main.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{ DisjointSlice, @@ -269,8 +269,8 @@ fn main() -> Result<(), Box> { let stream = context.default_stream(); let module = kernels::load(&context)?; - let rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); - let rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); + let rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); + let rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); let mut old_values = DeviceBuffer::<(u32, u64)>::zeroed(&stream, N)?; // SAFETY: the launch covers exactly N unique one-dimensional indices. @@ -287,8 +287,8 @@ fn main() -> Result<(), Box> { } stream.synchronize()?; - let got_u32 = rmw_u32.cast_elem::().to_host_vec(&stream)?[0]; - let got_u64 = rmw_u64.cast_elem::().to_host_vec(&stream)?[0]; + let got_u32 = rmw_u32.peel_elem::().to_host_vec(&stream)?[0]; + let got_u64 = rmw_u64.peel_elem::().to_host_vec(&stream)?[0]; let old_values = old_values.to_host_vec(&stream)?; if got_u32 != N as u32 || got_u64 != N as u64 || !old_values_are_permutations(&old_values) { return Err(format!( @@ -298,8 +298,8 @@ fn main() -> Result<(), Box> { .into()); } - let cas_u32 = DeviceBuffer::from_host(&stream, &[7_u32])?.cast_elem::(); - let cas_u64 = DeviceBuffer::from_host(&stream, &[7_u64])?.cast_elem::(); + let cas_u32 = DeviceBuffer::from_host(&stream, &[7_u32])?.wrap_elem::(); + let cas_u64 = DeviceBuffer::from_host(&stream, &[7_u64])?.wrap_elem::(); let mut observed_u32 = DeviceBuffer::<(u32, u32)>::zeroed(&stream, 1)?; let mut observed_u64 = DeviceBuffer::<(u64, u64)>::zeroed(&stream, 1)?; @@ -319,8 +319,8 @@ fn main() -> Result<(), Box> { let observed_u32 = observed_u32.to_host_vec(&stream)?[0]; let observed_u64 = observed_u64.to_host_vec(&stream)?[0]; - let final_u32 = cas_u32.cast_elem::().to_host_vec(&stream)?[0]; - let final_u64 = cas_u64.cast_elem::().to_host_vec(&stream)?[0]; + let final_u32 = cas_u32.peel_elem::().to_host_vec(&stream)?[0]; + let final_u64 = cas_u64.peel_elem::().to_host_vec(&stream)?[0]; if observed_u32 != (7, 11) || observed_u64 != (7, 11) || final_u32 != 11 || final_u64 != 11 { return Err(format!( "legacy integer CAS mismatch: observed_u32={observed_u32:?}, observed_u64={observed_u64:?}, final_u32={final_u32}, final_u64={final_u64}" @@ -328,10 +328,10 @@ fn main() -> Result<(), Box> { .into()); } - let block_rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); - let block_rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); - let system_rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); - let system_rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.cast_elem::(); + let block_rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); + let block_rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); + let system_rmw_u32 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); + let system_rmw_u64 = DeviceBuffer::::zeroed(&stream, 1)?.wrap_elem::(); let mut block_old_values = DeviceBuffer::<(u32, u64)>::zeroed(&stream, N)?; let mut system_old_values = DeviceBuffer::<(u32, u64)>::zeroed(&stream, N)?; let scoped_cfg = LaunchConfig { @@ -358,10 +358,10 @@ fn main() -> Result<(), Box> { } stream.synchronize()?; - let block_final_u32 = block_rmw_u32.cast_elem::().to_host_vec(&stream)?[0]; - let block_final_u64 = block_rmw_u64.cast_elem::().to_host_vec(&stream)?[0]; - let system_final_u32 = system_rmw_u32.cast_elem::().to_host_vec(&stream)?[0]; - let system_final_u64 = system_rmw_u64.cast_elem::().to_host_vec(&stream)?[0]; + let block_final_u32 = block_rmw_u32.peel_elem::().to_host_vec(&stream)?[0]; + let block_final_u64 = block_rmw_u64.peel_elem::().to_host_vec(&stream)?[0]; + let system_final_u32 = system_rmw_u32.peel_elem::().to_host_vec(&stream)?[0]; + let system_final_u64 = system_rmw_u64.peel_elem::().to_host_vec(&stream)?[0]; let block_old_values = block_old_values.to_host_vec(&stream)?; let system_old_values = system_old_values.to_host_vec(&stream)?; if block_final_u32 != N as u32 @@ -379,12 +379,12 @@ fn main() -> Result<(), Box> { .into()); } - let block_cas = DeviceBuffer::from_host(&stream, &[7_u32])?.cast_elem::(); - let system_cas = DeviceBuffer::from_host(&stream, &[7_u32])?.cast_elem::(); + let block_cas = DeviceBuffer::from_host(&stream, &[7_u32])?.wrap_elem::(); + let system_cas = DeviceBuffer::from_host(&stream, &[7_u32])?.wrap_elem::(); let device_strong_cas = - DeviceBuffer::from_host(&stream, &[7_u32])?.cast_elem::(); + DeviceBuffer::from_host(&stream, &[7_u32])?.wrap_elem::(); let device_ordered_cas = - DeviceBuffer::from_host(&stream, &[7_u32])?.cast_elem::(); + DeviceBuffer::from_host(&stream, &[7_u32])?.wrap_elem::(); let mut block_observed = DeviceBuffer::<(u32, u32)>::zeroed(&stream, 1)?; let mut system_observed = DeviceBuffer::<(u32, u32)>::zeroed(&stream, 1)?; let mut device_observed = DeviceBuffer::<(u32, u32)>::zeroed(&stream, 1)?; @@ -413,10 +413,10 @@ fn main() -> Result<(), Box> { let system_observed = system_observed.to_host_vec(&stream)?[0]; let device_observed = device_observed.to_host_vec(&stream)?[0]; let device_ordered_observed = device_ordered_observed.to_host_vec(&stream)?[0]; - let block_cas_final = block_cas.cast_elem::().to_host_vec(&stream)?[0]; - let system_cas_final = system_cas.cast_elem::().to_host_vec(&stream)?[0]; - let device_cas_final = device_strong_cas.cast_elem::().to_host_vec(&stream)?[0]; - let device_ordered_final = device_ordered_cas.cast_elem::().to_host_vec(&stream)?[0]; + let block_cas_final = block_cas.peel_elem::().to_host_vec(&stream)?[0]; + let system_cas_final = system_cas.peel_elem::().to_host_vec(&stream)?[0]; + let device_cas_final = device_strong_cas.peel_elem::().to_host_vec(&stream)?[0]; + let device_ordered_final = device_ordered_cas.peel_elem::().to_host_vec(&stream)?[0]; if block_observed != (7, 11) || system_observed != (7, 11) || device_observed != (7, 11) diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.lock index 1aa0595d80..c5069a42a3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/src/main.rs index c2ab6b9545..f79fa3613a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/src/main.rs @@ -10,7 +10,7 @@ //! and `u32`. Together, these cases check that legacy typed pointers work //! throughout a function, not only at loads and stores. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.lock index d285772669..311d3a4c17 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/src/main.rs index 1c97759edd..8cc396fa5b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/src/main.rs @@ -9,7 +9,7 @@ //! example checks SwiGLU plus `sin`, `exp`, `sqrt`, `atan`, `atan2`, `acos`, //! and `tan` on both legacy and modern NVVM targets. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.lock index 460643529b..184c547eee 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/src/main.rs index aeebfe45ac..8b4bd44e53 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/src/main.rs @@ -23,7 +23,7 @@ //! Run with: //! cargo oxide run libm_math -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.lock index 3f172b6f05..00d5c80cb8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/src/main.rs index 12222c65d3..a17ad8431c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/src/main.rs @@ -9,7 +9,7 @@ //! instead of `#[cuda_module]`. It keeps the explicit API covered while the //! typed embedded-module API is the preferred ergonomic path. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_launch, load_kernel_module}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.lock index 4118e90b10..95655f56da 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/src/main.rs index 43856a4012..9b249fb53d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/src/main.rs @@ -5,7 +5,7 @@ //! Manual low-level launch API regression test for NVVM IR/libdevice artifacts. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_launch, load_kernel_module}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.lock index 9924014ffa..9e3354bfeb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/src/main.rs index 29d7d1ed7a..76e1ff293b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/src/main.rs @@ -32,7 +32,7 @@ //! Usage: //! cargo oxide run map_sum -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock index 8ce1894329..1112e209dc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/src/main.rs index 5c6f7e97df..27e4f4185f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/src/main.rs @@ -15,7 +15,7 @@ //! //! Exits 0 on PASS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.lock index f7f1e6c274..cbfeb0fd77 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/src/main.rs index 85131b60aa..8714b5d7a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/src/main.rs @@ -32,7 +32,7 @@ // style lints that do not improve a numeric smoke test. #![allow(clippy::too_many_arguments, clippy::type_complexity)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.lock index b7718a33a7..9e01577d3b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/src/main.rs index 78faf4af2e..5893ec0b3a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/src/main.rs @@ -35,7 +35,7 @@ //! //! Exits 0 on SUCCESS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.lock index 552212e5e4..5f259d40fb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/src/main.rs index 89544b519a..51e7953ee5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/src/main.rs @@ -23,7 +23,7 @@ //! //! Exits 0 on PASS, 1 on FAIL. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::{cuda_module, ltoir}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.lock index 76fc9daba6..ccfd2d0f5b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.lock @@ -75,6 +75,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -148,7 +168,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen 0.69.5", "libloading", @@ -160,18 +180,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -182,7 +204,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -191,6 +215,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -211,6 +236,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -275,6 +304,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/src/main.rs index 9848dabb43..335448298d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/src/main.rs @@ -873,7 +873,7 @@ fn build_pipeline(mathdx_root: &Path) -> Result { // Host-side test infrastructure // ============================================================================= -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use std::sync::Arc; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.lock index 47eec7302e..fd841846e5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/src/main.rs index 90cb8dee3d..3c83ce4c84 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/src/main.rs @@ -13,7 +13,7 @@ //! Build and run: //! cargo oxide run mcast_barrier_test -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::barrier::{ Barrier, fence_proxy_async_shared_cta, mbarrier_arrive_cluster, mbarrier_init, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.lock index 611d7a182b..30e5e66d25 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/src/main.rs index d5effc1497..ab6bb8bf59 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/src/main.rs @@ -20,7 +20,7 @@ //! cargo oxide run mem_swap //! CUDA_OXIDE_NO_OPT=1 cargo oxide run mem_swap -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.lock index 8791af0505..792f79c62a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/src/main.rs index 4a1ccfe81a..d927a1f103 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/src/main.rs @@ -45,7 +45,7 @@ //! cargo oxide run mir_projected_call_destination use core::intrinsics::mir::*; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.lock index da7974c0f7..a24fcaff9e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/src/main.rs index 9c52f8c0fd..852eff1a8b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/src/main.rs @@ -26,7 +26,7 @@ //! Build and run with: //! cargo oxide run mma_mxf8f6f4 --arch sm_120a -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread, wmma}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.lock index 4bd0371e5b..aa8b93853e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/src/main.rs index ffb8bad124..d007f4f61d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/src/main.rs @@ -13,7 +13,7 @@ // Kernels lock in both directions for u32 and f32, plus `f32::to_bits`, // a same-width scalar transmute that must stay on the plain bitcast path. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.lock index c52786e005..538dc654b7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/src/main.rs index 5472665507..223fbf5123 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/src/main.rs @@ -12,7 +12,7 @@ //! through it instead of rejecting the assignment before the generic //! projection walker can handle it. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock index bfd2fa4115..7169f10e20 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/src/main.rs index 3eb12d458d..06901ef3e9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/src/main.rs @@ -14,7 +14,7 @@ //! //! Run: cargo oxide run numeric_stress -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.lock index 781a8d3f06..8762fa7551 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/README.md index 2a11e96506..6a544416aa 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/README.md @@ -1,7 +1,7 @@ # nvls_all_reduce Single-process NVLS (NVLink SHARP) all-reduce across every multicast-capable -GPU in the machine, built on the `cuda_core::simt::vmm` multicast wrappers +GPU in the machine, built on the `cuda_core::vmm` multicast wrappers (`MulticastObject`, `Mapping::new_multicast`). ## What it demonstrates @@ -54,7 +54,7 @@ integers, so the sums are exact and order-independent). ## See also -- `cuda-core/tests/simt_vmm_multicast.rs` in cutile-rs: host-side team plumbing test +- `cuda-core/tests/vmm_multicast.rs` in cutile-rs: host-side team plumbing test (bind/map/unbind lifecycle) that runs without issuing `multimem`. - PTX ISA section 9.7.13.4 (`multimem.ld_reduce`, `multimem.st`, `multimem.red`). diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/src/main.rs index a38077d7c3..7c8771b558 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/src/main.rs @@ -26,8 +26,8 @@ use cuda_core::CudaContext; use cuda_core::IntoResult; -use cuda_core::simt::LaunchConfig; -use cuda_core::simt::vmm; +use cuda_core::LaunchConfig; +use cuda_core::vmm; use cuda_device::{cuda_module, kernel, ptx_asm}; use std::mem::MaybeUninit; use std::sync::Arc; @@ -149,7 +149,7 @@ fn main() { let input: Vec = (0..ELEMS).map(|i| input_value(g, i)).collect(); let stream = ctx.default_stream(); unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( views[g].0.base(), input.as_ptr(), bytes, @@ -197,7 +197,7 @@ fn main() { let mut result = vec![0f32; ELEMS]; let stream = ctx.default_stream(); unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( result.as_mut_ptr(), views[g].0.base(), bytes, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock index b5af8f30b4..479a2908b1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/src/main.rs index 4267dcb23c..57c45246b7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/src/main.rs @@ -25,7 +25,7 @@ use core::mem; use core::ptr; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.lock index f5893ba8a7..c83dbaa4f3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/src/main.rs index 155eee9460..f2283c1ea9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/src/main.rs @@ -30,7 +30,7 @@ //! Run with: //! cargo oxide run option_ref_unwrap_or -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.lock index 85a0cbe346..a6e94aa576 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/src/main.rs index 658c5fa9d5..dae19c9b04 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/src/main.rs @@ -9,7 +9,7 @@ //! which on the GPU means the kernel must trap and the launch must fail. //! If the launch succeeds instead, the panic path was compiled away. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.lock index b80b1f9dcb..396740952e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/src/main.rs index c8019f6b74..ec4d88d8c3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/src/main.rs @@ -21,7 +21,7 @@ //! //! Run: cargo oxide run ord_cmp -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.lock index f4c8348b24..cde51b4d0b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/src/main.rs index 0b857b8fbb..ef18af3bd7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/src/main.rs @@ -28,7 +28,7 @@ //! `b` at byte offset 2. The padding byte itself is intentionally not checked: //! Rust does not guarantee a stable value for padding bytes. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{SharedArray, cuda_module, device, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.lock index 5c9aae69d3..f8f59f0159 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/src/main.rs index 17eae1e0a8..520d69662c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/src/main.rs @@ -10,7 +10,7 @@ //! old values for each 16-bit half independently because PTX does not promise //! that the two returned halves form one coherent 32-bit snapshot. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::atomic::{atom_add_bf16x2, atom_add_f16x2}; use cuda_device::{kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.lock index 7621003859..67e2eba6f5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/src/main.rs index ff825f45fc..229cc59b56 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/src/main.rs @@ -27,7 +27,7 @@ //! Run with: //! cargo oxide run padded_constants -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.lock index 1dd0ec55a8..e7ffd1f055 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock index f00ea7f41c..a67726271c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs index efccbe1b28..34195693cc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs @@ -12,7 +12,7 @@ use std::error::Error; use std::sync::Arc; use std::time::{Duration, Instant}; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaEvent, CudaStream, DeviceBuffer, PinnedHostBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.lock index 63206cdbab..da897d6dde 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/src/main.rs index 6bc95a1793..500ac71691 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/src/main.rs @@ -18,7 +18,7 @@ //! //! Run: cargo oxide run place_read_fallbacks -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.lock index ac860e426a..ea37ec21f1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.lock index f07ca73faa..643985bc0d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/src/main.rs index 0d5d196cbc..bc2521704b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/src/main.rs @@ -25,7 +25,7 @@ //! //! Run: `cargo oxide run primitive_stress` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock index cfaa742b4f..66ce726532 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/src/main.rs index 5900191f64..7891e0621a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/src/main.rs @@ -204,7 +204,7 @@ mod kernels { // ============================================================================= fn main() -> Result<(), Box> { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== GPU Printf Test (Unified) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.lock index 30bc83e4e5..4edcd8538a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/src/main.rs index 1f0cc06a07..731f1b01a8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/src/main.rs @@ -14,7 +14,7 @@ //! //! Run: cargo oxide run projected_place_read -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.lock index 42b08210f7..ee40f4f558 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.lock index de5ac4ae48..2a86336656 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/src/main.rs index f80c309259..75e4cb7a90 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/src/main.rs @@ -18,7 +18,7 @@ //! cargo oxide run ptr_copy //! CUDA_OXIDE_NO_OPT=1 cargo oxide run ptr_copy -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock index 0ec7cfb44f..0ab5ec5bbb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/src/main.rs index 5da1ff8bb5..62ce6bacc6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/src/main.rs @@ -13,7 +13,7 @@ //! offsets by the pointee type, retain pointer mutability through libcore's cast, //! and leave zero-sized pointees stationary. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.lock index fa5df66742..3a59743e5e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/src/main.rs index e3fba3047f..88d591bcaf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/src/main.rs @@ -8,7 +8,7 @@ //! raw address exactly; an old Debug-string byte parser mis-decoded these //! partially-initialized constant bytes. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock index 8d7de86367..267c574d08 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/src/main.rs index eb1675311b..81e50df688 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/src/main.rs @@ -35,7 +35,7 @@ #![feature(reborrow)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.lock index a74777042e..16d965ac09 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/src/main.rs index 420de7f079..1cefb4e5fa 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/src/main.rs @@ -77,7 +77,7 @@ mod kernels { // ============================================================================= fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== redux.sync f32 family (Blackwell family 10x) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.lock index bbb321da29..ae6b6c9847 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/src/main.rs index 6bdc88c801..b7ec14f342 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/src/main.rs @@ -126,7 +126,7 @@ mod kernels { // ============================================================================= fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== redux.sync integer family (sm_80+) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.lock index 3495a7289a..2e89d63014 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/src/main.rs index fdce5eeddd..4604acb2ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/src/main.rs @@ -74,7 +74,7 @@ mod kernels { // ============================================================================= fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== redux.sync.add Warp Reduction (sm_80+) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.lock index 58a388e07a..e4f85effad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/src/main.rs index aefcae951f..cdabb3d269 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/src/main.rs @@ -32,7 +32,7 @@ * projections, returning the field's base address). */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread, warp}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.lock index ec2404f932..b6859f4d61 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/src/main.rs index 5bc36bed83..bc50a9120d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/src/main.rs @@ -22,7 +22,7 @@ //! //! Run: cargo oxide run ref_operand_mul -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_host::cuda_module; use std::sync::Arc; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.lock index cc0078e643..3aec6380b4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/src/main.rs index 36f3678093..2568dfe2a9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/src/main.rs @@ -18,7 +18,7 @@ //! //! Run: `cargo oxide run repr_align_block_args` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.lock index fa4c02c424..aaa44498fb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/src/main.rs index 7ff826c025..0ef5c34b07 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/src/main.rs @@ -13,7 +13,7 @@ use core::marker::PhantomData; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{Uniform, cuda_module, device, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.lock index 6af97378f8..0866ab5d61 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/src/main.rs index d71f86e0c5..a268c968ec 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/src/main.rs @@ -52,7 +52,7 @@ * cargo oxide run repr_u32_enum_stride */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DeviceCopy}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.lock index 85b2f5f477..9ad7bc6ac4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/src/main.rs index ff14f911b2..02449c809e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/src/main.rs @@ -31,7 +31,7 @@ //! //! Run: cargo oxide run row_width_slice -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.lock index 67e55271c0..97f534a5b5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -259,6 +288,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/src/main.rs index bb566b6bc5..46a5a3fc25 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/src/main.rs @@ -34,7 +34,7 @@ //! cargo oxide run rustlantis-smoke use core::intrinsics::mir::*; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.lock index 013a3379fb..9938ba9184 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/src/main.rs index 8674c63f1c..c3c937cbc2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/src/main.rs @@ -34,7 +34,7 @@ //! Build and run with: //! cargo oxide run scoped_atomic_load_store --materialize-cubin -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.lock index 86eafa2b6e..1cc8f528af 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/src/main.rs index b62ec4c4d1..f614ec9cc3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/src/main.rs @@ -24,7 +24,7 @@ //! Usage: //! cargo oxide run select_unpredictable -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock index bbd354e38b..e0eadda39b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/src/main.rs index ed742c6bf6..81744c4132 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/src/main.rs @@ -17,7 +17,7 @@ #![allow(internal_features)] use core::intrinsics::mir::*; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock index e157fea2ac..da57a76a78 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/src/main.rs index aab836ad43..b003b6946f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/src/main.rs @@ -16,7 +16,7 @@ use core::intrinsics::mir::*; use core::num::NonZeroU32; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer, DeviceCopy}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock index 3195c5052d..14477e40aa 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/src/main.rs index a5ae803dc3..002d1f3b1a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/src/main.rs @@ -7,7 +7,7 @@ //! Persistent AS3 debug-info fixture for cuda-gdb/T27. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::barrier::Barrier; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.lock index e72dcc60f5..5536ebd051 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md index a022732b05..8947f5a809 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md @@ -13,58 +13,6 @@ unsafe fn helper(smem: *mut u8, offset: usize) -> u64 { } ``` -The no-feature build also runs the same probe through `SharedPtr`, a -CTA-shared pointer typed as such (the feature builds fail MIR verification -before any probe runs): - -```rust -let smem = DynamicSharedArray::::shared_ptr(); -shared_ptr_probe(Region { base: smem, offset }); - -#[inline(never)] -unsafe fn shared_ptr_probe(region: Region) -> u64 { - let full = region.base.cast::().byte_add(1024).cast::().as_ptr(); - // ... -} -``` - -`SharedPtr` lowers to an `addrspace(3)` pointer by its type, so the -generic helper's parameter, nested in a struct, is one too, and the call -needs no address-space adaptation. - -Two more shapes run after that probe, in `shared_ptr_shapes_probe`: a struct -with two `SharedPtr` fields passed by value to the non-inlined -`write_pair`, which writes through both, and a `SharedPtr` returned from -the non-inlined `select_slot` and used as a store target. Each of 64 lanes -writes its own slots, then overwrites one slot of its partner lane through the -returned pointer, so the final values prove both stores landed. - -`verify-code-shape.sh` checks the LLVM IR and PTX the build leaves next to -this file: - -| Helper | Required code shape | -| --- | --- | -| `shared_ptr_probe` | first parameter is `ptr addrspace(3)`, and its call site passes one | -| `write_pair` | both `Pair` fields arrive as `ptr addrspace(3)` parameters | -| `select_slot` | return type is `ptr addrspace(3)` | -| `shared_ptr_probe` in PTX | a `.func` that the kernel calls once, with `mbarrier.init.shared` inside its body | -| whole module | no `cvta.to.shared` and no `isspacep` in the PTX | - -The PTX half reads the unoptimized build. `scripts/smoketest.sh` lists this -example in `NO_OPT_SHAPE_EXAMPLES`, so after the optimized build decides the -verdict it rebuilds once with `CUDA_OXIDE_NO_OPT=1` and runs the script on -those artifacts. To run the script by hand, build the same way first: - -```bash -CUDA_OXIDE_NO_OPT=1 cargo oxide build shared_memory_helper -bash crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh -``` - -Without the surviving `.func`, a module-wide search for `cvta.to.shared` -proves little: an optimized build inlines the helper, and the generic-pointer -control compiles to the same instructions. The call into the `.func` is where -a generic pointer would have needed a conversion. - The launch supplies 2048 dynamic shared bytes. Four independent CTAs each use one thread. Nine launches cover three data offsets and three input values, including wrapping arithmetic. Every CTA checks both a volatile shared-memory @@ -131,11 +79,6 @@ RTX 5090, NVIDIA driver 580.173.02, and CUDA tools 13.3. | Local helper, all three inline policies | MIR verification rejects the call | | Dependency helper, all three inline policies | Same rejection | -With no features, the `SharedPtr` probe runs after the direct control and -reports `PASS shared-ptr: 9 launches, 36 CTA data/barrier checks`, and the -shapes probe reports `PASS shared-ptr-shapes: 3 launches, 384 slot checks`, -on the LLVM NVPTX route and with `--emit-nvvm-ir --arch sm_120`. - Both Compute Sanitizer runs of the direct control reported: ```text diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/src/main.rs index 6b25119cb0..6b7e30a2b0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/src/main.rs @@ -7,10 +7,9 @@ //! `inline(never)` requests a MIR call boundary; inspect generated code to //! determine whether the backend also keeps the calls. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; -use cuda_device::{DynamicSharedArray, SharedPtr, barrier::Barrier}; -use cuda_device::{cuda_module, kernel, launch_bounds, thread}; +use cuda_device::{DynamicSharedArray, cuda_module, kernel, launch_bounds, thread}; #[cfg(all(feature = "local-helper", feature = "cross-crate-helper"))] compile_error!("Select at most one helper feature; no features selects the direct control"); @@ -28,82 +27,6 @@ unsafe fn local_probe(smem: *mut u8, offset: usize) -> u64 { shared_memory_helper_lib::shared_probe_body!(smem, offset) } -/// The shared memory a typed probe works on, as a struct field. -#[derive(Clone, Copy)] -struct Region { - base: SharedPtr, - offset: usize, -} - -/// The probe body over a typed shared pointer, in a generic helper that is not -/// inlined. Unlike the raw `*mut u8` helpers above, it compiles: `SharedPtr` -/// is a CTA-shared pointer by its type, so the helper's parameter is one too, -/// and the barrier is addressed without converting a generic pointer. -/// -/// # Safety -/// As for `shared_probe_body!`: 2048 live shared bytes aligned to 1024, one -/// thread per CTA, and an initialized u64 at `offset`, below byte 1024. -#[inline(never)] -unsafe fn shared_ptr_probe(region: Region) -> u64 { - unsafe { - let bytes = region.base.cast::(); - let full: *mut Barrier = bytes.byte_add(1024).cast::().as_ptr(); - let data = bytes.byte_add(region.offset).cast::().as_ptr(); - cuda_device::barrier::mbarrier_init(full, 1); - let pending = !cuda_device::barrier::mbarrier_try_wait_parity(full, 0); - let _token = cuda_device::barrier::mbarrier_arrive(full); - let mut complete = 0u64; - let mut attempts = 0u32; - while attempts < 1024 { - if cuda_device::barrier::mbarrier_try_wait_parity(full, 0) { - complete = if pending { 1 } else { 2 }; - break; - } - attempts += 1; - } - let value = core::ptr::read_volatile(data); - core::ptr::write_volatile(data, value.wrapping_mul(3) ^ 0x1234_5678); - complete - } -} - -/// Two typed shared pointers in one struct, passed by value. -#[derive(Clone, Copy)] -struct Pair { - a: SharedPtr, - b: SharedPtr, -} - -/// Writes through both fields of a by-value `Pair`. The call boundary splits -/// the struct into two parameters; both must stay CTA-shared pointers. -/// -/// # Safety -/// Both fields must point at live shared memory with at least `index + 1` -/// `f32` slots. -#[inline(never)] -unsafe fn write_pair(pair: Pair, index: usize, value: f32) { - unsafe { - pair.a.add(index).as_ptr().write(value); - pair.b.add(index).as_ptr().write(value + 1000.0); - } -} - -/// Returns one slot of a `Pair` as a typed shared pointer. The return type -/// must stay a CTA-shared pointer so the caller's store needs no conversion. -/// -/// # Safety -/// As for `write_pair`. -#[inline(never)] -unsafe fn select_slot(pair: Pair, index: usize, second: bool) -> SharedPtr { - unsafe { - if second { - pair.b.add(index) - } else { - pair.a.add(index) - } - } -} - #[cuda_module] mod kernels { use super::*; @@ -129,54 +52,6 @@ mod kernels { *output.add(block * 2 + 1) = completed; } } - - /// The same probe through `SharedPtr`, a typed CTA-shared pointer, passed - /// to a non-inlined generic helper inside a struct. - #[kernel] - #[launch_bounds(1)] - pub fn shared_ptr_memory_probe(output: *mut u64, seed: u64, byte_offset: u64) { - let smem = DynamicSharedArray::::shared_ptr(); - let block = thread::blockIdx_x() as usize; - let offset = byte_offset as usize; - unsafe { - let data = smem.byte_add(offset).cast::().as_ptr(); - core::ptr::write_volatile(data, seed.wrapping_add(block as u64)); - let completed = shared_ptr_probe(Region { base: smem, offset }); - *output.add(block * 2) = core::ptr::read_volatile(data); - *output.add(block * 2 + 1) = completed; - } - } - - /// `SharedPtr` shapes a helper boundary must keep: a two-field struct of - /// typed shared pointers passed by value, and a typed shared pointer - /// returned from a helper and used as a store target. - #[kernel] - #[launch_bounds(64)] - pub fn shared_ptr_shapes_probe(output: *mut f32, seed: f32) { - const LANES: usize = 64; - let lane = thread::threadIdx_x() as usize; - let smem = DynamicSharedArray::::shared_ptr(); - // SAFETY: the launch supplies 2048 shared bytes, enough for the two - // 64-slot regions, and every lane writes only its own index and one - // slot of its partner lane, after a barrier. - unsafe { - let pair = Pair { - a: smem, - b: smem.add(LANES), - }; - let value = seed + lane as f32; - write_pair(pair, lane, value); - thread::sync_threads(); - // Each lane overwrites one slot of its partner lane through the - // pointer a helper returns; odd lanes pick the second field. - let partner = LANES - 1 - lane; - let slot = select_slot(pair, partner, lane % 2 == 1); - slot.as_ptr().write(-value); - thread::sync_threads(); - *output.add(lane) = *pair.a.add(lane).as_ptr(); - *output.add(LANES + lane) = *pair.b.add(lane).as_ptr(); - } - } } fn main() -> Result<(), Box> { @@ -204,85 +79,41 @@ fn main() -> Result<(), Box> { block_dim: (1, 1, 1), shared_mem_bytes: 2048, }; - for typed in [false, true] { - let variant = if typed { "shared-ptr" } else { variant }; - let mut cases = 0; - for seed in [0u64, 17, u64::MAX] { - for offset in [0u64, 64, 248] { - let out = output.cu_deviceptr() as *mut u64; - // SAFETY: four independent CTAs, one thread each; output holds - // two u64s per CTA. Data offsets are aligned and disjoint from - // the barrier. Every launch receives a fresh shared-memory - // allocation. - unsafe { - if typed { - module.shared_ptr_memory_probe(&stream, config, out, seed, offset)?; - } else { - module.shared_memory_probe(&stream, config, out, seed, offset)?; - } - } - let got = output.to_host_vec(&stream)?; - for block in 0..BLOCKS { - let expected = seed.wrapping_add(block as u64).wrapping_mul(3) ^ 0x1234_5678; - assert_eq!( - got[block * 2], - expected, - "{variant}: data, seed={seed}, offset={offset}, block={block}" - ); - assert_eq!( - got[block * 2 + 1], - 1, - "{variant}: barrier, seed={seed}, offset={offset}, block={block}" - ); - } - cases += 1; + let mut cases = 0; + for seed in [0u64, 17, u64::MAX] { + for offset in [0u64, 64, 248] { + // SAFETY: four independent CTAs, one thread each; output holds two + // u64s per CTA. Data offsets are aligned and disjoint from the + // barrier. Every launch receives a fresh shared-memory allocation. + unsafe { + module.shared_memory_probe( + &stream, + config, + output.cu_deviceptr() as *mut u64, + seed, + offset, + )?; } - } - println!( - "PASS {variant}: {cases} launches, {} CTA data/barrier checks", - cases * BLOCKS - ); - } - - const LANES: usize = 64; - let shapes_output = DeviceBuffer::::zeroed(&stream, LANES * 2)?; - let shapes_config = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (LANES as u32, 1, 1), - shared_mem_bytes: 2048, - }; - let seeds = [0.0f32, 17.0, 1024.0]; - for seed in seeds { - let out = shapes_output.cu_deviceptr() as *mut f32; - // SAFETY: one CTA of 64 lanes; output holds two f32 per lane, and the - // kernel uses 512 of the 2048 shared bytes supplied. - unsafe { module.shared_ptr_shapes_probe(&stream, shapes_config, out, seed)? }; - let got = shapes_output.to_host_vec(&stream)?; - for lane in 0..LANES { - let partner = LANES - 1 - lane; - // Lane `partner` overwrote one of this lane's slots: the second - // field when `partner` is odd, the first field otherwise. - let overwritten = -(seed + partner as f32); - let (want_a, want_b) = if partner % 2 == 1 { - (seed + lane as f32, overwritten) - } else { - (overwritten, seed + lane as f32 + 1000.0) - }; - assert_eq!( - got[lane], want_a, - "shared-ptr-shapes: field a, seed={seed}, lane={lane}" - ); - assert_eq!( - got[LANES + lane], - want_b, - "shared-ptr-shapes: field b, seed={seed}, lane={lane}" - ); + let got = output.to_host_vec(&stream)?; + for block in 0..BLOCKS { + let expected = seed.wrapping_add(block as u64).wrapping_mul(3) ^ 0x1234_5678; + assert_eq!( + got[block * 2], + expected, + "{variant}: data, seed={seed}, offset={offset}, block={block}" + ); + assert_eq!( + got[block * 2 + 1], + 1, + "{variant}: barrier, seed={seed}, offset={offset}, block={block}" + ); + } + cases += 1; } } println!( - "PASS shared-ptr-shapes: {} launches, {} slot checks", - seeds.len(), - seeds.len() * LANES * 2 + "PASS {variant}: {cases} launches, {} CTA data/barrier checks", + cases * BLOCKS ); Ok(()) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh deleted file mode 100755 index 5bc0dc0612..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh +++ /dev/null @@ -1,103 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# `SharedPtr` keeps a CTA-shared pointer typed `addrspace(3)` across a call. -# -# LLVM IR: the non-inlined `shared_ptr_probe` helper takes one as its -# parameter and its call site passes one, the two fields of a by-value `Pair` -# reach `write_pair` as two such parameters, and `select_slot` returns one. -# -# PTX: `shared_ptr_probe` survives as a `.func` that the kernel calls once, -# and that `.func` initializes the barrier with `mbarrier.init.shared` on its -# parameter. No instruction converts a generic address to a shared one -# (`cvta.to.shared`) or asks which window an address is in (`isspacep`). The -# call into the `.func` is where a generic pointer would have needed one; a -# module-wide search alone proves little once the helper is inlined, because -# the generic-pointer control compiles to the same instructions. -# -# The PTX half reads the unoptimized build. scripts/smoketest.sh lists this -# example in NO_OPT_SHAPE_EXAMPLES and rebuilds it with CUDA_OXIDE_NO_OPT=1 -# before running this script; by hand, build the same way first. -set -euo pipefail - -dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ll="${dir}/shared_memory_helper.ll" -ptx="${dir}/shared_memory_helper.ptx" - -for file in "${ll}" "${ptx}"; do - [[ -f "${file}" ]] || { echo "FAIL: ${file} not found; build the example first" >&2; exit 1; } -done - -# --- LLVM IR ----------------------------------------------------------------- - -if ! grep -Eq '^define .*shared_ptr_probe.*\(ptr addrspace\(3\) ' "${ll}"; then - echo "FAIL: shared_ptr_probe does not take a ptr addrspace(3) parameter" >&2 - grep -E '^define .*shared_ptr_probe' "${ll}" >&2 || true - exit 1 -fi - -if ! grep -Eq 'call i64 @[^(]*shared_ptr_probe[^(]*\(ptr addrspace\(3\) ' "${ll}"; then - echo "FAIL: the call to shared_ptr_probe does not pass a ptr addrspace(3)" >&2 - grep -E 'call i64 @[^(]*shared_ptr_probe' "${ll}" >&2 || true - exit 1 -fi - -if ! grep -Eq '^define .*write_pair.*\(ptr addrspace\(3\) %[A-Za-z0-9_.]+, ptr addrspace\(3\) %[A-Za-z0-9_.]+, ' "${ll}"; then - echo "FAIL: write_pair does not take both Pair fields as ptr addrspace(3) parameters" >&2 - grep -E '^define .*write_pair' "${ll}" >&2 || true - exit 1 -fi - -if ! grep -Eq '^define ptr addrspace\(3\) @[^(]*select_slot' "${ll}"; then - echo "FAIL: select_slot does not return ptr addrspace(3)" >&2 - grep -E '^define .*select_slot' "${ll}" >&2 || true - exit 1 -fi - -# --- PTX --------------------------------------------------------------------- - -# The body of the PTX `.func` whose name contains the symbol: from the header -# line, past the parameter list, to the closing brace. A forward declaration -# ends in `;` before any `{` and is skipped. -func_body() { - awk -v needle="$1" ' - !emit && !candidate && index($0, ".func") && index($0, needle) { candidate = 1 } - candidate && $0 ~ /^[[:space:]]*;[[:space:]]*$/ { candidate = 0; next } - candidate && $0 ~ /^[[:space:]]*\{/ { emit = 1; candidate = 0 } - emit { print } - emit && $0 ~ /^\}/ { exit } - ' "${ptx}" -} - -probe_body="$(func_body shared_ptr_probe)" -if [[ -z "${probe_body}" ]]; then - echo "FAIL: no PTX .func for shared_ptr_probe in ${ptx}" >&2 - echo "hint: this half reads the unoptimized build; rebuild with" >&2 - echo " CUDA_OXIDE_NO_OPT=1 cargo oxide build shared_memory_helper" >&2 - exit 1 -fi - -if ! grep -Eq 'mbarrier\.init\.shared(::cta)?\.b64' <<<"${probe_body}"; then - echo "FAIL: the shared_ptr_probe .func does not initialize the barrier with mbarrier.init.shared" >&2 - exit 1 -fi - -# One call into the helper: the symbol on a `call` line, outside `.func` -# headers and comments. -calls="$(grep -E 'shared_ptr_probe' "${ptx}" | grep -vE '\.func|^[[:space:]]*//' | grep -cE '(^|[[:space:]])call' || true)" -if [[ "${calls}" -ne 1 ]]; then - echo "FAIL: expected one PTX call into shared_ptr_probe, found ${calls}" >&2 - grep -nE 'call.*shared_ptr_probe' "${ptx}" >&2 || true - exit 1 -fi - -for forbidden in 'cvta\.to\.shared' 'isspacep'; do - if grep -Eq "${forbidden}" "${ptx}"; then - echo "FAIL: PTX contains ${forbidden}: a generic address was converted or classified" >&2 - grep -nE "${forbidden}" "${ptx}" >&2 - exit 1 - fi -done - -echo "PASS: SharedPtr stays addrspace(3) across shared_ptr_probe, write_pair and select_slot; the shared_ptr_probe .func is called once and initializes its barrier in shared memory; no cvta.to.shared or isspacep" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml deleted file mode 100644 index f7fd6f703c..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "shared_ptr_api" -version = "0.1.0" -edition = "2024" -authors = ["NVIDIA Corporation"] -license = "Apache-2.0" -description = "SharedPtr through helper returns, references, closures, arrays, Option, shared storage and transmute" - -# Mark as standalone crate (not part of parent workspace) -[workspace] - -# Build with: -# cargo oxide run shared_ptr_api - -[dependencies] -# Device-side intrinsics (thread indexing, barriers, etc.) -cuda-device = { path = "../../../cuda-device" } - -# Host-side utilities (CudaKernel trait, etc.) -cuda-host = { path = "../../../cuda-host" } - -# Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md deleted file mode 100644 index 52917d7c09..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md +++ /dev/null @@ -1,74 +0,0 @@ -# shared_ptr_api - -`SharedPtr` is a raw pointer into the issuing block's shared memory whose -address space (`addrspace(3)`) is part of its type. A generic `*mut T` loses -that fact at the first helper boundary; a `SharedPtr` keeps it through -parameters, return values, struct fields, closures and arrays, so the compiler -can use the shared-memory forms of the instructions that take it. This example -moves one through each of those boundaries and checks every value on the host. - -| Kernel | Boundary | Check | -| :------------------ | :-------------------------------------------------------------- | :-------------------------------------- | -| `static_round_trip` | `SharedArray::shared_ptr`, `from_raw(as_ptr())`, negative `offset` | each thread reads back its own slot | -| `closure_capture` | a closure that captured a `SharedPtr`, called through a generic `#[inline(never)]` helper | `3 * tid` | -| `addr_u32_deltas` | `addr_u32` of `shared_ptr_at(64)`, `add(3)`, `byte_add(7)`, and of `get_raw()` through `cvta_generic_to_shared_u32` | `[64, 12, 1, 7]` | -| `behind_references` | `&Region` (a struct holding the pointer) and `&mut SharedPtr` through non-inlined helpers | three reads summed | -| `returned_pointer` | a non-inlined helper returning `SharedPtr` | odd lanes read the pool, even the array | -| `array_by_value` | `[SharedPtr; 2]` passed by value | `7 + tid` | -| `struct_pointee` | `SharedPtr` and `cast::()` to its first field | `a + b + a` | -| `byte_views` | `byte_add` into a `u8` pool, viewed as `f64` and `u16` | `0.25 * tid + tid + 7` | -| `option_pointer` | `Option>` returned and passed | `Some` stores, `None` does not | -| `stored_in_shared` | `SharedArray, 64>`: pointers kept in shared memory | the mirrored slot's value | -| `transmuted` | `transmute` to `usize` and back in non-inlined helpers | the store lands and the stride is 4 | - -The helpers carry `#[inline(never)]`, so each boundary is a real call. In the -LLVM IR the default build leaves next to this file, `smem_fill`, `smem_load`, -`pick`, `maybe` and the closure's `apply` take `ptr addrspace(3)`, -`store_via_array` takes `[2 x ptr addrspace(3)]`, `pick` and `from_bits` -return `ptr addrspace(3)`, and `SLOTS` is a `[64 x ptr addrspace(3)]` global -in address space 3. The PTX converts no generic address to a shared one -except in `from_bits`: the integer view of a shared pointer is its generic -address, as for unions, so the transmute back is one `cvta.to.shared`. - -Structs that mix a `SharedPtr` with a narrower field, such as -`{ p: SharedPtr, x: u32 }`, stay out of this example. On the libNVVM -route a shared pointer is 4 bytes wide, and such a struct is not yet -supported in arrays or in shared memory; the book's shared memory chapter -lists this limit. - -## Build and run - -```bash -cargo oxide run shared_ptr_api -``` - -The compile-only lane of `scripts/smoketest.sh` also compiles this example -through libNVVM (`NVVM_VERIFY_EXAMPLES`), where the shared pointer is 4 bytes -wide: - -```bash -cargo oxide emit-ltoir shared_ptr_api --arch=sm_120 -``` - -## Expected output - -```text -=== SharedPtr API (64 threads) === - ok static_round_trip - ok closure_capture - ok addr_u32_deltas - ok behind_references - ok returned_pointer - ok array_by_value - ok struct_pointee - ok byte_views - ok option_pointer - ok stored_in_shared - ok transmuted -SUCCESS: shared_ptr_api, 11 kernels, 0 mismatches -``` - -## Hardware requirements - -Any CUDA GPU; the kernels use plain shared memory and `sync_threads()`. The -launch supplies 1024 bytes of dynamic shared memory. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs deleted file mode 100644 index 12ba608b2b..0000000000 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs +++ /dev/null @@ -1,535 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! `SharedPtr` across every boundary a helper call can put it behind. -//! -//! `SharedPtr` is a raw pointer into the issuing block's shared memory -//! whose address space is part of its type. Each kernel here moves one -//! through a boundary where a generic `*mut T` would lose that fact, or where -//! the lowering has to rebuild the value from parts: -//! -//! - a `from_raw` round trip and a negative `offset` on a static array, -//! - a closure that captured one, called through a generic helper, -//! - `addr_u32` deltas against `cvta_generic_to_shared_u32(get_raw())`, -//! - `&Region` and `&mut SharedPtr` through non-inlined helpers, -//! - a non-inlined helper that returns one, -//! - `[SharedPtr; 2]` passed by value, -//! - a struct pointee and a `cast` to its first field, -//! - `f64` and `u16` views of a dynamic byte pool, -//! - `Option>` returned and passed, -//! - `SharedPtr` values stored in shared memory itself, -//! - `transmute` to `usize` and back. -//! -//! The host checks every value and prints one SUCCESS line. Structs that mix -//! a `SharedPtr` with a narrower field stay out of this example: on the -//! libNVVM route a shared pointer is 4 bytes wide, and such a struct is not -//! yet supported in arrays or in shared memory. -//! -//! Usage: -//! cargo oxide run shared_ptr_api - -use cuda_core::simt::LaunchConfig; -use cuda_core::{CudaContext, DeviceBuffer}; -use cuda_device::shared::cvta_generic_to_shared_u32; -use cuda_device::{DynamicSharedArray, SharedArray, SharedPtr, cuda_module, kernel, thread}; - -const THREADS: usize = 64; - -/// A pointee with two fields; a `cast` to `f32` reaches the first. -#[derive(Clone, Copy)] -#[repr(C)] -struct Pair { - a: f32, - b: f32, -} - -/// A typed shared pointer next to an element offset, read through `&Region`. -#[derive(Clone, Copy)] -struct Region { - base: SharedPtr, - offset: usize, -} - -static mut TILE_F32: SharedArray = SharedArray::UNINIT; -static mut TILE_U32: SharedArray = SharedArray::UNINIT; -static mut PAIRS: SharedArray = SharedArray::UNINIT; -static mut SLOTS: SharedArray, THREADS> = SharedArray::UNINIT; - -/// Stores through a typed shared pointer inside a non-inlined generic helper. -/// -/// # Safety -/// `p.add(i)` must be in bounds of live shared memory. -#[inline(never)] -unsafe fn smem_fill(p: SharedPtr, i: usize, v: T) { - unsafe { p.add(i).as_ptr().write(v) } -} - -/// Loads through a typed shared pointer inside a non-inlined generic helper. -/// -/// # Safety -/// `p.add(i)` must be in bounds of initialized shared memory. -#[inline(never)] -unsafe fn smem_load(p: SharedPtr, i: usize) -> T { - unsafe { p.add(i).as_ptr().read() } -} - -/// Returns one of two typed shared pointers. The return type lowers to -/// `ptr addrspace(3)`. -#[inline(never)] -fn pick(a: SharedPtr, b: SharedPtr, second: bool) -> SharedPtr { - if second { b } else { a } -} - -/// Calls a closure that captured a `SharedPtr` by value. -#[inline(never)] -fn apply(f: F, i: usize) { - f(i) -} - -/// Reads through `&Region`: a deref, a field of `Region`, then the pointer. -/// -/// # Safety -/// `r.base.add(r.offset + i)` must be in bounds of initialized shared memory. -#[inline(never)] -unsafe fn read_region(r: &Region, i: usize) -> T { - unsafe { r.base.add(r.offset + i).as_ptr().read() } -} - -/// Advances a `SharedPtr` in place through `&mut SharedPtr`. -/// -/// # Safety -/// The advanced pointer must stay inside the same allocation. -#[inline(never)] -unsafe fn bump(slot: &mut SharedPtr, by: usize) { - unsafe { *slot = slot.add(by) } -} - -/// Stores through one of two `SharedPtr`s passed in an array by value. -/// -/// # Safety -/// `ptrs[which].add(i)` must be in bounds of live shared memory. -#[inline(never)] -unsafe fn store_via_array(ptrs: [SharedPtr; 2], which: usize, i: usize, v: T) { - unsafe { ptrs[which].add(i).as_ptr().write(v) } -} - -/// Stores through an `Option>` and reports whether it did. -/// -/// # Safety -/// A `Some` pointer must be in bounds of live shared memory at `i`. -#[inline(never)] -unsafe fn store_opt(p: Option>, i: usize, v: u32) -> u32 { - match p { - Some(p) => { - unsafe { p.add(i).as_ptr().write(v) }; - 1 - } - None => 0, - } -} - -/// Wraps a `SharedPtr` in `Some` or drops it, in a non-inlined helper. -#[inline(never)] -fn maybe(p: SharedPtr, take: bool) -> Option> { - if take { Some(p) } else { None } -} - -/// A `SharedPtr` as its address bits. -#[inline(never)] -fn to_bits(p: SharedPtr) -> usize { - // SAFETY: `SharedPtr` is `repr(transparent)` over a thin pointer. - unsafe { core::mem::transmute::, usize>(p) } -} - -/// Address bits back to a `SharedPtr`. -/// -/// # Safety -/// `bits` must come from `to_bits` on a live `SharedPtr`. -#[inline(never)] -unsafe fn from_bits(bits: usize) -> SharedPtr { - unsafe { core::mem::transmute::>(bits) } -} - -#[cuda_module] -mod kernels { - use super::*; - - /// A static array through `shared_ptr`, a `from_raw` round trip, and a - /// negative `offset` back to this thread's slot. - #[kernel] - pub fn static_round_trip(out: *mut f32) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: TILE_F32 is a `static mut SharedArray` of this kernel; each - // thread writes and reads its own slot, with a barrier between. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_F32); - smem_fill(tile, tid, tid as f32 + 0.5); - thread::sync_threads(); - let again = SharedPtr::from_raw(tile.as_ptr()); - let last = again.add(THREADS - 1); - let mine = last.offset(-((THREADS - 1 - tid) as isize)); - *out.add(tid) = smem_load(mine, 0); - } - } - - /// A closure that captured a `SharedPtr`, called through a generic - /// non-inlined helper. - #[kernel] - pub fn closure_capture(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel. - let tile = unsafe { SharedArray::shared_ptr(&raw mut TILE_U32) }; - // SAFETY: each thread writes its own slot. - let store = move |i: usize| unsafe { tile.add(i).as_ptr().write(i as u32 * 3) }; - apply(store, tid); - thread::sync_threads(); - // SAFETY: the slot was written before the barrier. - unsafe { *out.add(tid) = tile.add(tid).as_ptr().read() }; - } - - /// `addr_u32` deltas within the dynamic pool, and equality with - /// `cvta_generic_to_shared_u32` on the pool's raw pointer. - #[kernel] - pub fn addr_u32_deltas(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - let base = DynamicSharedArray::::shared_ptr(); - let at64 = DynamicSharedArray::::shared_ptr_at(64); - // SAFETY: `get_raw` points into the dynamic pool; thread 0 writes the - // four output slots. - unsafe { - let raw = cvta_generic_to_shared_u32(DynamicSharedArray::::get_raw()); - if tid == 0 { - *out.add(0) = at64.addr_u32() - base.addr_u32(); - *out.add(1) = base.add(3).addr_u32() - base.addr_u32(); - *out.add(2) = (base.addr_u32() == raw) as u32; - *out.add(3) = base.cast::().byte_add(7).addr_u32() - base.addr_u32(); - } - } - } - - /// `SharedPtr` behind references: `&Region` and `&mut SharedPtr`. - #[kernel] - pub fn behind_references(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - let pool = DynamicSharedArray::::shared_ptr(); - // SAFETY: the launch supplies 1024 dynamic shared bytes, enough for - // the 64 `u32` slots; every slot is written before the barrier. - unsafe { - pool.add(tid).as_ptr().write(100 + tid as u32); - thread::sync_threads(); - let region = Region { - base: pool, - offset: 0, - }; - let a = read_region(®ion, tid); - let mut cursor = pool; - bump(&mut cursor, tid); - let b = cursor.as_ptr().read(); - let shifted = Region { - base: pool, - offset: THREADS / 2, - }; - let c = read_region(&shifted, tid % (THREADS / 2)); - *out.add(tid) = a + b + c; - } - } - - /// A non-inlined helper that returns a `SharedPtr` picks between a static - /// array and the dynamic pool. - #[kernel] - pub fn returned_pointer(out: *mut f32) { - let tid = thread::threadIdx_x() as usize; - let dynp = DynamicSharedArray::::shared_ptr(); - // SAFETY: TILE_F32 is a `static mut SharedArray` of this kernel and - // the pool holds 64 `f32`; each thread writes its own slot in both. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_F32); - smem_fill(tile, tid, 1000.0 + tid as f32); - smem_fill(dynp, tid, 2000.0 + tid as f32); - thread::sync_threads(); - let chosen = pick(tile, dynp, tid % 2 == 1); - *out.add(tid) = smem_load(chosen, tid); - } - } - - /// `[SharedPtr; 2]` passed by value to a non-inlined helper. - #[kernel] - pub fn array_by_value(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - let pool = DynamicSharedArray::::shared_ptr(); - // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel and - // the pool holds 64 `u32`; each thread writes its own slot in one. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_U32); - store_via_array([tile, pool], tid % 2, tid, 7 + tid as u32); - thread::sync_threads(); - let v = if tid.is_multiple_of(2) { - tile.add(tid).as_ptr().read() - } else { - pool.add(tid).as_ptr().read() - }; - *out.add(tid) = v; - } - } - - /// A struct pointee, and a `cast` to its first field. - #[kernel] - pub fn struct_pointee(out: *mut f32) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: PAIRS is a `static mut SharedArray` of this kernel; each - // thread writes and reads its own element. - unsafe { - let pairs = SharedArray::shared_ptr(&raw mut PAIRS); - smem_fill( - pairs, - tid, - Pair { - a: tid as f32, - b: 2.0 * tid as f32, - }, - ); - thread::sync_threads(); - let p = smem_load(pairs, tid); - let first: f32 = smem_load(pairs.add(tid).cast::(), 0); - *out.add(tid) = p.a + p.b + first; - } - } - - /// `byte_add` into a dynamic byte pool, viewed as `f64` and as `u16`. - #[kernel] - pub fn byte_views(out: *mut f64) { - let tid = thread::threadIdx_x() as usize; - let pool = DynamicSharedArray::::shared_ptr(); - // SAFETY: the pool holds 1024 bytes; the `f64` view uses the first - // 512 and the `u16` view the next 128, each thread its own slot. - unsafe { - let d = pool.byte_add(8 * tid).cast::(); - smem_fill(d, 0, tid as f64 * 0.25); - let h = pool.byte_add(8 * THREADS + 2 * tid).cast::(); - smem_fill(h, 0, tid as u16 + 7); - thread::sync_threads(); - let dv: f64 = smem_load(pool.cast::(), tid); - let hv: u16 = smem_load(pool.byte_add(8 * THREADS).cast::(), tid); - *out.add(tid) = dv + hv as f64; - } - } - - /// `Option>` returned from and passed into non-inlined - /// helpers. - #[kernel] - pub fn option_pointer(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel; each - // thread writes and reads its own slot. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_U32); - let took = store_opt(maybe(tile, true), tid, 5 + tid as u32); - let skipped = store_opt(maybe(tile, false), tid, 0); - thread::sync_threads(); - *out.add(tid) = tile.add(tid).as_ptr().read() + 10 * took + 100 * skipped; - } - } - - /// `SharedPtr` values stored in shared memory: a `SharedArray` of them. - #[kernel] - pub fn stored_in_shared(out: *mut u32) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: TILE_U32 and SLOTS are `static mut SharedArray`s of this - // kernel; each thread writes its own slot in both and reads one slot - // after the barrier. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_U32); - let slots = SharedArray::shared_ptr(&raw mut SLOTS); - tile.add(tid).as_ptr().write(tid as u32 * 11); - slots.add(tid).as_ptr().write(tile.add(THREADS - 1 - tid)); - thread::sync_threads(); - let p: SharedPtr = slots.add(tid).as_ptr().read(); - *out.add(tid) = p.as_ptr().read(); - } - } - - /// `transmute` to `usize` and back, in non-inlined helpers. The high word - /// records whether the bits of this thread's slot sit `4 * tid` bytes - /// past the bits of the array base. - #[kernel] - pub fn transmuted(out: *mut u64) { - let tid = thread::threadIdx_x() as usize; - // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel; each - // thread writes and reads its own slot. - unsafe { - let tile = SharedArray::shared_ptr(&raw mut TILE_U32); - let bits = to_bits(tile.add(tid)); - let back = from_bits(bits); - back.as_ptr().write(tid as u32 * 13); - thread::sync_threads(); - let v = tile.add(tid).as_ptr().read(); - let strided = bits == to_bits(tile) + 4 * tid; - *out.add(tid) = v as u64 | ((strided as u64) << 32); - } - } -} - -/// Counts the mismatches of one kernel and prints its line. -fn check(name: &str, bad: usize, mismatches: &mut usize, kernels: &mut usize) { - if bad == 0 { - println!(" ok {name}"); - } else { - println!(" BAD {name}: {bad} mismatches"); - } - *mismatches += bad; - *kernels += 1; -} - -fn main() -> Result<(), Box> { - let ctx = CudaContext::new(0)?; - let stream = ctx.default_stream(); - let module = kernels::load(&ctx)?; - let config = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (THREADS as u32, 1, 1), - shared_mem_bytes: 1024, - }; - let mut mismatches = 0usize; - let mut kernels = 0usize; - println!("=== SharedPtr API ({THREADS} threads) ==="); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - // SAFETY: one block of THREADS threads; every output buffer holds one - // slot per thread (four slots for addr_u32_deltas, written by thread 0). - unsafe { module.static_round_trip(&stream, config, out.cu_deviceptr() as *mut f32)? }; - let got = out.to_host_vec(&stream)?; - check( - "static_round_trip", - (0..THREADS).filter(|&i| got[i] != i as f32 + 0.5).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.closure_capture(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - check( - "closure_capture", - (0..THREADS).filter(|&i| got[i] != i as u32 * 3).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, 4)?; - unsafe { module.addr_u32_deltas(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - let want = [64u32, 12, 1, 7]; - check( - "addr_u32_deltas", - (0..4).filter(|&i| got[i] != want[i]).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.behind_references(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - check( - "behind_references", - (0..THREADS) - .filter(|&i| { - let t = i as u32; - let half = (THREADS / 2) as u32; - got[i] != (100 + t) + (100 + t) + (100 + half + t % half) - }) - .count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.returned_pointer(&stream, config, out.cu_deviceptr() as *mut f32)? }; - let got = out.to_host_vec(&stream)?; - check( - "returned_pointer", - (0..THREADS) - .filter(|&i| { - let base = if i % 2 == 1 { 2000.0 } else { 1000.0 }; - got[i] != base + i as f32 - }) - .count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.array_by_value(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - check( - "array_by_value", - (0..THREADS).filter(|&i| got[i] != 7 + i as u32).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.struct_pointee(&stream, config, out.cu_deviceptr() as *mut f32)? }; - let got = out.to_host_vec(&stream)?; - check( - "struct_pointee", - (0..THREADS).filter(|&i| got[i] != 4.0 * i as f32).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.byte_views(&stream, config, out.cu_deviceptr() as *mut f64)? }; - let got = out.to_host_vec(&stream)?; - check( - "byte_views", - (0..THREADS) - .filter(|&i| got[i] != i as f64 * 0.25 + (i as f64 + 7.0)) - .count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.option_pointer(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - check( - "option_pointer", - (0..THREADS).filter(|&i| got[i] != 15 + i as u32).count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.stored_in_shared(&stream, config, out.cu_deviceptr() as *mut u32)? }; - let got = out.to_host_vec(&stream)?; - check( - "stored_in_shared", - (0..THREADS) - .filter(|&i| got[i] != (THREADS - 1 - i) as u32 * 11) - .count(), - &mut mismatches, - &mut kernels, - ); - - let out = DeviceBuffer::::zeroed(&stream, THREADS)?; - unsafe { module.transmuted(&stream, config, out.cu_deviceptr() as *mut u64)? }; - let got = out.to_host_vec(&stream)?; - check( - "transmuted", - (0..THREADS) - .filter(|&i| got[i] != (i as u64 * 13) | (1u64 << 32)) - .count(), - &mut mismatches, - &mut kernels, - ); - - if mismatches == 0 { - println!("SUCCESS: shared_ptr_api, {kernels} kernels, 0 mismatches"); - Ok(()) - } else { - println!("FAIL: shared_ptr_api, {mismatches} mismatches across {kernels} kernels"); - std::process::exit(1) - } -} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock index 42647acf94..7b563c7c9c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/src/main.rs index bc9bbac296..e0d8362744 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/src/main.rs @@ -45,7 +45,7 @@ #![allow(static_mut_refs)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, cuda_module, device, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.lock index 7517a55603..70d27c0686 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/src/main.rs index d70b97e58d..c402529f99 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/src/main.rs @@ -22,7 +22,7 @@ #![allow(static_mut_refs)] -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, device, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.lock index 98f07a9a47..619380b460 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/src/main.rs index edadf18425..1941d2d3a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/src/main.rs @@ -24,7 +24,7 @@ //! //! Run: `cargo oxide run shared_stage_select` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock index e60e6cf190..57605c1e94 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/src/main.rs index 7852b13c0f..1f82d5374d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/src/main.rs @@ -24,7 +24,7 @@ //! Build and run with: //! cargo oxide run sharedmem -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.lock index 46b3b08f05..b0d71c63c3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/src/main.rs index c25cc2b641..2f3e326536 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/src/main.rs @@ -126,7 +126,7 @@ mod kernels { const WARP: usize = 32; fn main() { - use cuda_core::simt::LaunchConfig; + use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; println!("=== 64-bit warp shuffle (u64 / f64) ===\n"); diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.lock index 52b8d52887..7f50b5cc06 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/src/main.rs index 99fa6d3b48..d912a46282 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/src/main.rs @@ -20,7 +20,7 @@ //! Usage: //! cargo oxide run slice_from_raw_parts_cast -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock index 20152e80f8..5652f97f02 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/src/main.rs index 1f8b0a9b52..66c6f7ccda 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/src/main.rs @@ -42,7 +42,7 @@ * address through fat-pointer deref"). */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.lock index e28e9d5ddf..7f787283d9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/src/main.rs index a7d7ca1c3b..fff2c5c707 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/src/main.rs @@ -14,7 +14,7 @@ //! or the launch must fail (thread trapped). //! If the sentinel is overwritten, the kernel performed an out-of-bounds global read. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.lock index 0ac8bde6f2..83f7bf3a01 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/src/main.rs index 5256f4f31f..4fab01a5c9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/src/main.rs @@ -20,7 +20,7 @@ //! cargo oxide run slice_iterators //! CUDA_OXIDE_NO_OPT=1 cargo oxide run slice_iterators -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock index fe2e92d2b2..54da984b1f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/src/main.rs index 9dd1471cd8..c2ffe061cb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/src/main.rs @@ -12,7 +12,7 @@ // pointer with `Rvalue::Aggregate(AggregateKind::RawPtr(..), [data_ptr, len])`. // The importer had no arm for that aggregate kind, so any kernel that passed // a sub-slice to a helper function failed to compile. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.lock index d446e6625f..95fe6dfc1f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/src/main.rs index 02f3142588..d48de0e790 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/src/main.rs @@ -336,7 +336,7 @@ fn build_pipeline() -> Result { // the results. // ============================================================================= -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use std::sync::Arc; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.lock index c781eb7524..58b69c66dd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/src/main.rs index 45c9b4ba13..4216c4f43f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/src/main.rs @@ -8,7 +8,7 @@ //! See the PTX ISA's "Matrix Fragments for sparse mma.m16n8k64" for the layout. //! `--probe` additionally compares candidate layouts on the current GPU. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread, wmma}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.lock index 03d67c5e3a..fcb4ffff52 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/src/main.rs index 7b02c14e88..3e939edca8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/src/main.rs @@ -21,7 +21,7 @@ //! lane's metadata register also changes the result; //! - values stay small (A<=7, B<=8, C<=31, sums < 1024) so f16, bf16, and //! tf32 all accumulate exactly and the host reference is integer math. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread, wmma}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.lock index 4324edf483..78fee50137 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/src/main.rs index 4cdde3d73a..58e7a5b3de 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/src/main.rs @@ -21,7 +21,7 @@ //! Usage: //! cargo oxide run split_at_mut -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock index bbf2681720..aa84e68767 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock index 15f0537070..8acb78b133 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/src/main.rs index ea50b15c3f..c1b2c09974 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/src/main.rs @@ -11,7 +11,7 @@ //! //! Run: `cargo oxide run static_slice_addend` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::kernel; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock index c074858de1..0904ae90d8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/src/main.rs index 0443d1661b..158eada8e4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/src/main.rs @@ -10,7 +10,7 @@ //! //! Run: `cargo oxide run static_slice_unsize` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::kernel; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock index 84cba98cab..9c6abcacda 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/src/main.rs index 544c5e9714..e189c7782b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/src/main.rs @@ -23,7 +23,7 @@ //! //! Run: cargo oxide run step_by -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.lock index 690c24a327..948b9bbcb7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/src/main.rs index 6a3bdca1b0..9c0acc0e8b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/src/main.rs @@ -9,7 +9,7 @@ //! metadata word. The importer must preserve both pieces when the constant is //! materialized for GPU code. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock index fe53acdb76..275ac3023b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/src/main.rs index c94cc5edf4..d317d8dc13 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/src/main.rs @@ -25,7 +25,7 @@ //! //! Run: `cargo oxide run struct_field_layout` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.lock index 7ef72e0da5..dd2c8dc7eb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/src/main.rs index f05dd15c0c..6c87da4f20 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/src/main.rs @@ -22,7 +22,7 @@ * cargo oxide pipeline subslice_projection */ -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.lock index 03811356e6..af26a73542 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/src/main.rs index 8fb1f1e5d4..f8cc1a2e68 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/src/main.rs @@ -25,7 +25,7 @@ //! Usage: //! cargo oxide run switchint_128bit -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock index f152cecc4c..0199a285b0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/src/main.rs index 93b01ee4b4..bf9dc33c5f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/src/main.rs @@ -12,7 +12,7 @@ //! Build and run with: //! cargo oxide run swizzle_smem -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, kernel, swizzle::Swizzle, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock index 8c6054ce23..edb8093db5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/src/main.rs index acb06ddc36..c0a8b831c0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/src/main.rs @@ -32,7 +32,7 @@ //! Build and run with: //! cargo oxide run tcgen05 -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer, sys}; use cuda_device::barrier::Barrier; use cuda_device::shared::{SharedArray, cvta_generic_to_shared_offset}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.lock index 651ffe661e..0f27e83d8e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/src/main.rs index 88fd1e2778..b1a8b21e82 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/src/main.rs @@ -10,7 +10,7 @@ //! Build and run with: //! cargo oxide run tcgen05_matmul -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{ CudaContext, CudaStream, DeviceBuffer, sys::{ diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.lock index 75c208fc87..328f417235 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock index cf9576c158..3db2c23a61 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/src/main.rs index 87fa66fa63..d0a37f3383 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/src/main.rs @@ -15,7 +15,7 @@ //! Build and run with: //! cargo oxide run tiled_gemm -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.lock index c86e4ed6e2..d9a057a886 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -176,6 +200,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -240,6 +268,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.lock index 5a9c81ad8f..10c09a6636 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md index 0f8b32dba3..2fcae23cee 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md @@ -11,13 +11,13 @@ For TMA multicast (also sm_90+), see the [`tma_multicast`](../tma_multicast/) ex 1. **tma_copy_2d_test**: Async 2D tile copy from global to shared memory 2. **tma_copy_2d_cta_test**: The same copy with `cp_async_bulk_tensor_2d_g2s_cta`, - whose destination is a `SharedPtr`, the issuing CTA's own shared memory by type + targeting the issuing CTA's own shared memory (`.shared::cta`) 3. **tma_pipeline_test**: TMA with mbarrier for completion tracking `cp_async_bulk_tensor_2d_g2s` addresses cluster shared memory, so its destination may belong to another CTA of the cluster. The `_cta` form takes a -`SharedPtr` destination, here `SharedArray::shared_ptr(&raw mut TILE)`, and -emits `cp.async.bulk.tensor.2d.shared::cta.global` without converting any +`*mut u8` destination into this CTA's shared memory (`&raw mut TILE as *mut u8`) +and emits `cp.async.bulk.tensor.2d.shared::cta.global` without converting any address; `verify-code-shape.sh` checks both forms in the generated PTX. All kernels use `#[grid_constant] tensor_map: &TmaDescriptor`. The generated diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/src/main.rs index aa21c06446..c24c593f95 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/src/main.rs @@ -10,8 +10,7 @@ //! Demonstrates TMA (Tensor Memory Accelerator) usage: //! - cp_async_bulk_tensor_2d_g2s: Async 2D tensor copy global → shared //! - cp_async_bulk_tensor_2d_g2s_cta: the same copy into the issuing CTA's -//! own shared memory, whose destination is a `SharedPtr`, a CTA-shared -//! pointer by type +//! own shared memory (`.shared::cta`) //! - mbarrier: Barrier-based completion tracking //! - `#[grid_constant]`: the descriptor is passed by value in the launch packet, //! without allocating or uploading a separate descriptor buffer. @@ -22,7 +21,7 @@ //! Build and run with: //! cargo oxide run tma_copy -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{ CudaContext, CudaStream, DeviceBuffer, sys::{ @@ -161,7 +160,7 @@ mod kernels { // before any thread reads TILE. unsafe { cp_async_bulk_tensor_2d_g2s_cta( - SharedArray::shared_ptr(&raw mut TILE).cast::(), + &raw mut TILE as *mut u8, tensor_map, tile_x, tile_y, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh index 818506f8f9..f7aad86e32 100755 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh @@ -4,8 +4,8 @@ # cp_async_bulk_tensor_2d_g2s_cta must copy into the issuing CTA's shared # memory: its kernel emits the `.shared::cta` destination form and, since its -# `SharedPtr` destination is shared by type, converts no address to the shared -# or cluster window. The cluster API must keep its +# `&raw mut TILE as *mut u8` destination stays in `addrspace(3)`, converts no +# address to the shared or cluster window. The cluster API must keep its # `.shared::cluster` form, since its destination may belong to another CTA. set -euo pipefail diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.lock index 3977ab9ddc..3f421473f7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/src/main.rs index 5a1e325a5d..58805648dc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/src/main.rs @@ -28,7 +28,7 @@ //! Build and run with: //! cargo oxide run tma_multicast -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{ CudaContext, CudaStream, DeviceBuffer, sys::{ diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.lock index 29f8b4dec6..c7625cb91c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.lock index 0e2998f8d5..7f1557e42f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/src/main.rs index 2c3dbba3c6..fde67b4e40 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/src/main.rs @@ -18,7 +18,7 @@ //! //! Run: cargo oxide run transmute_scalar_struct -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.lock index de10cae507..bed50145cd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/src/main.rs index e95958c280..7cced6c01b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/src/main.rs @@ -10,7 +10,7 @@ //! //! Run: `cargo oxide run tuple_array_provenance` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock index 48212e960c..1497542c61 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/src/main.rs index e167273513..49f046b9c1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/src/main.rs @@ -3,7 +3,7 @@ //! Regression for pointer provenance in a direct tuple constant. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock index 47aad8a282..bb303a576f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/src/main.rs index 4fd3d31fab..6148e7b040 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/src/main.rs @@ -28,7 +28,7 @@ //! //! Run: `cargo oxide run tuple_field_lookup` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.lock index 2ac68321ee..8ed26af0fd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/src/main.rs index 733e2f727d..e9a75c72e1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/src/main.rs @@ -20,7 +20,7 @@ //! Build and run with: //! cargo oxide run tuple_return -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.lock index 4affc8b321..aefbe9aa32 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/src/main.rs index 50c2279e20..cdc68b231b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/src/main.rs @@ -12,7 +12,7 @@ //! cargo oxide run unaligned_memory //! CUDA_OXIDE_NO_OPT=1 cargo oxide run unaligned_memory -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock index 53f368594d..8182f8bc56 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/src/main.rs index 89e954f0f1..760f1a27ea 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/src/main.rs @@ -39,7 +39,7 @@ //! Build and run with: //! cargo oxide run unchecked_indexing -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.lock index 1517384a3b..6f5e3243f6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/src/main.rs index 6e68befde0..8e8514c3de 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/src/main.rs @@ -14,7 +14,7 @@ //! Run: cargo oxide run uninit_const use core::mem::MaybeUninit; -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.lock index c2a4fec2d7..2c08f8b6a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/src/main.rs index b3d340c3a1..d563abae45 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/src/main.rs @@ -11,7 +11,7 @@ //! passing a union through an ordinary device function. It also checks that //! an over-aligned union keeps its size and array stride. -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.lock index 71e089c287..b393294873 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/src/main.rs index 37ecf93b63..2277bcdc2a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/src/main.rs @@ -6,7 +6,7 @@ //! `#[unroll]` must preserve checked indexing and division-by-zero checks. //! Run: cargo oxide run unroll_bounds_check -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.lock index 780f6f3e9d..3a825c4b57 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/src/main.rs index 2469ed630a..9265bf3acf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/src/main.rs @@ -27,7 +27,7 @@ //! //! Run: cargo oxide run unroll_smoke -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock index 9d76a571bd..42dc3b1172 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/src/main.rs index 01fb68414b..75e740b7cc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/src/main.rs @@ -20,7 +20,7 @@ // No #![cfg_attr(cuda_device, no_std)] - this compiles as ONE unit! -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock index 6cb5e3d2a9..d5fca34564 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs index 255b16ace7..9861b39880 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/src/main.rs @@ -14,7 +14,7 @@ //! //! Run: `cargo oxide run vectorization` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{access, cuda_module}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.lock index 51689b847a..f70c5a69ab 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/src/main.rs index 1e70397774..1e80bd6c30 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/src/main.rs @@ -5,7 +5,7 @@ //! //! Run: `cargo oxide run volatile_memory` -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.lock index 49805d2fe3..a3f82dae4c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/src/main.rs index 3cf6f60042..034265603c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/src/main.rs @@ -14,7 +14,7 @@ //! Build and run with: //! cargo oxide run warp_reduce -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::cooperative_groups::{ ThreadGroup, block_reduce, block_scan, ops::Sum, this_thread_block, diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.lock index 1b446984c3..31cc22f4b7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock index 11fbfa1d79..fa3bb7a8e8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/src/main.rs index faa4e622a9..090ab40ee5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/src/main.rs @@ -16,7 +16,7 @@ //! Build and run with: //! cargo oxide run wgmma -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; use cuda_device::shared::SharedArray; use cuda_device::wgmma::{make_smem_desc, wgmma_commit_group, wgmma_fence, wgmma_wait_group}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.lock index e5e504178f..007576977d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -128,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +160,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +184,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +195,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +216,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -255,6 +284,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs index 97e5b7c45a..078099ebaf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs @@ -126,7 +126,7 @@ //! Kernel export names are separate — they use `compute_kernel_export_name` //! with human-readable base names derived from the `#[kernel]` macro. -use mir_importer::is_panic_entry_path; +use mir_importer::{is_cute_compiler_hook_path, is_panic_entry_path}; use rustc_hir::def_id::{DefId, LOCAL_CRATE}; use rustc_index::{Idx, bit_set::DenseBitSet}; use rustc_middle::mir::visit::Visitor; @@ -135,7 +135,8 @@ use rustc_middle::mir::{ }; use rustc_middle::mono::{CodegenUnit, MonoItem}; use rustc_middle::ty::{ - Instance, InstanceKind, ShimKind, Ty, TyCtxt, TyKind, TypeVisitableExt, TypingEnv, + GenericArgsRef, Instance, InstanceKind, ShimKind, Ty, TyCtxt, TyKind, TypeVisitableExt, + TypingEnv, }; use rustc_span::Span; use std::collections::{HashMap, HashSet, VecDeque}; @@ -235,6 +236,144 @@ enum CollectDecision { Forbidden { crate_name: String, fn_path: String }, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CuteImplOwner { + Tensor, + TensorMut, + /// `TensorMut<'_, T, Tile1D, T>`: the impl whose `store` and + /// `store_at` the importer lowers to tile operations. The Contiguous1D + /// impl's `store_at` is ordinary Rust whose body must be collected. + TensorMutTile1D, + BlockScaledTensor, + BlockScaledThreadRow, + BlockScaledTile64, + LoadedBlockScaledTile64, + StaticPersistentTileScheduler, + WorkTileInfo, + PipelineTmaAsync, + PipelineState, + SharedTensor, + Mxf4TiledMma, + Mxf4AccumulatorTile2x8, + Mxf4ScaleTile128, + Sm120Epilogue128x128, + Sm120EpilogueWarp128x128, + PipelineTmaStore, +} + +fn is_cute_recognition_stub_identity( + crate_name: &str, + definition_path: &str, + impl_owner: Option, +) -> bool { + if !matches!(crate_name, "cute_rs" | "cute-rs") { + return false; + } + // Free functions and `__compiler` hooks: the importer's registry is the + // one list of recognized defining paths, so the collector asks it rather + // than keeping a copy that drifts when cute-rs moves a module. + if is_cute_compiler_hook_path(definition_path) { + return true; + } + + let Some(impl_owner) = impl_owner else { + return false; + }; + let module = match impl_owner { + CuteImplOwner::Tensor | CuteImplOwner::TensorMut | CuteImplOwner::TensorMutTile1D => { + "tensor" + } + CuteImplOwner::BlockScaledTensor => "block_scaled::tensor", + CuteImplOwner::BlockScaledThreadRow + | CuteImplOwner::BlockScaledTile64 + | CuteImplOwner::LoadedBlockScaledTile64 => "block_scaled::gemv", + CuteImplOwner::StaticPersistentTileScheduler | CuteImplOwner::WorkTileInfo => "scheduler", + CuteImplOwner::PipelineTmaAsync + | CuteImplOwner::PipelineState + | CuteImplOwner::PipelineTmaStore => "pipeline", + CuteImplOwner::SharedTensor => "copy::tiled", + CuteImplOwner::Mxf4TiledMma => "mma::block_scaled::tiled", + CuteImplOwner::Mxf4AccumulatorTile2x8 | CuteImplOwner::Mxf4ScaleTile128 => { + "mma::block_scaled::fragments" + } + CuteImplOwner::Sm120Epilogue128x128 | CuteImplOwner::Sm120EpilogueWarp128x128 => { + "epilogue::sm120" + } + }; + let Some(rest) = definition_path + .strip_prefix("::") + .and_then(|path| path.strip_prefix(module)) + .and_then(|path| path.strip_prefix("::")) + else { + return false; + }; + let mut segments = rest.split("::"); + let Some(impl_segment) = segments.next() else { + return false; + }; + let Some(method) = segments.next() else { + return false; + }; + if segments.next().is_some() || !impl_segment.starts_with("{impl#") { + return false; + } + + match impl_owner { + CuteImplOwner::Tensor => matches!( + method, + "from_slice" | "zipped_divide" | "local_tile" | "slice" | "is_full" | "base" | "load" + ), + CuteImplOwner::TensorMut => { + matches!( + method, + "from_disjoint_slice" | "zipped_divide" | "local_tile" | "slice" + ) + } + CuteImplOwner::TensorMutTile1D => matches!(method, "store" | "store_at"), + CuteImplOwner::BlockScaledTensor => matches!(method, "from_slices" | "thread_row"), + CuteImplOwner::BlockScaledThreadRow => method == "k_tile", + CuteImplOwner::BlockScaledTile64 => method == "load", + CuteImplOwner::LoadedBlockScaledTile64 => method == "dot_accumulate", + CuteImplOwner::StaticPersistentTileScheduler => { + matches!( + method, + "new_1d" | "has_work" | "get_current_work" | "advance_to_next_work" + ) + } + CuteImplOwner::WorkTileInfo => method == "tile_idx", + CuteImplOwner::PipelineTmaAsync => matches!( + method, + "from_raw_base" + | "from_raw_parts" + | "init" + | "producer_acquire" + | "producer_expect_tx" + | "consumer_wait" + | "consumer_release" + | "producer_tail" + ), + CuteImplOwner::PipelineState => matches!(method, "new" | "index" | "advance"), + CuteImplOwner::SharedTensor => method == "from_raw_parts", + CuteImplOwner::Mxf4TiledMma => matches!( + method, + "get_slice" | "load_scale_atom_128" | "load_a_128" | "get_b_tile_k64" + ), + CuteImplOwner::Mxf4AccumulatorTile2x8 => matches!(method, "zero" | "gemm"), + CuteImplOwner::Mxf4ScaleTile128 => method == "pairs_at_unchecked", + CuteImplOwner::Sm120Epilogue128x128 => matches!( + method, + "from_raw" | "get_slice" | "sync_reusable" | "sync_ready_for_tma" | "tma_half" + ), + CuteImplOwner::Sm120EpilogueWarp128x128 => method == "store_tile", + CuteImplOwner::PipelineTmaStore => { + matches!( + method, + "producer_acquire" | "producer_commit" | "producer_tail" + ) + } + } +} + // The prefix constants and substring/extractor helpers used below // (`KERNEL_PREFIX`, `is_kernel_symbol`, `kernel_base_name`, etc.) live in // the workspace-internal `reserved-oxide-symbols` crate. That crate is the @@ -1385,7 +1524,7 @@ impl<'tcx> DeviceCollector<'tcx> { // kind (a reified `#[target_feature]` fn, say) must be checked. if !matches!( func.instance.def, - InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAsPtr(..)) + InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAddr(..)) ) { self.check_host_cpu_only(def_id, ctx.user_span, &ctx, None); } @@ -1401,7 +1540,7 @@ impl<'tcx> DeviceCollector<'tcx> { // `instance_mir` body: drop glue, and (since // nightly-2026-08-28) the `::as_ptr` // shim that backs `FnPtr::addr`. - InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAsPtr(..)) + InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAddr(..)) ); if has_mir { // Use instance_mir for monomorphized MIR. @@ -1915,7 +2054,7 @@ impl<'tcx> DeviceCollector<'tcx> { if !matches!( resolved.def, InstanceKind::Item(_) - | InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAsPtr(..)) + | InstanceKind::Shim(ShimKind::DropGlue(..) | ShimKind::FnPtrAddr(..)) ) { return; } @@ -1978,6 +2117,13 @@ impl<'tcx> DeviceCollector<'tcx> { return; } + if self.is_cute_recognition_stub(resolved.def_id()) { + if self.verbose { + eprintln!("[collector] Skipping CuTe recognition stub: {raw_name}"); + } + return; + } + // Check if this is a device extern declaration (FFI with external LTOIR). // These have no MIR body but should be emitted as LLVM `declare` statements. if is_device_extern_symbol(&raw_name) { @@ -2032,7 +2178,7 @@ impl<'tcx> DeviceCollector<'tcx> { // `is_mir_available` is false for their def_id, but `instance_mir` // synthesises their body; do not skip those. if !self.tcx.is_mir_available(resolved.def_id()) - && !matches!(resolved.def, InstanceKind::Shim(ShimKind::FnPtrAsPtr(..))) + && !matches!(resolved.def, InstanceKind::Shim(ShimKind::FnPtrAddr(..))) { if self.verbose { eprintln!( @@ -2143,6 +2289,19 @@ impl<'tcx> DeviceCollector<'tcx> { .emit() } + if self.is_cute_recognition_stub(instance.def_id()) { + self.tcx + .dcx() + .struct_span_fatal( + call_span, + format!( + "`{target}` cannot be used as a function item in device code because CuTe operations require direct call-site lowering" + ), + ) + .with_help("call the CuTe function directly from a local device function") + .emit() + } + match self.should_collect_from_crate(instance.def_id()) { CollectDecision::Collect => {} CollectDecision::SkipIntentional => { @@ -2450,6 +2609,70 @@ impl<'tcx> DeviceCollector<'tcx> { false } + fn is_cute_recognition_stub(&self, def_id: DefId) -> bool { + let crate_name = self.tcx.crate_name(def_id.krate); + if !matches!(crate_name.as_str(), "cute_rs" | "cute-rs") { + return false; + } + let definition_path = self.tcx.def_path(def_id).to_string_no_crate_verbose(); + let impl_owner = self.tcx.impl_of_assoc(def_id).and_then(|impl_id| { + // `type_of(..).instantiate_identity()` returns `Unnormalized` + // since the eager-normalization refactor (rust-lang/rust#155345). + // An impl block's self type here is the nominal cute-rs ADT as + // written, so unwrapping without normalizing is sufficient. + match self + .tcx + .type_of(impl_id) + .instantiate_identity() + .skip_normalization() + .kind() + { + TyKind::Adt(def, args) => { + match self.tcx.item_name(def.did()).to_string().as_str() { + "Tensor" => Some(CuteImplOwner::Tensor), + "TensorMut" if self.tensor_impl_layout_is_tile_1d(args) => { + Some(CuteImplOwner::TensorMutTile1D) + } + "TensorMut" => Some(CuteImplOwner::TensorMut), + "BlockScaledTensor" => Some(CuteImplOwner::BlockScaledTensor), + "BlockScaledThreadRow" => Some(CuteImplOwner::BlockScaledThreadRow), + "BlockScaledTile64" => Some(CuteImplOwner::BlockScaledTile64), + "LoadedBlockScaledTile64" => Some(CuteImplOwner::LoadedBlockScaledTile64), + "StaticPersistentTileScheduler" => { + Some(CuteImplOwner::StaticPersistentTileScheduler) + } + "WorkTileInfo" => Some(CuteImplOwner::WorkTileInfo), + "PipelineTmaAsync" => Some(CuteImplOwner::PipelineTmaAsync), + "PipelineState" => Some(CuteImplOwner::PipelineState), + "SharedTensor" => Some(CuteImplOwner::SharedTensor), + "Mxf4TiledMma" => Some(CuteImplOwner::Mxf4TiledMma), + "Mxf4AccumulatorTile2x8" => Some(CuteImplOwner::Mxf4AccumulatorTile2x8), + "Mxf4ScaleTile128" => Some(CuteImplOwner::Mxf4ScaleTile128), + "Sm120Epilogue128x128" => Some(CuteImplOwner::Sm120Epilogue128x128), + "Sm120EpilogueWarp128x128" => Some(CuteImplOwner::Sm120EpilogueWarp128x128), + "PipelineTmaStore" => Some(CuteImplOwner::PipelineTmaStore), + _ => None, + } + } + _ => None, + } + }); + is_cute_recognition_stub_identity(crate_name.as_str(), &definition_path, impl_owner) + } + + /// Whether a cute-rs tensor impl's `Layout` generic argument is `Tile1D`. + /// + /// `Tensor` and `TensorMut` are `<'a, Element, Layout, Storage>`, so the + /// layout is the second type argument of the impl's self type. + fn tensor_impl_layout_is_tile_1d(&self, args: GenericArgsRef<'tcx>) -> bool { + args.types() + .nth(1) + .is_some_and(|layout| match layout.kind() { + TyKind::Adt(def, _) => self.tcx.item_name(def.did()).as_str() == "Tile1D", + _ => false, + }) + } + /// Recognize a generated intrinsic placeholder or reject a mismatched raw /// intrinsic crate before its panic-only host body enters collection. fn is_generated_intrinsic_placeholder_or_report_mismatch( @@ -2503,7 +2726,7 @@ impl<'tcx> DeviceCollector<'tcx> { .with_help( "use the cuda-intrinsics crate generated by this cuda-oxide revision, then rebuild the device crate", ) - .emit_fatal() + .emit() } /// Emits the heap-allocation diagnostic (issue #108) and aborts. @@ -2539,7 +2762,7 @@ impl<'tcx> DeviceCollector<'tcx> { "store the data in a fixed-size array `[T; N]` instead, or in a \ `SharedArray` when the scratch space should be shared by the thread block", ) - .emit_fatal() + .emit() } /// Refuses a device-reachable function that only makes sense on the host @@ -2979,9 +3202,9 @@ pub fn dump_device_mir_info<'tcx>(tcx: TyCtxt<'tcx>, functions: &[CollectedFunct #[cfg(test)] mod tests { use super::{ - device_runtime_checks_target, is_host_cpu_arch_def_path, is_kernel_entry_def_path, - is_spin_loop_expansion, is_std_float_inherent_method, truncate_path_for_box, - unsupported_codegen_protocol_root, + CuteImplOwner, device_runtime_checks_target, is_cute_recognition_stub_identity, + is_host_cpu_arch_def_path, is_kernel_entry_def_path, is_spin_loop_expansion, + is_std_float_inherent_method, truncate_path_for_box, unsupported_codegen_protocol_root, }; use reserved_oxide_symbols::{ DEVICE_PREFIX, KERNEL_PREFIX, LEGACY_DEVICE_PREFIX, LEGACY_KERNEL_PREFIX, @@ -3102,6 +3325,302 @@ mod tests { assert_eq!(device_runtime_checks_target(&targets), false_target); } + #[test] + fn cute_tensor_recognition_is_closed_over_helpers_and_typed_methods() { + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::make_tensor_read", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#7}::load", + Some(CuteImplOwner::Tensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#11}::store_at", + Some(CuteImplOwner::TensorMutTile1D), + )); + // `TensorMut::store_at` shares the method name but is + // ordinary Rust; only the Tile1D impl's scalar store is a stub. + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#10}::store_at", + Some(CuteImplOwner::TensorMut), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#11}::store", + Some(CuteImplOwner::TensorMutTile1D), + )); + // The Contiguous1D impl's view constructor and view changes, including + // the fused `local_tile`, stay stubs the importer intercepts. + for method in ["from_disjoint_slice", "zipped_divide", "local_tile"] { + assert!( + is_cute_recognition_stub_identity( + "cute_rs", + &format!("::tensor::{{impl#10}}::{method}"), + Some(CuteImplOwner::TensorMut), + ), + "{method}" + ); + } + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#6}::local_tile", + Some(CuteImplOwner::Tensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::__compiler::block_scaled_load_k64", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::tensor::{impl#17}::from_slices", + Some(CuteImplOwner::BlockScaledTensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::gemv::{impl#19}::k_tile", + Some(CuteImplOwner::BlockScaledThreadRow), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::gemv::{impl#21}::dot_accumulate", + Some(CuteImplOwner::LoadedBlockScaledTile64), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::scheduler::__compiler::scheduler_current_tile", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::scheduler::{impl#4}::has_work", + Some(CuteImplOwner::StaticPersistentTileScheduler), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::scheduler::{impl#2}::tile_idx", + Some(CuteImplOwner::WorkTileInfo), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::__compiler::pipeline_producer_expect_tx", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::{impl#9}::producer_expect_tx", + Some(CuteImplOwner::PipelineTmaAsync), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::{impl#3}::advance", + Some(CuteImplOwner::PipelineState), + )); + assert!( + is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::{impl#9}::from_raw_parts", + Some(CuteImplOwner::PipelineTmaAsync), + ), + "the rejected split-ring constructor must not be body-translated" + ); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::copy::tiled::__compiler::shared_tensor_overlay", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::copy::tiled::{impl#3}::from_raw_parts", + Some(CuteImplOwner::SharedTensor), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::mma::block_scaled::__compiler::tiled_gemm", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::mma::block_scaled::tiled::{impl#7}::load_a_128", + Some(CuteImplOwner::Mxf4TiledMma), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::mma::block_scaled::fragments::{impl#4}::zero", + Some(CuteImplOwner::Mxf4AccumulatorTile2x8), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::mma::block_scaled::fragments::{impl#5}::pairs_at_unchecked", + Some(CuteImplOwner::Mxf4ScaleTile128), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::sm120::__compiler::epilogue_sync_ready_for_tma", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::sm120::{impl#2}::tma_half", + Some(CuteImplOwner::Sm120Epilogue128x128), + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::sm120::{impl#9}::store_tile", + Some(CuteImplOwner::Sm120EpilogueWarp128x128), + )); + // Free functions with host-panicking bodies are recognized through the + // importer registry at their defining paths. + for path in [ + "::copy::vector::load_tile", + "::copy::vector::assume", + "::copy::cooperative::copy_g2s", + "::copy::tma::copy_tma_2d", + "::copy::tma::copy_tma_s2g_2d", + "::mma::load_matrix_a", + "::sm100::__compiler::sm100_tiled_mma", + "::attention::sm100::__compiler::tma_load", + ] { + assert!( + is_cute_recognition_stub_identity("cute_rs", path, None), + "{path}" + ); + } + // Import paths are never defining paths, so none of these may be + // recognized: the root module aliases cute_rs once had (`tile`, `tma`, + // `cooperative`, `tiled_copy`, `block_scaled_mma`, retired) and the + // retired `epilogue::__compiler` re-export. + for path in [ + "::tile::load_tile", + "::tma::copy_tma_2d", + "::cooperative::copy_g2s", + "::tiled_copy::__compiler::shared_tensor_overlay", + "::block_scaled_mma::__compiler::tiled_gemm", + "::epilogue::__compiler::epilogue_half", + ] { + assert!( + !is_cute_recognition_stub_identity("cute_rs", path, None), + "{path}" + ); + } + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::__compiler::tma_store_producer_commit", + None, + )); + assert!(is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::{impl#4}::producer_tail", + Some(CuteImplOwner::PipelineTmaStore), + )); + // The SM100 store pipeline shares the item name but lives in + // cute_rs::sm100::pipeline; its methods are ordinary Rust over the + // sm100 hooks, so the owner table must not claim them. + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::sm100::pipeline::{impl#4}::producer_tail", + Some(CuteImplOwner::PipelineTmaStore), + )); + + // Same-spelled methods on the layout ADT are ordinary Rust. + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#2}::is_full", + None, + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::tensor::{impl#7}::store", + Some(CuteImplOwner::Tensor), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::block_scaled::gemv::{impl#21}::value_pair", + Some(CuteImplOwner::LoadedBlockScaledTile64), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::scheduler::{impl#4}::for_worker", + Some(CuteImplOwner::StaticPersistentTileScheduler), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::pipeline::{impl#3}::phase", + Some(CuteImplOwner::PipelineState), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::mma::block_scaled::tiled::{impl#7}::load_b_pair", + Some(CuteImplOwner::Mxf4TiledMma), + )); + assert!(!is_cute_recognition_stub_identity( + "cute_rs", + "::epilogue::sm120::{impl#3}::store_atom", + Some(CuteImplOwner::Sm120EpilogueWarp128x128), + )); + assert!(!is_cute_recognition_stub_identity( + "other", + "::tensor::make_tensor_read", + None, + )); + } + + #[test] + fn cute_nested_impl_owners_reject_old_and_near_miss_paths() { + for (module, owner, method) in [ + ( + "block_scaled::tensor", + CuteImplOwner::BlockScaledTensor, + "from_slices", + ), + ( + "block_scaled::gemv", + CuteImplOwner::BlockScaledThreadRow, + "k_tile", + ), + ( + "mma::block_scaled::tiled", + CuteImplOwner::Mxf4TiledMma, + "load_a_128", + ), + ( + "mma::block_scaled::fragments", + CuteImplOwner::Mxf4AccumulatorTile2x8, + "zero", + ), + ( + "epilogue::sm120", + CuteImplOwner::Sm120EpilogueWarp128x128, + "store_tile", + ), + ("copy::tiled", CuteImplOwner::SharedTensor, "from_raw_parts"), + ] { + assert!(is_cute_recognition_stub_identity( + "cute_rs", + &format!("::{module}::{{impl#7}}::{method}"), + Some(owner) + )); + for path in [ + format!( + "::{}::{{impl#7}}::{method}", + module.split("::").next().unwrap() + ), + format!("::{module}_extra::{{impl#7}}::{method}"), + format!("::{module}::nested::{{impl#7}}::{method}"), + ] { + assert!( + !is_cute_recognition_stub_identity("cute_rs", &path, Some(owner)), + "{path}" + ); + } + } + } + #[test] fn kernel_def_paths_are_entry_points() { // Bare and fully-qualified kernel names: the marker is the final segment. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs index 941dd11d10..f22d0937d9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs @@ -103,7 +103,6 @@ use rustc_middle::mir::interpret::{AllocId, GlobalAlloc, Scalar}; use rustc_middle::mir::visit::Visitor; use rustc_middle::mir::{ConstOperand, ConstValue, Location}; use rustc_middle::mono::MonoItem; -use rustc_middle::ty::consts::ConstExt; use rustc_middle::ty::layout::{LayoutCx, LayoutOf}; use rustc_middle::ty::{EarlyBinder, Instance, InstanceKind, ShimKind, TypingEnv}; use rustc_middle::ty::{Ty, TyCtxt, TyKind}; @@ -218,7 +217,7 @@ fn rustc_ty_to_device_extern_type<'tcx>( "failed to compute layout for repr(transparent) device-extern type `{ty}`: {err:?}" ) })?; - let peeled = layout.peel_transparent_wrappers_from_non_1zst(&layout_cx); + let peeled = layout.peel_transparent_wrappers(&layout_cx); // A transparent type with no peelable non-1ZST field is not a // scalar/pointer wrapper that this device-extern ABI can represent. @@ -350,6 +349,10 @@ pub struct DeviceCodegenConfig { pub output_dir: PathBuf, /// Base name for output files (e.g., "kernel" → kernel.ll, kernel.ptx). pub output_name: String, + /// Optional exact path for the prepared semantic CUTLASS MLIR module. + pub mlir_output: Option, + /// Device artifact backend selected by the rustc frontend. + pub device_backend: mir_importer::DeviceBackend, /// Print verbose progress to stderr. pub verbose: bool, /// Dump raw rustc MIR before translation. @@ -365,6 +368,8 @@ impl Default for DeviceCodegenConfig { Self { output_dir: std::env::current_dir().unwrap_or_else(|_| ".".into()), output_name: "kernel".to_string(), + mlir_output: None, + device_backend: mir_importer::DeviceBackend::Native, verbose: false, dump_rustc_mir: false, dump_mir_dialect: false, @@ -1021,6 +1026,8 @@ pub fn generate_device_code<'tcx>( let output_dir = config.output_dir.clone(); let output_name = config.output_name.clone(); + let mlir_output = config.mlir_output.clone(); + let device_backend = config.device_backend.clone(); let verbose = config.verbose; let show_rustc_mir = config.dump_rustc_mir; let show_mir = config.dump_mir_dialect; @@ -1247,6 +1254,8 @@ pub fn generate_device_code<'tcx>( let pipeline_config = mir_importer::PipelineConfig { output_dir: output_dir.clone(), output_name: output_name.clone(), + mlir_output: mlir_output.clone(), + device_backend: device_backend.clone(), verbose, show_mir_dialect: show_mir, show_llvm_dialect: show_llvm, @@ -1409,6 +1418,8 @@ mod tests { let config = DeviceCodegenConfig::default(); assert!(!config.verbose); assert_eq!(config.output_name, "kernel"); + assert_eq!(config.mlir_output, None); + assert_eq!(config.device_backend, mir_importer::DeviceBackend::Native); } #[test] diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs index 635436eee3..09bf66fe4d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/generated_intrinsics.rs @@ -4,7 +4,7 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 3f56e2a1acd3d2b767a2a5c91c66a7f2e7276516404e8385da9d2dece37cf0d7. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. // LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated collector predicates for intrinsic placeholder functions. @@ -252,6 +252,7 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1016", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0819", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0314", @@ -366,6 +367,7 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0576", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0098", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1017", @@ -375,6 +377,7 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0309", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0310", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830", @@ -530,6 +533,7 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0124", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0123", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0520", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0153", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0146", @@ -1293,6 +1297,7 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::float::ex2_approx_f16", "cuda_intrinsics::float::ex2_approx_f32", "cuda_intrinsics::float::ex2_approx_ftz_f32", + "cuda_intrinsics::barrier::fence_acquire_shared_cluster_cluster", "cuda_intrinsics::barrier::fence_mbarrier_init_release_cluster", "cuda_intrinsics::barrier::fence_proxy_async_generic_acquire_shared_cluster_cluster", "cuda_intrinsics::barrier::fence_proxy_async_generic_release_shared_cta_cluster", @@ -1407,6 +1412,7 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::float::max_xorsign_abs_f32", "cuda_intrinsics::barrier::mbarrier_arrive", "cuda_intrinsics::barrier::mbarrier_arrive_cluster", + "cuda_intrinsics::barrier::mbarrier_arrive_cluster_relaxed", "cuda_intrinsics::barrier::mbarrier_arrive_expect_tx", "cuda_intrinsics::barrier::mbarrier_arrive_expect_tx_cluster", "cuda_intrinsics::barrier::mbarrier_arrive_no_complete", @@ -1416,6 +1422,7 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::barrier::mbarrier_try_wait", "cuda_intrinsics::barrier::mbarrier_try_wait_parity", "cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster", + "cuda_intrinsics::barrier::mbarrier_try_wait_parity_cluster_relaxed", "cuda_intrinsics::bf16::min_bf16", "cuda_intrinsics::bf16x2::min_bf16x2", "cuda_intrinsics::f16::min_f16", @@ -1571,6 +1578,7 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::matrix::mma_m16n8k32_s32_u8_s8_satfinite", "cuda_intrinsics::matrix::mma_m16n8k32_s32_u8_satfinite", "cuda_intrinsics::matrix::mma_m16n8k4_f32_tf32", + "cuda_intrinsics::matrix::mma_m16n8k64_mxf4_f32", "cuda_intrinsics::matrix::mma_m16n8k64_s32_s4", "cuda_intrinsics::matrix::mma_m16n8k64_s32_s4_satfinite", "cuda_intrinsics::matrix::mma_m16n8k64_s32_s4_u4", @@ -2340,6 +2348,7 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1016" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0819" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0820" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1042" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0313" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0315" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0314" @@ -2454,6 +2463,7 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0576" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0098" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0308" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1040" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0306" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0307" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1017" @@ -2463,6 +2473,7 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0309" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0310" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0311" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1041" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0838" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0067" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0830" @@ -2618,6 +2629,7 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0124" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0123" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0520" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1039" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0145" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0153" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0146" @@ -3388,6 +3400,7 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::float::ex2_approx_f16" | "cuda_device::float::ex2_approx_f32" | "cuda_device::float::ex2_approx_ftz_f32" + | "cuda_device::barrier::fence_acquire_shared_cluster_cluster" | "cuda_device::barrier::fence_mbarrier_init_release_cluster" | "cuda_device::barrier::fence_proxy_async_generic_acquire_shared_cluster_cluster" | "cuda_device::barrier::fence_proxy_async_generic_release_shared_cta_cluster" @@ -3503,6 +3516,7 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::float::max_xorsign_abs_f32" | "cuda_device::barrier::mbarrier_arrive" | "cuda_device::barrier::mbarrier_arrive_cluster" + | "cuda_device::barrier::mbarrier_arrive_cluster_relaxed" | "cuda_device::barrier::mbarrier_arrive_expect_tx" | "cuda_device::barrier::mbarrier_arrive_expect_tx_cluster" | "cuda_device::barrier::mbarrier_arrive_no_complete" @@ -3512,6 +3526,7 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::barrier::mbarrier_try_wait" | "cuda_device::barrier::mbarrier_try_wait_parity" | "cuda_device::barrier::mbarrier_try_wait_parity_cluster" + | "cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed" | "cuda_device::bf16::min_bf16" | "cuda_device::bf16x2::min_bf16x2" | "cuda_device::f16::min_f16" @@ -3667,6 +3682,7 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::wmma::mma_m16n8k32_s32_u8_s8_satfinite" | "cuda_device::wmma::mma_m16n8k32_s32_u8_satfinite" | "cuda_device::wmma::mma_m16n8k4_f32_tf32" + | "cuda_device::wmma::mma_m16n8k64_mxf4_f32" | "cuda_device::wmma::mma_m16n8k64_s32_s4" | "cuda_device::wmma::mma_m16n8k64_s32_s4_satfinite" | "cuda_device::wmma::mma_m16n8k64_s32_s4_u4" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs index 600b298617..8564f84dbe 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/lib.rs @@ -285,6 +285,9 @@ //! | `CUDA_OXIDE_DUMP_MIR` | Dump the `dialect-mir` module | //! | `CUDA_OXIDE_DUMP_LLVM` | Dump the LLVM dialect module | //! | `CUDA_OXIDE_PTX_DIR` | Override PTX output directory | +//! | `CUDA_OXIDE_MLIR_OUTPUT` | Write prepared CUTLASS MLIR | +//! | `CUDA_OXIDE_DEVICE_BACKEND` | Select `native` or `cutlass-mlir` | +//! | `CUDA_OXIDE_CUTLASS_OPT_LEVEL` | Official compiler level 2 or 3 (default 3) | //! | `CUDA_OXIDE_TARGET` | Override GPU target (e.g., `sm_90a`) | //! | `CUDA_OXIDE_DEVICE_CODEGEN_CRATE` | Filter device owner crate names | //! @@ -338,7 +341,7 @@ use rustc_middle::dep_graph::WorkProductMap; use rustc_middle::ty::TyCtxt; use rustc_middle::ty::print::with_no_trimmed_paths; use rustc_session::config::OutputFilenames; -use rustc_session::{CodegenBackendInit, EarlySession, IncrCompSession, Session}; +use rustc_session::{IncrCompSession, Session}; use std::any::Any; use std::collections::BTreeSet; use std::path::{Path, PathBuf}; @@ -375,7 +378,7 @@ static ARTIFACT_OBJECT_COUNTER: AtomicU64 = AtomicU64::new(0); /// /// All configuration is read from environment variables at backend load time. /// This avoids the need to thread configuration through rustc's argument parsing. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct CudaCodegenConfig { /// Print detailed compilation progress to stderr. pub verbose: bool, @@ -387,6 +390,10 @@ pub struct CudaCodegenConfig { pub dump_llvm_dialect: bool, /// Override PTX output directory (defaults to current directory). pub ptx_output_dir: Option, + /// Exact output path for the optional prepared CUTLASS MLIR module. + pub mlir_output: Option, + /// Parsed device backend selection, or the configuration error to report. + pub device_backend: Result, /// When set, emit device code only for these normalized local crate names. /// Host code still goes through the wrapped LLVM backend for every crate. pub device_codegen_crates: Option>, @@ -402,7 +409,11 @@ impl CudaCodegenConfig { /// | `CUDA_OXIDE_DUMP_MIR` | `dump_mir_dialect` | /// | `CUDA_OXIDE_DUMP_LLVM` | `dump_llvm_dialect` | /// | `CUDA_OXIDE_PTX_DIR` | `ptx_output_dir` | - /// | `CUDA_OXIDE_DEVICE_CODEGEN_CRATE` | `device_codegen_crates` | + /// | `CUDA_OXIDE_MLIR_OUTPUT` | `mlir_output` | + /// | `CUDA_OXIDE_DEVICE_BACKEND` | `device_backend` | + /// | `CUDA_OXIDE_CUTLASS_COMPILER` | `device_backend` | + /// | `CUDA_OXIDE_MLIR_PROFILE` | `device_backend` | + /// | `CUDA_OXIDE_DEVICE_CODEGEN_CRATE` | `device_codegen_crates` | pub fn from_env() -> Self { Self { verbose: std::env::var("CUDA_OXIDE_VERBOSE").is_ok(), @@ -412,6 +423,8 @@ impl CudaCodegenConfig { ptx_output_dir: std::env::var("CUDA_OXIDE_PTX_DIR") .ok() .map(std::path::PathBuf::from), + mlir_output: std::env::var_os("CUDA_OXIDE_MLIR_OUTPUT").map(std::path::PathBuf::from), + device_backend: device_backend_from_env(), device_codegen_crates: parse_device_codegen_crates( std::env::var(reserved_oxide_symbols::DEVICE_CODEGEN_CRATE_ENV) .ok() @@ -427,6 +440,93 @@ impl CudaCodegenConfig { } } +impl Default for CudaCodegenConfig { + fn default() -> Self { + Self { + verbose: false, + dump_rustc_mir: false, + dump_mir_dialect: false, + dump_llvm_dialect: false, + ptx_output_dir: None, + mlir_output: None, + device_backend: Ok(mir_importer::DeviceBackend::Native), + device_codegen_crates: None, + } + } +} + +fn device_backend_from_env() -> Result { + let selector = match std::env::var("CUDA_OXIDE_DEVICE_BACKEND") { + Ok(value) => Some(value), + Err(std::env::VarError::NotPresent) => None, + Err(std::env::VarError::NotUnicode(_)) => { + return Err("CUDA_OXIDE_DEVICE_BACKEND is not valid Unicode".to_owned()); + } + }; + let profile = match std::env::var("CUDA_OXIDE_MLIR_PROFILE") { + Ok(value) => Some(value), + Err(std::env::VarError::NotPresent) => None, + Err(std::env::VarError::NotUnicode(_)) => { + return Err("CUDA_OXIDE_MLIR_PROFILE is not valid Unicode".to_owned()); + } + }; + let mut backend = parse_device_backend( + selector.as_deref(), + std::env::var_os("CUDA_OXIDE_CUTLASS_COMPILER").map(std::path::PathBuf::from), + profile.as_deref(), + )?; + if let mir_importer::DeviceBackend::CutlassMlir(config) = &mut backend { + match std::env::var("CUDA_OXIDE_CUTLASS_OPT_LEVEL") { + Ok(value) => { + config.optimization_level = value + .parse() + .map_err(|_| "CUDA_OXIDE_CUTLASS_OPT_LEVEL must be 2 or 3")?; + if !matches!(config.optimization_level, 2 | 3) { + return Err("CUDA_OXIDE_CUTLASS_OPT_LEVEL must be 2 or 3".into()); + } + } + Err(std::env::VarError::NotPresent) => {} + Err(std::env::VarError::NotUnicode(_)) => { + return Err("CUDA_OXIDE_CUTLASS_OPT_LEVEL is not valid Unicode".into()); + } + } + } + Ok(backend) +} + +fn parse_device_backend( + selector: Option<&str>, + compiler_library: Option, + profile: Option<&str>, +) -> Result { + match selector.map(str::trim).filter(|value| !value.is_empty()) { + None | Some("native") => Ok(mir_importer::DeviceBackend::Native), + Some("cutlass-mlir") => { + let compiler_library = compiler_library + .filter(|path| !path.as_os_str().is_empty()) + .ok_or_else(|| { + "CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir requires CUDA_OXIDE_CUTLASS_COMPILER=" + .to_owned() + })?; + let mut config = mir_importer::CutlassBackendConfig::new(compiler_library); + if let Some(profile) = profile { + config.profile = profile.to_owned(); + } + Ok(mir_importer::DeviceBackend::CutlassMlir(config)) + } + Some(value) => Err(format!( + "unknown CUDA_OXIDE_DEVICE_BACKEND value {value:?}; expected `native` or `cutlass-mlir`. CUDA_OXIDE_BACKEND is intentionally not used here because cargo-oxide reserves it for the rustc backend library path" + )), + } +} + +fn device_backend_conflicts_with_materialization( + backend: &mir_importer::DeviceBackend, + materialization_requested: bool, +) -> bool { + materialization_requested && matches!(backend, mir_importer::DeviceBackend::CutlassMlir(_)) +} + fn normalize_device_crate_name(name: &str) -> String { name.trim().replace('-', "_") } @@ -462,12 +562,12 @@ impl CodegenBackend for CudaCodegenBackend { "cuda" } - fn init(&mut self, sess: &EarlySession) -> CodegenBackendInit { + fn init(&self, sess: &Session) { // Note: Don't log here - init() is called for ALL crates including dependencies. // We log in codegen_crate() only when there are kernels to compile. // Initialize the underlying LLVM backend - self.llvm_backend.init(sess) + self.llvm_backend.init(sess); } fn print_version(&self) { @@ -482,7 +582,7 @@ impl CodegenBackend for CudaCodegenBackend { self.llvm_backend.target_cpu(sess) } - fn target_config(&self, sess: &EarlySession) -> rustc_codegen_ssa::TargetConfig { + fn target_config(&self, sess: &Session) -> rustc_codegen_ssa::TargetConfig { self.llvm_backend.target_config(sess) } @@ -523,6 +623,11 @@ impl CodegenBackend for CudaCodegenBackend { // This is necessary because we use tcx.def_path_str() and other functions that // trigger trimmed_def_paths. rust-gpu uses the same pattern. with_no_trimmed_paths!({ + let device_backend = self.config.device_backend.clone().unwrap_or_else(|error| { + tcx.dcx().fatal(format!( + "[rustc_codegen_cuda] Invalid device backend configuration: {error}" + )) + }); // Step 1: Analyze for device code let mono_partitions = tcx.collect_and_partition_mono_items(()); let kernel_count = collector::count_kernels_in_cgus(tcx, mono_partitions.codegen_units); @@ -616,6 +721,14 @@ impl CodegenBackend for CudaCodegenBackend { "[rustc_codegen_cuda] Invalid cubin materialization request: {error}" )) }); + if device_backend_conflicts_with_materialization( + &device_backend, + materialization_request.is_some(), + ) { + tcx.dcx().fatal( + "[rustc_codegen_cuda] CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir already produces a final cubin and cannot be combined with cargo oxide --materialize-cubin; choose exactly one cubin producer", + ); + } if self.config.verbose { eprintln!("[rustc_codegen_cuda] Compiling device code via cuda-oxide..."); } @@ -659,6 +772,8 @@ impl CodegenBackend for CudaCodegenBackend { std::env::current_dir().unwrap_or_else(|_| ".".into()) }), output_name: tcx.crate_name(rustc_hir::def_id::LOCAL_CRATE).to_string(), + mlir_output: self.config.mlir_output.clone(), + device_backend: device_backend.clone(), verbose: self.config.verbose, dump_rustc_mir: self.config.dump_rustc_mir, dump_mir_dialect: self.config.dump_mir_dialect, @@ -851,6 +966,7 @@ impl CodegenBackend for CudaCodegenBackend { assembly: None, llvm_ir: None, global_asm_object: None, + links_from_incr_cache: Vec::new(), }); } (compiled_modules, work_products) @@ -1274,6 +1390,76 @@ mod tests { assert!(both.contains("32-bit"), "{both}"); } + #[test] + fn device_backend_selector_keeps_native_as_the_explicit_default() { + assert_eq!( + parse_device_backend(None, None, None).unwrap(), + mir_importer::DeviceBackend::Native + ); + assert_eq!( + parse_device_backend(Some("native"), None, None).unwrap(), + mir_importer::DeviceBackend::Native + ); + } + + #[test] + fn cutlass_selector_requires_the_official_compiler_library() { + let error = parse_device_backend(Some("cutlass-mlir"), None, None).unwrap_err(); + assert!(error.contains("CUDA_OXIDE_CUTLASS_COMPILER")); + + let backend = parse_device_backend( + Some("cutlass-mlir"), + Some("/opt/cutlass/lib/libCutlassCompiler.so".into()), + None, + ) + .unwrap(); + assert_eq!( + backend, + mir_importer::DeviceBackend::CutlassMlir(mir_importer::CutlassBackendConfig::new( + "/opt/cutlass/lib/libCutlassCompiler.so".into(), + )) + ); + } + + #[test] + fn device_backend_selector_rejects_unknown_values() { + let error = parse_device_backend(Some("unknown"), None, None).unwrap_err(); + assert!(error.contains("expected `native` or `cutlass-mlir`")); + assert!(error.contains("CUDA_OXIDE_BACKEND")); + } + + #[test] + fn device_backend_preserves_explicit_primitives_profile() { + let backend = parse_device_backend( + Some("cutlass-mlir"), + Some("/opt/cutlass/lib/libCutlassCompiler.so".into()), + Some("cutlass-primitives-mlir22"), + ) + .unwrap(); + assert!( + matches!(backend, mir_importer::DeviceBackend::CutlassMlir(c) + if c.profile == "cutlass-primitives-mlir22") + ); + } + + #[test] + fn cutlass_backend_conflicts_with_a_second_cubin_producer() { + let native = mir_importer::DeviceBackend::Native; + let cutlass = + mir_importer::DeviceBackend::CutlassMlir(mir_importer::CutlassBackendConfig::new( + "/opt/cutlass/lib/libCutlassCompiler.so".into(), + )); + assert!(!device_backend_conflicts_with_materialization( + &native, true + )); + assert!(!device_backend_conflicts_with_materialization( + &cutlass, false + )); + assert!(device_backend_conflicts_with_materialization( + &cutlass, true + )); + } + #[test] fn device_codegen_owner_filter_normalizes_and_matches_crate_names() { let owners = parse_device_codegen_crates(Some(" gpu-kernels,math_gpu , gpu-kernels ")) diff --git a/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.lock new file mode 100644 index 0000000000..dd4f819e21 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.lock @@ -0,0 +1,802 @@ +# 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", + "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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[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 = "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.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn 2.0.119", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[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.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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 = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[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 = "primitives_api_compile" +version = "0.0.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "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.6", +] + +[[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 = "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.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +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.6", +] + +[[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.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" + +[[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.59" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6df92bf3d9227be3d53173901ddbffac2babc27ae50f397776ffd6dc33f800cb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.59" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac4f328cf2f05d084e496c3e9c3f33ed0a183656a16e1fcec4d464d8373aec82" +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/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.toml new file mode 100644 index 0000000000..ce7eeef4df --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "primitives_api_compile" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +ctm-rs = { path = "../../../../../cute-rs-preview/ctm-rs" } +cuda-device = { path = "../../../../cuda-device" } +cuda-host = { path = "../../../../cuda-host" } +cuda-core = { path = "../../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/kernel.rs b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/kernel.rs new file mode 100644 index 0000000000..4c10a774d9 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/kernel.rs @@ -0,0 +1,250 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use ctm_rs::cuda_module; + +// Keep this boundary present: concrete shared pointers must adapt to a +// generic Rust helper ABI without changing pointee, mutability or pointer kind. +#[inline(never)] +fn generic_shared_identity(pointer: *mut u8) -> *mut u8 { + pointer +} + +#[cuda_module] +mod probes { + use ctm_rs::{ + async_copy, barrier, + cluster::{self, CtaMask}, + convert, fence, float, + groups::CtaGroup2, + kernel, launch_bounds, launch_contract, + mbarrier::{self, Barrier, BarrierAddressExt, ManagedBarrier, Ready, TmaBarrier, Uninit}, + memory::{ + DisjointSlice, DynamicSharedArray, F16Tile128, F16TileB80, F16TileB96, F16TileC, + SharedRegion, + }, + tcgen05::{ + self, + descriptors::*, + mma::{self, collector_a}, + tmem::{self, ColumnOffset, TmemGuard, TmemReady, TmemUninit}, + }, + thread, + tma::{self, TensorMap}, + warp, + }; + use cuda_device::{SharedArray, cluster_launch}; + + static mut SLOT: SharedArray = SharedArray::UNINIT; + static mut BAR: Barrier = Barrier::UNINIT; + static mut A: SharedArray = SharedArray::UNINIT; + static mut B0: SharedArray = SharedArray::UNINIT; + static mut B1: SharedArray = SharedArray::UNINIT; + static mut C: SharedArray = SharedArray::UNINIT; + + #[kernel] + pub fn primitives_grid_constant(#[grid_constant] map: &TensorMap) { + unsafe { tma::prefetch_tensormap(map) }; + } + + // These instruction probes are compile-only. Their straight-line issue + // sequence intentionally does not implement a launchable collective kernel. + #[kernel] + #[cluster_launch(2, 1, 1)] + pub fn primitives_lifecycle() { + unsafe { + let slot = SharedArray::as_raw_mut_ptr(&raw mut SLOT); + let guard = TmemGuard::::from_raw_slot(slot); + let guard = guard.alloc_collective(); + tcgen05::relinquish_alloc_permit_cg2(); + let _dead = guard.dealloc_collective(); + let bar = + ManagedBarrier::::from_static(&raw mut BAR).init_instruction(1); + let _dead = bar.invalidate_instruction(); + } + } + + #[kernel] + #[cluster_launch(2, 1, 1)] + pub fn primitives_tma( + #[grid_constant] a_map: &TensorMap, + #[grid_constant] b0_map: &TensorMap, + #[grid_constant] b1_map: &TensorMap, + #[grid_constant] c_map: &TensorMap, + out: *mut u64, + ) { + unsafe { + let Ok(a) = SharedRegion::::from_array(&*(&raw const A)) else { + return; + }; + let Ok(b0) = SharedRegion::::from_array(&*(&raw const B0)) else { + return; + }; + let Ok(b1) = SharedRegion::::from_array(&*(&raw const B1)) else { + return; + }; + let Ok(c) = SharedRegion::::from_array(&*(&raw const C)) else { + return; + }; + let bar = ManagedBarrier::::adopt_ready(&raw mut BAR); + let completion = cluster::mapa(bar.shared_address(), 0); + let Ok(mask) = CtaMask::::try_from_bits(0b11) else { + return; + }; + let coords = tma::Coordinates2D::new(-32, 4096); + tma::cp_async_bulk_tensor_shared_cluster_global(&a, a_map, coords, completion, mask); + tma::cp_async_bulk_tensor_shared_cluster_global(&b0, b0_map, coords, completion, mask); + tma::cp_async_bulk_tensor_shared_cluster_global(&b1, b1_map, coords, completion, mask); + tma::cp_async_bulk_tensor_global_shared_cta( + &c, + c_map, + tma::Coordinates2D::new(64, 128), + ); + tma::cp_async_bulk_commit_group(); + tma::cp_async_bulk_wait_group_read::<0>(); + tma::cp_async_bulk_wait_group_read::<7>(); + out.write(mbarrier::arrive(&bar).raw()); + out.add(1) + .write(mbarrier::arrive_expect_tx(&bar, 16384).raw()); + out.add(2).write(mbarrier::try_wait_parity(&bar, 1) as u64); + out.add(3) + .write(mbarrier::arrive_expect_tx_cluster(&bar, 77824).raw()); + out.add(4) + .write(mbarrier::try_wait_parity_cluster(&bar, 0) as u64); + cluster::arrive_remote(cluster::mapa(bar.shared_address(), 1)); + tcgen05::commit_multicast_cg2(bar.shared_address(), mask); + } + } + + #[kernel] + #[cluster_launch(2, 1, 1)] + pub fn primitives_sync(out: *mut u64) { + unsafe { + let shared = DynamicSharedArray::::get_raw(); + out.add(3) + .write(super::generic_shared_identity(shared) as u64); + fence::fence_mbarrier_init(); + fence::fence_proxy_async_shared_cta(); + tcgen05::fence_before_thread_sync(); + tcgen05::fence_after_thread_sync(); + barrier::barrier_cta_sync(3, 128); + barrier::barrier_cta_sync_aligned(3, 128); + barrier::barrier_cta_arrive(3, 128); + barrier::barrier_cta_arrive_aligned(3, 128); + cluster::barrier_cluster_arrive(); + cluster::barrier_cluster_wait(); + cluster::barrier_cluster_arrive_aligned(); + cluster::barrier_cluster_wait_aligned(); + cluster::barrier_cluster_arrive_relaxed(); + cluster::barrier_cluster_arrive_relaxed_aligned(); + let (lane, elected) = warp::elect_sync(u32::MAX); + out.write(lane as u64); + out.add(1).write(elected as u64); + out.add(2).write( + (warp::warp_id() + + thread::blockIdx_x() + + cluster::cluster_ctarank() + + cluster::cluster_size()) as u64, + ); + } + } + + // Distinct constants expose operand permutations and exact math modifiers + // in the exported IR; these probes are never launched. + #[kernel] + #[launch_bounds(32)] + #[launch_contract(domain = 1, coordinates = u32, block = (32, 1, 1))] + pub fn primitives_attention(mut out: DisjointSlice, packed: *mut u32) { + use ctm_rs::mma::{Accumulator16x8, F16A, F16B, mma_f16}; + unsafe { + async_copy::cp_async_commit_group(); + async_copy::cp_async_wait_group::<0>(); + async_copy::cp_async_wait_group::<7>(); + let out = out.as_mut_ptr(); + out.write(float::ex2_approx_f32(1.25)); + out.add(1).write(float::ex2_approx_ftz_f32(2.5)); + out.add(2).write(float::rcp_approx_ftz_f32(4.0)); + out.add(3).write(warp::shuffle_xor_f32(8.0, 2)); + packed.write(convert::cvt_f16x2_f32(1.25, 2.5)); + let c = Accumulator16x8::from_registers([1.0, 2.0, 3.0, 4.0]); + let a = F16A::from_registers([0x3c003c00, 0x40004000, 0x42004200, 0x44004400]); + let b = F16B::from_registers([0x46004600, 0x48004800]); + let registers = mma_f16(c, a, b).registers(); + out.add(4).write(registers[0]); + out.add(5).write(registers[1]); + out.add(6).write(registers[2]); + out.add(7).write(registers[3]); + } + } + + #[kernel] + #[cluster_launch(2, 1, 1)] + pub fn primitives_mma(out: *mut f32) { + type Left = DenseSharedF16; + type Right = DenseSharedF16; + unsafe { + let Ok(a) = SharedRegion::::from_array(&*(&raw const A)) else { + return; + }; + let Ok(b0) = SharedRegion::::from_array(&*(&raw const B0)) else { + return; + }; + let Ok(b1) = SharedRegion::::from_array(&*(&raw const B1)) else { + return; + }; + let Ok(a) = SmemOperand::::from_region(&a) else { + return; + }; + let Ok(b0) = SmemOperand::::from_region(&b0) else { + return; + }; + let Ok(b1) = SmemOperand::::from_region(&b1) else { + return; + }; + let guard = TmemGuard::::adopt_ready( + SharedArray::as_raw_mut_ptr(&raw mut SLOT), + ); + let Ok(left) = guard.view::<::DstFormat>(ColumnOffset::new(0)) else { + return; + }; + let Ok(right) = guard.view::<::DstFormat>(ColumnOffset::new(192)) + else { + return; + }; + let l = InstructionDescriptor::::new(); + let r = InstructionDescriptor::::new(); + mma::mma::(&left, &a.k16::<0>(), &b0.k16::<0>(), &l, false); + mma::mma::(&left, &a.k16::<1>(), &b0.k16::<1>(), &l, true); + mma::mma::(&left, &a.k16::<2>(), &b0.k16::<2>(), &l, true); + mma::mma::(&left, &a.k16::<3>(), &b0.k16::<3>(), &l, true); + mma::mma::(&right, &a.k16::<0>(), &b1.k16::<0>(), &r, false); + mma::mma::(&right, &a.k16::<1>(), &b1.k16::<1>(), &r, true); + mma::mma::( + &right, + &a.k16::<2>(), + &b1.k16::<2>(), + &r, + true, + ); + mma::mma::( + &right, + &a.k16::<3>(), + &b1.k16::<3>(), + &r, + true, + ); + let Ok(read) = guard.read_32x32(96, ColumnOffset::new(192)) else { + return; + }; + let raw = mma::load_32x32b_x32(&read); + tmem::load_wait(); + let regs = raw.decode_f32(); + macro_rules! store { ($($i:literal),*) => { $(out.add($i).write(regs.get::<$i>());)* }; } + store!( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, + 23, 24, 25, 26, 27, 28, 29, 30, 31 + ); + let c = SharedArray::as_raw_mut_ptr(&raw mut C); + c.write(regs.get::<0>() as f16); + } + } +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/main.rs new file mode 100644 index 0000000000..54c26e02bd --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/src/main.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![feature(f16)] +mod kernel; + +fn main() { + // Compile-only instruction/ABI probes. These kernels are never launched. +} diff --git a/cuda-oxide/cuda-oxide-book/advanced/cluster-programming.md b/cuda-oxide/cuda-oxide-book/advanced/cluster-programming.md index 9ac70e4d2d..c2d8671e1a 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/cluster-programming.md +++ b/cuda-oxide/cuda-oxide-book/advanced/cluster-programming.md @@ -83,12 +83,15 @@ let ncls = cluster::num_clusters(); // total clusters in the grid block's shared memory to read. Think of it as a "block-local device ID" within the cluster. -For the full 3D cluster position: +For the full 3D position of the block in its cluster, and of the cluster in +the grid: ```rust let cx = cluster::cluster_ctaidX(); // block's X position in the cluster let cy = cluster::cluster_ctaidY(); // block's Y position in the cluster let cz = cluster::cluster_ctaidZ(); // block's Z position in the cluster +let gx = cluster::cluster_idx_x(); // cluster's X position in the grid +let nx = cluster::cluster_grid_dim_x(); // clusters along X in the grid ``` --- diff --git a/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg b/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg index fc3a56b807..9f509220ab 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg +++ b/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/advanced/images/shared-memory-alignment.svg b/cuda-oxide/cuda-oxide-book/advanced/images/shared-memory-alignment.svg index 761f93ba5d..a66298136b 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/images/shared-memory-alignment.svg +++ b/cuda-oxide/cuda-oxide-book/advanced/images/shared-memory-alignment.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.svg b/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.svg deleted file mode 100644 index acbd534d5a..0000000000 --- a/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.svg +++ /dev/null @@ -1,139 +0,0 @@ - - - - - -CTA-local TMA loads: a shared pointer that is shared by type, and the .shared::cta copy form -The cluster form converts a generic *mut u8 to the cluster window; on sm_120a ptxas adds a runtime CTA check whose call drops setmaxnreg. The CTA form takes a SharedPtr<u8>. -TODAY (cp_async_bulk_tensor_2d_g2s, generic *mut u8) -THE TYPE (SharedPtr<T>, from #1341) -NEW (cp_async_bulk_tensor_2d_g2s_cta, SharedPtr<u8>) - -tma_copy_2d_test: the cluster form -static mut TILE: SharedArray<f32, TILE_SIZE, 128> - = SharedArray::UNINIT; -static mut BAR: Barrier = Barrier::UNINIT; - -// dst: a generic *mut u8, any CTA of the cluster -cp_async_bulk_tensor_2d_g2s( - &raw mut TILE as *mut u8, - tensor_map, tile_x, tile_y, - &raw mut BAR, -); -Unchanged by this PR: 39 in-tree callers keep this form and -its cluster contract; the PTX is byte for byte the same. - -tma_copy.ll: the generic pointer is converted -%v21 = addrspacecast ptr addrspace(3) @TILE to ptr -%v22 = addrspacecast ptr %v21 to ptr addrspace(7) -call void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d( - ptr addrspace(7) %v22, ptr addrspace(3) @BAR, ptr %v7, - i32 %v9, i32 %v10, i16 0, i64 0, i1 0, i1 0, i32 0) -The destination arrives as a generic pointer (the user's cast), -so the lowering can only name the cluster window, addrspace(7). -On sm_100a the cluster form is native. On sm_120a ptxas cannot -prove a non-multicast .shared::cluster destination local, so -it inserts a runtime CTA check with a syscall fallback, and -that call makes it ignore setmaxnreg (C7506): the #1238 kernel -ran with 24 registers instead of 168. - -SharedPtr<T>: addrspace(3) by type -let t = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; -let d = DynamicSharedArray::<u8, 128>::shared_ptr_at(1024); -let p = unsafe { SharedPtr::from_raw(generic) }; // asserts -t.add(i) t.byte_add(b) t.cast::<u8>() // keep the type -t.as_ptr() // the generic *mut T for loads and stores -t.addr_u32() // the 32-bit .shared::cta address -repr(transparent) over *mut SharedSlot<T>; the importer types -that pointee as shared, next to SharedArray and Barrier, and -the newtype lowers as one scalar, like ThreadIndex. - -shared_memory_helper: across helper boundaries -struct Region<T> { base: SharedPtr<T>, offset: usize } -struct Pair { a: SharedPtr<f32>, b: SharedPtr<f32> } -#[inline(never)] -unsafe fn shared_ptr_probe<T>(region: Region<T>) -> u64 -#[inline(never)] -unsafe fn write_pair(pair: Pair, index: usize, value: f32) -#[inline(never)] -unsafe fn select_slot(pair: Pair, index: usize, second: bool) - -> SharedPtr<f32> - -the .ll: what each boundary keeps -define i64 @..shared_ptr_probe..( - ptr addrspace(3) %v0, i64 %v1) -define void @..write_pair( - ptr addrspace(3) %v0, ptr addrspace(3) %v1, i64, float) -define ptr addrspace(3) @..select_slot( - ptr addrspace(3) %v0, ptr addrspace(3) %v1, i64, i1) -Both routes; 0 cvta.to.shared in the PTX. A SharedPtr kernel -parameter, bare or in a struct or array, is a compile error. - -tma_copy_2d_cta_test: the CTA form (new) -// dst: SharedPtr<u8>, this CTA's shared memory by type -cp_async_bulk_tensor_2d_g2s_cta( - SharedArray::shared_ptr(&raw mut TILE).cast::<u8>(), - tensor_map, tile_x, tile_y, - &raw mut BAR, -); -cp_async_bulk_tensor_{1..5}d_g2s_cta, catalog records -i1034 to i1038, PTX 8.6, sm_90 and newer; tile mode only. - -LLVM NVPTX route: typed intrinsic (LLVM 22+) -call void @llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.2d( - ptr addrspace(3) @TILE, ptr addrspace(3) @BAR, ptr %v7, - i32 %v9, i32 %v10, i64 0, i1 0) -No addrspacecast: the destination is already addrspace(3). - -libNVVM route: inline PTX, u32 addresses -asm sideeffect - "cp.async.bulk.tensor.2d.shared::cta.global.tile - .mbarrier::complete_tx::bytes [$0], [$2, {$3, $4}], [$1];", - "r,r,l,r,r,~{memory}" - -refused at compile time -a SharedPtr kernel parameter, bare or inside a struct or array -a barrier passed straight from map_shared_rank_mut (a located -error that names the cluster form to use instead) -the caller's contract: from_raw on a cluster pointer; the raw -cuda_intrinsics layer keeps *mut u8 and converts with -cvta.to.shared, and its safety text says CTA-local - - -THE PTX AND WHAT PTXAS MAKES OF IT ON sm_120a (tma_copy.ptx built for sm_90a; mnemonic column aligned; the copy line wrapped before .mbarrier; register counts from ptxas 13.4) - -cluster form: tma_copy_2d_test (unchanged) -mov.b64 %rd15, TILE; -cvta.shared.u64 %rd16, %rd15; -cvta.to.shared::cluster.u64 %rd17, %rd16; -cp.async.bulk.tensor.2d.shared::cluster.global.tile - .mbarrier::complete_tx::bytes [%rd17], [%rd2, {%r3, %r4}], [BAR]; -// ptxas sm_120a: 24 registers, 8-byte stack; with setmaxnreg: C7506 ignored -// SASS: CALL.ABS.NOINC, then UTMALDG.2D; no USETMAXREG - -CTA form: tma_copy_2d_cta_test -// no address conversion: the shared symbols are the operands -cp.async.bulk.tensor.2d.shared::cta.global.tile - .mbarrier::complete_tx::bytes [TILE], [%rd2, {%r3, %r4}], [BAR]; -// ptxas sm_120a: 10 registers, no stack; with setmaxnreg: 168 registers -// no C7506; SASS: USETMAXREG.DEALLOC.CTAPOOL 0x28, UTMALDG.2D, -// USETMAXREG.TRY_ALLOC.CTAPOOL 0xe8; no CALL -What this proves: the cluster form converts the generic destination twice (to generic, then to the -cluster window) and ptxas on sm_120a adds the runtime CTA check whose call drops setmaxnreg. The CTA -form names the shared symbols directly and keeps register reallocation. On the 5090, tma_copy copies -4096 of 4096 values through the new kernel next to the unchanged ones, through both routes, with -memcheck and synccheck clean; ranks 1 to 5 are probed and assembled, and the cluster PTX is unchanged. - diff --git a/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md b/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md index 5871909e1a..9a0c9f011f 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md +++ b/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md @@ -270,69 +270,6 @@ the strongest typed requirement over the kernel's `get`, `get_raw` and The declaration aligns the pool base only; keeping an `offset(byte_offset)` result aligned for `T` is still the caller's job. -### Typed shared pointers: `SharedPtr` - -`DynamicSharedArray::get()` and `SharedArray::as_mut_ptr()` return `*mut T`, -which in device code is a *generic* pointer: it may address global, local or -shared memory, including another block's shared memory in a cluster. Once -such a pointer is passed to a helper function, the compiler no longer knows it -points into this block's shared memory. - -`SharedPtr` is a raw pointer into the issuing block's shared memory -(`.shared::cta`) that carries that fact in its type. It lowers to an -`addrspace(3)` pointer wherever it goes, including helper parameters and -return values, struct fields and generic code, so instructions that take it -use their shared-memory forms without converting a generic address. - -```rust -use cuda_device::{DynamicSharedArray, SharedArray, SharedPtr, kernel, thread}; - -#[inline(never)] -unsafe fn fill(tile: SharedPtr, i: usize, value: f32) { - // Still a shared-memory pointer here, although `fill` is not inlined. - unsafe { tile.add(i).as_ptr().write(value) }; -} - -static mut TILE: SharedArray = SharedArray::UNINIT; - -#[kernel] -pub fn fill_tile() { - // SAFETY: TILE is a `static mut SharedArray` of this kernel. - let tile: SharedPtr = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; - // The same type for dynamic shared memory, 1024 bytes into the pool. - let pool: SharedPtr = DynamicSharedArray::::shared_ptr_at(1024); - // SAFETY: the launch supplies more than 1024 dynamic shared bytes, and - // each thread writes its own slot of TILE. - unsafe { - fill(tile, thread::threadIdx_x() as usize, 1.0); - pool.as_ptr().write(0); - } -} -``` - -| Operation | Result | -|:---------------------------------------|:----------------------------------------------------| -| `unsafe SharedArray::shared_ptr(&raw mut S)` | `SharedPtr` to the first element of `S`, a `static mut` of this kernel | -| `DynamicSharedArray::::shared_ptr()` / `shared_ptr_at(bytes)` | Like `get()` / `offset(bytes)` | -| `unsafe SharedPtr::from_raw(p)` | Asserts that generic `p` points into this block's shared memory | -| `add`, `offset`, `byte_add`, `cast` | Pointer arithmetic that keeps the type | -| `as_ptr()` | The generic `*mut T`, for loads, stores and raw-pointer APIs | -| `addr_u32()` | The 32-bit `.shared::cta` address, for `_shared_u32` forms | - -`SharedPtr` cannot be a kernel parameter, bare or inside a struct or array: -the host has no shared-memory address to pass, and the compiler refuses the -signature. Pointers from `cluster::map_shared_rank` may name another block's -shared memory and are not `SharedPtr`s. - -Inside aggregates the usual shared-pointer limits apply. A union or a -`MaybeUninit>` fails closed under the union storage rules, and a -`repr(packed)` struct that holds one follows the packed-carrier limits; both -are listed in [Supported Features](../appendix/supported-features.md). On the -libNVVM route a shared pointer is 4 bytes wide, so a struct that mixes a -`SharedPtr` with a narrower field, such as `{ p: SharedPtr, x: u32 }`, is -not yet supported in arrays or in shared memory. That gap predates -`SharedPtr` and applies to `*mut Barrier` fields in the same way. - --- ## Synchronization: sync_threads() diff --git a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md index f2397747bf..d885b6c131 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md +++ b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md @@ -123,17 +123,14 @@ These functions address cluster shared memory: the destination may belong to any CTA of the cluster, so the compiler converts it to the cluster window. When the destination is the issuing CTA's own shared memory, `cp_async_bulk_tensor_*_g2s_cta` emits the `.shared::cta` destination form -(PTX 8.6, sm_90+) without that conversion. Its destination is a -[`SharedPtr`](shared-memory-and-synchronization.md), -a CTA-shared pointer by type, so the locality holds wherever the pointer came -from, including through helper functions: +(PTX 8.6, sm_90+): ```rust // SAFETY: TILE and BAR are `static mut` shared-memory objects of this kernel, // and the copy completes on BAR before any thread reads TILE. unsafe { cp_async_bulk_tensor_2d_g2s_cta( - SharedArray::shared_ptr(&raw mut TILE).cast::(), + &raw mut TILE as *mut u8, tensor_map, tile_x, tile_y, &raw mut BAR, ); @@ -141,19 +138,8 @@ unsafe { ``` On SM12x this form matters for performance: ptxas implements a non-multicast -`.shared::cluster` TMA load it cannot prove CTA-local with a runtime check and -a system call, which also makes it ignore `setmaxnreg`. - -```{figure} images/tma-cta-local.svg -:align: center -:width: 100% - -The two G2S destination forms. Left, the cluster form converts a generic `*mut u8` to the cluster -window, and on sm_120a ptxas adds a runtime CTA check whose call makes it ignore `setmaxnreg`. -Middle, `SharedPtr` keeps the shared address space in its type across helper parameters, struct -fields and return values. Right, `cp_async_bulk_tensor_2d_g2s_cta` takes a `SharedPtr` and emits -the `.shared::cta` form with no address conversion, so the kernel keeps `setmaxnreg` and its 168 registers. -``` +`.shared::cluster` TMA load with a runtime CTA-id check and a system call +fallback, which also makes it ignore `setmaxnreg`. ### Only one thread issues the copy diff --git a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md index cc90fe89a5..4214f0fba1 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md +++ b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md @@ -312,7 +312,6 @@ pub fn dynamic(data: &[f32]) { | `SharedArray` | Compile-time sized, block-scoped shared memory | | `SharedArray` | With 128-byte alignment (required for TMA destinations) | | `DynamicSharedArray` | Runtime-sized shared memory (set via `LaunchConfig`) | -| `SharedPtr` | Pointer into this block's shared memory, `addrspace(3)` by type | Both are `!Sync` — concurrent access requires explicit barriers. @@ -451,7 +450,7 @@ cp_async_bulk_commit_group(); | Function | Description | |:----------------------------------------------|:-------------------------------------| | `cp_async_bulk_tensor_{1..5}d_g2s(...)` | Global → shared async bulk copy | -| `cp_async_bulk_tensor_{1..5}d_g2s_cta(...)` | Global → this CTA's `SharedPtr` | +| `cp_async_bulk_tensor_{1..5}d_g2s_cta(...)` | Global → this CTA's shared memory | | `cp_async_bulk_tensor_{1..5}d_s2g(...)` | Shared → global async bulk copy | | `cp_async_bulk_tensor_2d_g2s_multicast(...)` | Multicast to all CTAs in cluster | | `cp_async_bulk_commit_group()` | Commit outstanding copies | @@ -524,7 +523,7 @@ N_COLS must be a power of 2 in the range [32, 512]. ### Typed Synchronous ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; let ctx = CudaContext::new(0).unwrap(); @@ -545,7 +544,7 @@ unsafe { ### Typed Async ```rust -use cuda_async::simt::device_operation::DeviceOperation; +use cuda_async::device_operation::DeviceOperation; let module = kernels::load_async(0)?; // SAFETY: this is 1D, buffers contain n elements, and module/scheduler share a context. @@ -651,7 +650,7 @@ debug::prof_trigger::<7>(); // Nsight profiler trigger | `thread` | Thread/block IDs, `index_1d`, `sync_threads` | All | | `config` | Compile-time policies, shapes, tiles, atoms, layouts, memory spaces, and scopes | All | | `disjoint` | `DisjointSlice` — typed writes completed by a launch proof | All | -| `shared` | `SharedArray`, `DynamicSharedArray`, `SharedPtr` | All | +| `shared` | `SharedArray`, `DynamicSharedArray` | All | | `warp` | Shuffle, vote, match, lane/warp ID | All | | `atomic` | Scoped atomics (device/block/system) | sm_70+ | | `debug` | `clock64`, `trap`, `breakpoint`, `gpu_printf!` | All | diff --git a/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md b/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md index 56500c646a..f3049e37b5 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md +++ b/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md @@ -10,7 +10,7 @@ from a fresh checkout. If you just want to run an example, the | Dependency | Version | Purpose | |:-----------------|:----------------------------- |:------------------------------------------------------------| -| **Rust nightly** | `nightly-2026-10-11` (pinned) | Compiler toolchain with `rustc-dev` for the codegen backend | +| **Rust** | `1.99.1` (pinned stable) | Compiler toolchain with `rustc-dev` for the codegen backend | | **CUDA Toolkit** | 13.0+ (with cuRAND headers) | Driver API, `nvcc`, PTX assembler; `curand.h` for bindgen | | **Clang** | 21+ (`clang-21` pkg) | `bindgen` in host `cuda-bindings` needs clang's headers | | **Linux** | Tested on Ubuntu 24.04 | Windows and macOS are not supported | @@ -26,22 +26,21 @@ cd cuda-rust ## Install the Rust toolchain The CUDA Oxide component ships `cuda-oxide/rust-toolchain.toml`, which pins the -exact nightly version and components. Rustup picks it up automatically for -commands run inside `cuda-oxide/`; the repository root and `cutile-rs/` stay on -stable: +exact stable release and components. Rustup picks it up automatically for +commands run inside `cuda-oxide/`: ```toml # cuda-oxide/rust-toolchain.toml [toolchain] -channel = "nightly-2026-10-11" +channel = "1.99.1" components = ["rust-src", "rustc-dev", "rust-analyzer", "clippy", "rustfmt", "llvm-tools"] ``` If you need to install manually: ```bash -rustup toolchain install nightly-2026-10-11 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-10-11 +rustup toolchain install 1.99.1 +rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain 1.99.1 ``` `rust-src` provides the standard library source for cross-compilation, @@ -66,7 +65,7 @@ required for `ptxas` and header files, but you will not be able to run kernels. ## Install LLVM (usually optional) The codegen pipeline emits LLVM IR and invokes `llc` to produce PTX. The -pinned Rust toolchain (`nightly-2026-10-11`) already ships LLVM 23 with the +pinned Rust toolchain (`1.99.1`) already ships LLVM 23 with the NVPTX backend enabled via the `llvm-tools` component, so the recommended path is: @@ -139,7 +138,7 @@ crates (`cuda-device`, `cuda-host`, `cuda-macros`, etc.) and the build tooling (`cargo-oxide`). The host runtime (`cuda-bindings`, `cuda-core`, `cuda-async`) is shared with cutile-rs and lives at the repository root; the SIMT workspace depends on those crates by path. This section, "Verify the installation", and -"Common commands" run from `cuda-oxide/`, where the nightly pin and the +"Common commands" run from `cuda-oxide/`, where the SIMT toolchain pin and the `cargo oxide` workspace alias apply: ```bash @@ -149,17 +148,18 @@ cargo build ```{note} The codegen backend (`cuda-oxide/crates/rustc-codegen-cuda/`) is intentionally **not** a -workspace member because it requires special nightly features and a different -build process. `cargo-oxide` handles building it transparently. +workspace member because it links against `rustc_private` (via `rustc-dev` and +`RUSTC_BOOTSTRAP=1`) and uses a separate build process. `cargo-oxide` handles +building it transparently. ``` ## Install cargo-oxide `cargo-oxide` is the cargo subcommand that drives the full compilation -pipeline. Inside the repo, it works via a workspace alias. For standalone use, install it with the pinned nightly toolchain: +pipeline. Inside the repo, it works via a workspace alias. For standalone use, install it with the pinned toolchain: ```bash -cargo +nightly-2026-10-11 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide +cargo +1.99.1 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide ``` On first run, `cargo-oxide` automatically fetches and builds the codegen backend @@ -249,8 +249,8 @@ repository/ ├── cuda-async/ # Shared async runtime ├── cutile-rs/ # Tile DSL (nested import) └── cuda-oxide/ - ├── Cargo.toml # Nightly SIMT workspace - ├── rust-toolchain.toml # Pinned nightly + components + ├── Cargo.toml # SIMT workspace + ├── rust-toolchain.toml # Pinned stable toolchain (1.99.1) + components ├── crates/ │ ├── cuda-device/ # Device intrinsics (#![no_std]) │ ├── cuda-host/ # Host launch APIs @@ -287,7 +287,7 @@ repository/ `error[E0463]: can't find crate for rustc_middle` : You are missing the `rustc-dev` component. Run: - `rustup component add rustc-dev --toolchain nightly-2026-10-11`. + `rustup component add rustc-dev --toolchain 1.99.1`. CUDA driver version mismatch : The toolkit version (compile-time) and driver version (runtime) must be @@ -296,5 +296,5 @@ CUDA driver version mismatch `cargo oxide doctor` fails on codegen backend : The backend is built on first use. If the build fails, check that - `rust-src` is installed and that the nightly version matches - `rust-toolchain.toml`. + `rust-src` and `rustc-dev` are installed and that the active toolchain + matches `rust-toolchain.toml`. diff --git a/cuda-oxide/cuda-oxide-book/appendix/glossary.md b/cuda-oxide/cuda-oxide-book/appendix/glossary.md index 8c66a44cec..60a223afa5 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/glossary.md +++ b/cuda-oxide/cuda-oxide-book/appendix/glossary.md @@ -39,7 +39,7 @@ alongside the normal host binary. ## `cuda-async` The async execution layer, a crate shared with cutile-rs and maintained at -the repository root (`cuda-async/`). The cuda-oxide model lives under `cuda_async::simt`: +the repository root (`cuda-async/`). Tile and SIMT share the same types: `DeviceOperation` (lazy GPU work description), `DeviceFuture` (stream-bound execution), and `DeviceBox` (device-owned memory). Compose work with `zip!`, `and_then`, and `value()`. @@ -55,8 +55,9 @@ cores, atomics, and debug facilities. Safe RAII wrappers around the CUDA Driver API: `CudaContext`, `CudaStream`, `DeviceBuffer`, and module loading. Handles GPU context and memory management on the host side. Shared with cutile-rs and maintained at the repository root -(`cuda-core/`); the cuda-oxide surface is re-exported at the crate root except -`simt::LaunchConfig`, `simt::vmm`, `simt::peer`, and `simt::memory`. +(`cuda-core/`). Both programming models use its shared `Context`, `Stream`, +`LaunchConfig` and typed allocation storage. A `Device` identifies the physical GPU; +a `Context` identifies an execution context on that GPU. ## `DeviceOperation` diff --git a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md index 71a5df7b38..8f233ac77a 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md +++ b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md @@ -117,13 +117,6 @@ do not count toward that bound. Taking the in-place address of such a field (`u.ptr = p`, `&u.ptr`, `&mut u.ptr`) is refused with a diagnostic; build the union by value and read its fields by value. -`SharedPtr`, the typed shared pointer, follows the same limits: a union or a -`MaybeUninit>` fails closed under these union storage rules, a -`repr(packed)` struct that holds one follows the packed-carrier limits above, a -kernel parameter that holds one bare or inside a struct or array is refused, -and on the libNVVM route a struct that mixes one with a narrower field is not -yet supported in arrays or in shared memory. - Enum constants preserve payload relocations to device statics, including non-zero byte addends. This includes niche-encoded `Option<&T>` and direct-tagged enum layouts, both for direct thin-reference payloads and for diff --git a/cuda-oxide/cuda-oxide-book/async-programming/combinators-and-composition.md b/cuda-oxide/cuda-oxide-book/async-programming/combinators-and-composition.md index 53b0049677..a84d7ad2b5 100644 --- a/cuda-oxide/cuda-oxide-book/async-programming/combinators-and-composition.md +++ b/cuda-oxide/cuda-oxide-book/async-programming/combinators-and-composition.md @@ -254,7 +254,7 @@ every combinator working together. This function returns a single `DeviceOperation` describing the entire forward pass for one batch: ```rust -use cuda_async::simt::launch::{AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; +use cuda_async::launch::{AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; fn forward_pass( batch: Vec, diff --git a/cuda-oxide/cuda-oxide-book/async-programming/concurrent-execution.md b/cuda-oxide/cuda-oxide-book/async-programming/concurrent-execution.md index dd02e654f7..5ef2ddf26d 100644 --- a/cuda-oxide/cuda-oxide-book/async-programming/concurrent-execution.md +++ b/cuda-oxide/cuda-oxide-book/async-programming/concurrent-execution.md @@ -59,7 +59,7 @@ sets up the scheduling policy (round-robin with four streams by default), and makes the thread-local state available for `.sync()` and `.await`: ```rust -use cuda_async::simt::device_context::init_device_contexts; +use cuda_async::device_context::init_device_contexts; #[tokio::main] async fn main() -> Result<(), Box> { diff --git a/cuda-oxide/cuda-oxide-book/async-programming/scheduling-and-streams.md b/cuda-oxide/cuda-oxide-book/async-programming/scheduling-and-streams.md index 76e079dda1..0a770910a8 100644 --- a/cuda-oxide/cuda-oxide-book/async-programming/scheduling-and-streams.md +++ b/cuda-oxide/cuda-oxide-book/async-programming/scheduling-and-streams.md @@ -109,12 +109,9 @@ it -- the policy just works. ## `SingleStream` -- one lane, strict order -:::{note} -`SingleStream` is implemented in the scheduling internals but is **not -currently wired into `GlobalSchedulingPolicy`** or exposed through -`init_device_contexts`. The default setup always uses `StreamPoolRoundRobin`. -This section describes the design intent for a future API surface. -::: +Use `SingleStream::from_stream(stream.clone())` with `.schedule(&policy)` to +run Tile and SIMT operations on the same stream. The default global policy +continues to use `StreamPoolRoundRobin`. For debugging, or when you need guaranteed ordering across *all* operations, `SingleStream` routes everything to one stream. Every operation sees the results @@ -140,7 +137,7 @@ Before any async operation can run, you initialize the thread-local device context: ```rust -use cuda_async::simt::device_context::init_device_contexts; +use cuda_async::device_context::init_device_contexts; init_device_contexts(0, 1)?; ``` diff --git a/cuda-oxide/cuda-oxide-book/async-programming/the-device-operation-model.md b/cuda-oxide/cuda-oxide-book/async-programming/the-device-operation-model.md index 1a6d6e14db..df5bfa9138 100644 --- a/cuda-oxide/cuda-oxide-book/async-programming/the-device-operation-model.md +++ b/cuda-oxide/cuda-oxide-book/async-programming/the-device-operation-model.md @@ -74,8 +74,8 @@ method. It looks like the sync method, but without the stream argument, and it returns a recipe instead of cooking immediately: ```rust -use cuda_async::simt::device_context::init_device_contexts; -use cuda_core::simt::LaunchConfig; +use cuda_async::device_context::init_device_contexts; +use cuda_core::LaunchConfig; // One-time setup: create a stream pool for scheduling init_device_contexts(0, 1)?; @@ -177,7 +177,7 @@ Operations never create streams themselves. The scheduling policy (covered in separation: operations describe *what*, the context provides *where*. Inside an `execute` implementation, you access the stream with -`ctx.get_cuda_stream()` and the CUDA context with `ctx.get_cuda_context()`. +`ctx.get_cuda_stream()` and the CUDA context with `ctx.context()`. For most operations this is all you need -- enqueue a kernel or a memory copy on the stream, and you are done. @@ -265,7 +265,7 @@ weight vector -- into a chain of device operations. The `value()` function wraps any `Send` type in a no-op `DeviceOperation` that returns it immediately: ```rust -use cuda_async::simt::device_operation::value; +use cuda_async::device_operation::value; let weights = vec![1.0f32; 1024]; let op = value(weights); // impl DeviceOperation> @@ -297,8 +297,8 @@ The stream is assigned later, by the scheduling policy. until the `ExecutionContext` is available: ```rust -use cuda_async::simt::device_operation::{value, with_context}; -use cuda_core::simt::memory::{malloc_async, memcpy_htod_async}; +use cuda_async::device_operation::{value, with_context}; +use cuda_core::memory::{malloc_async, memcpy_htod_async}; fn h2d(host_data: Vec) -> impl DeviceOperation> { with_context(move |ctx| { @@ -309,7 +309,7 @@ fn h2d(host_data: Vec) -> impl DeviceOperation> { let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); memcpy_htod_async(dptr, host_data.as_ptr(), num_bytes, stream.cu_stream()) .unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) + value(DeviceBox::from_raw_parts(dptr, n, stream.clone())) } }) } diff --git a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md index 32fcf73dbc..04d44fd57b 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md +++ b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md @@ -129,18 +129,23 @@ cuda-oxide is split into focused crates. Here is every one and its role: | `rustc-codegen-cuda` | Custom rustc codegen backend -- intercepts `codegen_crate()`, splits host/device code | | `mir-importer` | Translates Stable MIR into `dialect-mir`, orchestrates the full pipeline | | `dialect-mir` | pliron dialect modeling Rust MIR semantics (places, rvalues, terminators) | +| `dialect-cute` | Backend-neutral semantic IR for CuTe tensor programs | | `mir-transforms` | Analyzes and optimizes `dialect-mir` before lowering; currently provides loop unrolling | | `nvvm-transforms` | Converts lowered LLVM operations to the form accepted by the selected NVVM dialect | | `llvm-export` | Re-exports `pliron-llvm`'s LLVM dialect + cuda-oxide's textual `.ll` exporter | +| `cuda-oxide-mlir-export` | CUTLASS 4.7 mappings and profiles; upstream MLIR printing | | `dialect-nvvm` | pliron dialect for NVIDIA GPU intrinsics (`tid`, `ntid`, barriers, TMA) | | `mir-lower` | Lowers `dialect-mir` to the LLVM dialect -- the main transformation pass | | `cargo-oxide` | CLI tool: `cargo oxide build`, `cargo oxide run`, `cargo oxide pipeline` | | `cuda-device` | Device-side API: intrinsics, `DisjointSlice`, barriers, shared memory, warp ops | +| `cute-layout` | CuTe-style layout algebra and coordinate mapping | +| `ctm-rs` | Typed instruction-level CUDA primitives | +| `cute-rs` | Typed CuTe tensors, tiles, copies, pipelines, and MMA operations | | `cuda-macros` | Proc macros: `#[kernel]`, `#[device]` | | `cuda-host` | Host-side typed module loading and launch helpers | -| `cuda-core` | Shared with cutile-rs: safe Driver API bindings; SIMT surface under `cuda_core::simt` | +| `cuda-core` | Shared with cutile-rs: safe Driver API bindings; shared context, stream and buffer types | | `cuda-core-derive` | Shared with cutile-rs: `#[derive(DeviceCopy)]` proc-macro for `cuda-core` | -| `cuda-async` | Shared with cutile-rs: `DeviceOperation` model and combinators under `cuda_async::simt` | +| `cuda-async` | Shared with cutile-rs: `DeviceOperation` model and combinators shared by Tile and SIMT | | `cuda-bindings` | Shared with cutile-rs: FFI to the CUDA driver; `libcuda` is loaded at run time | | `cuda-intrinsics` | Generated low-level CUDA intrinsic declarations | | `cuda-intrinsics-gen`| Extractor and deterministic source generator for the intrinsics | @@ -153,7 +158,7 @@ cuda-oxide is split into focused crates. Here is every one and its role: | `dialect-ptx` | pliron dialect modelling structured PTX | | `ptx-parse` | Lossless structural views over PTX source text | | `ptx-schedule` | PTX schedule-perturbation fuzzing (nanosleep injection campaigns) | -| `cuda-oxide-codegen` | Rustc-independent PTX backend | +| `cuda-oxide-codegen` | Rustc-independent PTX and CUTLASS cubin backend | | `libnvvm-sys` | `dlopen` bindings to libNVVM (used by `cuda-host::ltoir`) | | `cuda-target-spec` | Shared CUDA target parsing and recorded LLVM PTX-floor policy | | `nvjitlink-sys` | `dlopen` bindings to nvJitLink (used by `cuda-host::ltoir`) | diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/dev-workflow-shift.svg b/cuda-oxide/cuda-oxide-book/compiler/images/dev-workflow-shift.svg new file mode 100644 index 0000000000..1414728e81 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/dev-workflow-shift.svg @@ -0,0 +1,107 @@ + + + +Track 1: Shifting cuda-oxide development from a daily nightly pin to an official stable Rust release + + + + + + + + +Track 1: Moving Development & Repo Toolchains to Stable Rust +Pins an official stable Rust release (`1.99.1`) in `cuda-oxide/rust-toolchain.toml` while keeping `rustc_private` and `rustc_public`. + + +BEFORE · DAILY NIGHTLY PIN IN CUDA-OXIDE + +Nightly Pin in cuda-oxide/rust-toolchain.toml + +Root + cutile-rs: channel = "1.98.0" (Stable) +cuda-oxide/: channel = "nightly-2026-10-11" +Result: daily snapshot drift and mid-cycle rustc churn. + + + + +Why Nightly Was Used Originally + +• Needed #![feature(rustc_private)] for rustc-codegen-cuda +• Needed rustc_public (stable_mir) for mir-importer +• Needed -Z codegen-backend and -Z mir-enable-passes + + +NEW WORKFLOW · PINNED STABLE RELEASE (1.99.1) + + +Step 1: Pin Stable Toolchain + + +channel = "1.99.1" # Stable! +components = ["rustc-dev", ...] +Ships LLVM 23.1.1 with NVPTX. +librustc_public.rlib included! + + + + +Step 2: cargo-oxide Sets Env + + +cmd.env("RUSTC_BOOTSTRAP", "1") +-Z codegen-backend=... +Works for both building the .so +and compiling kernel crates. + + + + +Step 3: 6-Week Release Cadence + + +1.99.1 -> 1.100.0 -> 1.101.0 +Vetted through 6 weeks of beta +No daily nightly regressions. +Same model as Microsoft UTC. + diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/drop-proof-termination.svg b/cuda-oxide/cuda-oxide-book/compiler/images/drop-proof-termination.svg index c4e27ddac8..4d7976bf36 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/images/drop-proof-termination.svg +++ b/cuda-oxide/cuda-oxide-book/compiler/images/drop-proof-termination.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/ecosystem-and-bootstrap.svg b/cuda-oxide/cuda-oxide-book/compiler/images/ecosystem-and-bootstrap.svg new file mode 100644 index 0000000000..6586c2ffe9 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/ecosystem-and-bootstrap.svg @@ -0,0 +1,110 @@ + + + +How custom Rust backends ship in production and why stable rustc already contains everything cuda-oxide needs + + + + + + + + +The Ecosystem Playbook & Why Stable rustc Already Has Everything We Need +Official stable Rust releases ship the exact same backend hooks and rustc_public crates as nightly behind RUSTC_BOOTSTRAP=1. + + +1 · WHAT OTHER RUST COMPILER / BACKEND PROJECTS DO + + +A. Reimplement Frontend (rust-analyzer, cutile-rs) +• Zero rustc code; custom parser + type solver on stable. +• Works for IDEs, not codegen (no borrowck, mono, or LLVM). + + +B. Out-of-Tree Nightly Pin (rust-gpu, Kani, Charon) +• Uses #![feature(rustc_private)] on a daily nightly pin. +• Pain point: users need nightly; daily builds can hit ICEs. + + +C. In-Tree Upstream (rustc_codegen_cranelift / gcc) +• Lives inside rust-lang/rust; built by upstream Rust CI. +• Only possible if upstream Rust owns your target backend. + + +D. Stable-Tag SDK Bundle (MS codegen_utc, Ferrocene) +• Uses the exact same CodegenBackend trait we use! +• Builds against a vetted STABLE release (e.g. 1.90 / 1.99) + and ships rustc + .so + sysroot as one standalone bundle. + + +2 · WHAT IS ACTUALLY INSIDE AN OFFICIAL STABLE RUST RELEASE + + +Official rustup Stable Toolchain (e.g. 1.99.1) + + + +Startup Gatekeeper in rustc +if channel == "nightly" || RUSTC_BOOTSTRAP == "1" + + +Unlocks on Stable: + + + +1. Unstable Compiler Flags (-Z) +-Z codegen-backend=... -Z mir-enable-passes=... + + +2. rustc-dev Component (Ships on Stable rustup!) +librustc_driver.so + librustc_public.rlib +(Verified: rustc_public & bridge come for free!) + + +3. Scoped Feature Gate (-Z allow-features=...) +Unlocks rustc_private for us; user code stays stable. + diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/enum-alignment-anchor.svg b/cuda-oxide/cuda-oxide-book/compiler/images/enum-alignment-anchor.svg index 947f4ff900..d335217568 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/images/enum-alignment-anchor.svg +++ b/cuda-oxide/cuda-oxide-book/compiler/images/enum-alignment-anchor.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/toolkit-bundle-and-sizes.svg b/cuda-oxide/cuda-oxide-book/compiler/images/toolkit-bundle-and-sizes.svg new file mode 100644 index 0000000000..166767f0b3 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/toolkit-bundle-and-sizes.svg @@ -0,0 +1,113 @@ + + + +Track 2: Self-contained CUDA Toolkit bundle, measured binary sizes, and bug-fix workflow + + + + + + + + +Track 2: Bundling Everything for Users & the CUDA Toolkit (Option B) +Users never install rustc-dev or compile the backend. Everything ships prebuilt in one ~120 MB compressed package. + + +1 · MEASURED SIZE BREAKDOWN (x86_64-unknown-linux-gnu) + + +Stays in NVIDIA CI Only (Not Shipped): ~1.15 GB + +• rustc-dev (librustc_*.rlib compiler libs): ~930 MB +• llvm-tools & unstripped debug artifacts: ~220 MB + + +CI builds & strips bundle + + +Shipped in CUDA Toolkit / Standalone Bundle + + +librustc_codegen_cuda.so +34 MB (8.6 MB zst) + + +librustc_driver.so (frontend) +159 MB (~42 MB zst) + + +libLLVM.so (host LLVM backend) +199 MB (~52 MB zst) + + +Pre-encoded sysroot + bins +61 MB (~17 MB zst) + +Total Bundle: ~453 MB on disk | ~120 MB compressed + + +2 · ZERO-CONFIG USER FLOW & BUG OWNERSHIP + + +User Experience: Zero rustup or Nightly Setup + +$ cargo oxide build +1. cargo-oxide points RUSTC at bundled cuda-oxide-rustc +2. Uses bundled sysroot (built with -Z always-encode-mir) +3. Emits host binary + embedded PTX/CUBIN in one pass + + +How We Fix Bugs Without Maintaining a rustc Fork + + +Bucket 1 (~95% of bugs): In cuda-oxide crates +Fix in mir-importer/lower -> ship 8.6 MB .so hotfix only. + + +Bucket 2 (~5% of bugs): Upstream rustc passes +Toggle pass flag in driver (-JumpThreading) or bump tag. + diff --git a/cuda-oxide/cuda-oxide-book/compiler/images/upstream-rustc-public-flow.svg b/cuda-oxide/cuda-oxide-book/compiler/images/upstream-rustc-public-flow.svg new file mode 100644 index 0000000000..fbcf1855be --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/upstream-rustc-public-flow.svg @@ -0,0 +1,104 @@ + + + +Upstreaming rustc_public PRs on a stable release cycle without waiting or emergency nightly bumps + + + + + + + + +Upstreaming to rustc_public Without Waiting for the 6-Week Stable Train +Because rustc_public wraps existing rustc_middle data, facts.rs lets us ship on Day 1 and switch to rustc_public on schedule. + + +1 · OLD NIGHTLY WORKFLOW (EXAMPLE: PR #1482 + PR #1483) + + +1. Upstream PR to Rust + +rust-lang/rust#163242 +Exposes ArgAttributes getters + + + + +2. Forced Nightly Bump (PR #1482) + +nightly-08-28 -> nightly-10-11 +1,506 files changed due to unrelated churn! + + + + +3. Land Feature (PR #1483) + +facts.rs calls attrs.regular() +Blocked until step 2 finished + + +2 · NEW STABLE WORKFLOW (ZERO WAIT, ZERO EMERGENCY BUMPS) + + +Day 1: Ship Feature Immediately on Current Stable + + +// Inside facts.rs (quarantine boundary): +let internal = rustc_internal::internal(tcx, arg); +• Underlying rustc_abi data ALREADY exists in current stable! +• Upstream rust#163242 in parallel; 0 toolchain churn. + + +6 wks + + +Week 6: Scheduled Stable Release Upgrade + + +// Swap 5 lines in facts.rs to rustc_public: +let flags = attrs.regular(); +• rust#163242 arrives automatically in the next stable tag. +• Delete the 5-line shim in facts.rs; rest is unchanged. + diff --git a/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md b/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md index e227a9e59c..668013f72a 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md +++ b/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md @@ -587,7 +587,7 @@ that cuda-oxide emits. | 5th | `llc-21` on `PATH` | Distro / `apt.llvm.org` install of LLVM 21. | | 6th | `llc` on `PATH` | Reporting fallback only; rejected at runtime if older than LLVM 21. | -The pinned Rust toolchain (`nightly-2026-10-11`) ships LLVM 23 with NVPTX +The pinned Rust toolchain (`1.99.1`) ships LLVM 23 with NVPTX enabled, so `rustup component add llvm-tools` is the recommended onboarding path. The PATH probes for `llc-23` / `llc-22` / `llc-21` are kept as a fallback for users with an existing LLVM install. If none of the probes succeed the diff --git a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md index 19986fe39e..6225db2323 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md +++ b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md @@ -303,16 +303,16 @@ they become `call` instructions to `@llvm.nvvm.*` intrinsics. ### Architecture Coverage -At catalog SHA-256 `3f56e2a1` (the stamp in every `ops/generated/` file -header), the dialect holds 585 operations across 42 modules, and they come +At catalog SHA-256 `63f8d327` (the stamp in every `ops/generated/` file +header), the dialect holds 589 operations across 42 modules, and they come from two different places. The split is the first thing to know about it, because it decides where -- and whether -- you would add one. If the header -stamp no longer starts with `3f56e2a1`, the counts on this page predate the +stamp no longer starts with `63f8d327`, the counts on this page predate the catalog you are reading. **Hand-written**, directly under `cuda-oxide/crates/dialect-nvvm/src/ops/`. These are the ops with bespoke verification or lowering that the intrinsic catalog does not -describe. There are seven modules and 26 operations: +describe. There are seven modules and 27 operations: | Module | Description | Ops | | :-------- | :---------------------------------------------------------- | --: | @@ -321,13 +321,13 @@ describe. There are seven modules and 26 operations: | `cluster` | Cluster index and cluster-count registers | 2 | | `debug` | `assertfail`, `vprintf` | 2 | | `grid` | Cooperative `grid_sync` | 1 | -| `memory` | Generic-to-shared address conversion with a byte offset | 1 | +| `memory` | Generic-to-shared address conversion and global stores | 2 | | `wgmma` | Warpgroup MMA descriptors; bf16/f16 at m64n64k16, bf16 at m64n128k16, tf32 at m64n64k8 | 14 | **Generated**, under `ops/generated/`, from `intrinsics/catalog.json` by `cuda-intrinsics-gen`. Every file there opens with `// @generated ... DO NOT EDIT.`, and editing one by hand is undone by the next generator run. This is -the large majority -- 35 modules and 559 operations, resolved from 1038 catalog +the large majority -- 35 modules and 562 operations, resolved from 1042 catalog entries, since several intrinsics can share one structural op: | Area | Modules | Ops | @@ -336,7 +336,7 @@ entries, since several intrinsics can share one structural op: | Tensor Memory Accelerator | `tma` | 116 | | Special registers | `sreg` | 44 | | Packed (SIMD-in-register) | `packed_alu`, `packed_conversion`, `packed_atomic` | 55 | -| Async copy and barriers | `cp_async`, `mbarrier_extended`, `mbarrier_basic`, `sync` | 35 | +| Async copy and barriers | `cp_async`, `mbarrier_extended`, `mbarrier_basic`, `sync` | 38 | | Warp-level | `warp_shuffle`, `redux`, `vote`, `warp_match`, `warp_barrier`, `active_mask`, `elect` | 39 | | Matrix fragment movement | `ldmatrix`, `register_mma`, `stmatrix`, `wgmma_control`, `movmatrix`, `sparse_mma` | 22 | | Execution and debug control | `execution_control`, `debug_control` | 11 | diff --git a/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md b/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md index b157eda3d1..d42ac59af3 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md +++ b/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md @@ -277,42 +277,96 @@ already explicit. The `mir-importer` crate translates this into Pliron IR (an MLIR-like framework), and from there the [lowering pipeline](lowering-pipeline.md) takes it the rest of the way to PTX. -## Nightly pinning +## Stabilizing the toolchain pin -Even with `rustc_public` providing a stable *API*, the bridge layer -(`rustc_public_bridge`) is still compiled against the compiler's internals -and is not independently versioned. In practice, this means a specific -version of cuda-oxide works with a specific nightly. - -cuda-oxide pins to an exact nightly release via `rust-toolchain.toml`: +Even with `rustc_public` providing a stable API, the bridge layer +(`rustc_public_bridge`) is compiled against `librustc_driver.so` and is tied to +an exact compiler build. Instead of chasing daily `nightly` snapshots, +cuda-oxide pins an official **stable** Rust release via +`rust-toolchain.toml`: ```toml [toolchain] -channel = "nightly-2026-10-11" +channel = "1.99.1" components = ["rust-src", "rustc-dev", "rust-analyzer", "clippy", "rustfmt", "llvm-tools"] ``` -This pin guarantees reproducible builds: anyone who clones the repository -gets the same compiler, the same MIR shapes, and the same bridge behavior. -When updating the pin, the process is: - -1. Bump the nightly date in `cuda-oxide/rust-toolchain.toml`. Rustup walks up - from the standalone `rustc-codegen-cuda` workspace and finds that component - pin, so no nested copy is needed. `cuda-oxide/scripts/check-toolchain-parity.sh` - checks the scaffold `cargo oxide new` writes, the devcontainer, and the - blocks and dated commands quoted in this book and the READMEs -- run it to - find the stragglers. -2. Fix any `rustc_public` API changes (usually minor -- that is the whole - point of the stable API). -3. Run the full test suite to verify that all examples still compile and - produce correct PTX. -4. Celebrate, or revert and try next week's nightly. +When bumping this pin, run `cuda-oxide/scripts/check-toolchain-parity.sh` to +verify that `scaffold.rs`, `.devcontainer/devcontainer.json`, the book, and the +READMEs agree, and run `nix flake update rust-overlay` in `cuda-oxide/` so +`cuda-oxide/flake.lock` carries the new release manifest. -```{note} -As `rustc_public` matures and moves toward a crates.io release, the coupling -to a specific nightly will loosen. The long-term goal is for cuda-oxide to -work with any sufficiently recent stable Rust toolchain -- but we are not -there yet. +```{figure} images/ecosystem-and-bootstrap.svg +:align: center +:width: 100% + +How custom Rust compiler backends ship in production, and what lives inside an official stable rustup release. +``` + +### Why stable rustc already has rustc_private and rustc_public + +Official stable rustc binaries are compiled from the same source tree as +nightly; nothing is stripped out of the binary. In addition, rustup publishes +the `rustc-dev` component (`librustc_driver.so`, `librustc_public.rlib`, +`librustc_public_bridge.rlib`, and `librustc_middle.rlib`) for every stable +release. + +When cargo-oxide invokes Cargo (and inside the workspace via +`cuda-oxide/.cargo/config.toml`), it sets `RUSTC_BOOTSTRAP=1` in the process +environment. That unlocks `#![feature(rustc_private)]` and `-Z` flags +(`-Z codegen-backend=...`, `-Z mir-enable-passes=-JumpThreading`, +`-Z always-encode-mir`) on the pinned stable compiler while giving developers +and users a compiler that has already completed Rust's 12-week nightly and beta +stabilization cycle. + +```{figure} images/dev-workflow-shift.svg +:align: center +:width: 100% + +Pinning an official stable release (`1.99.1`) in `cuda-oxide/rust-toolchain.toml` with `RUSTC_BOOTSTRAP=1` managed by cargo-oxide. +``` + +### Upstreaming to rustc_public without emergency toolchain bumps + +When a feature requires a new `rustc_public` query that we upstream to +`rust-lang/rust` (such as `ArgAttributes` getters in +`rust-lang/rust#163242`), it takes 6 to 12 weeks for that commit to reach the +next stable Rust release. + +Because `rustc_public` only wraps data that already exists inside +`rustc_middle` and `rustc_abi`, `crates/mir-importer/src/translator/facts.rs` +acts as our quarantine boundary: + +1. **Day 1 (on the current stable pin):** Land the feature immediately by + reading the underlying `rustc_middle` / `rustc_target` field inside + `facts.rs` via `rustc_internal::internal(tcx, ...)`, while upstreaming the + `rustc_public` getter in parallel. +2. **Next scheduled stable upgrade:** Once the new stable Rust release ships + with the upstreamed `rustc_public` API, replace the temporary internal read + in `facts.rs` with the `rustc_public` method call. + +```{figure} images/upstream-rustc-public-flow.svg +:align: center +:width: 100% + +Decoupling feature delivery from toolchain bumps using `facts.rs` and scheduled stable upgrades. +``` + +### Self-contained bundling for standalone and CUDA Toolkit distribution + +For standalone or CUDA Toolkit packaging (following the same model as +Microsoft's `rustc_codegen_utc` and Ferrocene), users do not need `rustc-dev` +(~930 MB of compiler-internal `.rlib`s used only when building the backend +itself). A stripped runtime bundle containing `rustc`, `cargo`, `cargo-oxide`, +`librustc_codegen_cuda.so` (34 MB stripped, 8.6 MB compressed), +`librustc_driver.so`, `libLLVM.so`, and a pre-encoded `core`/`alloc`/`std` +sysroot totals roughly 453 MB on disk (~120 MB compressed). + +```{figure} images/toolkit-bundle-and-sizes.svg +:align: center +:width: 100% + +Measured bundle footprint on `x86_64-unknown-linux-gnu` and the two-bucket maintenance workflow. ``` --- diff --git a/cuda-oxide/cuda-oxide-book/cute/build-with-cutlass.md b/cuda-oxide/cuda-oxide-book/cute/build-with-cutlass.md new file mode 100644 index 0000000000..22a6d66365 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/build-with-cutlass.md @@ -0,0 +1,66 @@ +# Build with the CUTLASS Compiler + +CuTe kernels compile through NVIDIA's CUTLASS 4.7 compiler library, not through the +native LLVM path. `cargo oxide` installs the pinned release once, checks its digest on +every load, and embeds the resulting cubin in your binary. + +```text +Rust kernel + ──► one Pliron module: dialect-cute, dialect-mir, dialect-nvvm + ──► whole-module verification and the target gate + ──► CUTLASS MLIR text + ──► libCutlassCompiler 4.7.0 + ──► cubin embedded in the host binary +``` + +## Install once + +```bash +cargo oxide toolchain install cutlass +``` + +- Downloads the CUTLASS 4.7.0 release archive from GitHub and verifies the archive and + library SHA-256 digests. The bundle is under the NVIDIA EULA; the installer prints + the notice. +- Linux x86_64 with CUDA 13 only. The library lands in Cargo's data directory under + `cuda-oxide/toolchains/cutlass/4.7.0/x86_64-cu13`. +- `CUDA_OXIDE_CUTLASS_COMPILER=/abs/path/libCutlassCompiler.so` overrides discovery. + The digest check still applies, so only the byte-identical 4.7.0 library loads. + +## Build and run + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide run elementwise_cute --arch sm_120a +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build fp16_gemm_256x352_cute --arch sm_100a +``` + +- Every CuTe kernel needs `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir`. The native backend + has no CuTe lowering and stops with a verification error. +- Two targets: `sm_100a` and `sm_120a`. Anything else is rejected before compilation + and the error names the target. `tcgen05` operations need `sm_100a`; MXF4 register + MMA needs `sm_120a`. +- `cargo oxide run` without `--arch` uses the detected GPU, so an RTX 5090 gets + `sm_120a`. `cargo oxide build` needs `--arch`. +- Each example README states its device and shape requirements. The + [support matrix](support-matrix.md) lists what the two targets cover. + +## Options + +| Variable | Purpose | +| --- | --- | +| `CUDA_OXIDE_DEVICE_BACKEND` | `native` (default) or `cutlass-mlir`; other values fail | +| `CUDA_OXIDE_MLIR_PROFILE` | `cutlass-full-cute-mlir22` (default) or `cutlass-primitives-mlir22` for kernels that use `ctm-rs` alone; part of the build fingerprint | +| `CUDA_OXIDE_CUTLASS_OPT_LEVEL` | `2` or `3` (default); the attention examples compare at `2`; part of the build fingerprint | +| `CUDA_OXIDE_MLIR_OUTPUT` | Path for the prepared MLIR text, for example `$PWD/module.mlir` | +| `CUDA_OXIDE_CUTLASS_COMPILER` | Absolute path to a `libCutlassCompiler.so` with the pinned digest | + +`CUDA_OXIDE_BACKEND` selects the rustc backend library, not the device path. Keep MLIR +dumps and cubins out of version control. + +## Failures + +- Unsupported operations, rejected targets, verification errors and compiler errors + stop the build. There is no fallback to another backend. +- The compiler checks layouts, operand roles, pipeline state and the launch ABI. It + does not prove thread convergence, barrier phases, asynchronous completion or + pointer lifetimes; see [Checks and Obligations](contracts.md). diff --git a/cuda-oxide/cuda-oxide-book/cute/contracts.md b/cuda-oxide/cuda-oxide-book/cute/contracts.md new file mode 100644 index 0000000000..45e07fe5d3 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/contracts.md @@ -0,0 +1,65 @@ +# Checks and Obligations + +The compiler verifies the whole module before any MLIR is built. Types carry the +static facts, the verifier checks how operations connect, and everything dynamic stays +with the author at `unsafe` boundaries. + +```{figure} images/ctm-rs-at-a-glance.svg +:align: center +:width: 100% + +Where ctm-rs sits: typed carriers select instruction families, one of two compiler profiles consumes them, and the author keeps the schedule. +``` + +## What is checked + +| Layer | Checks | +| :--- | :--- | +| Rust types | layout, swizzle, operand role, element format, MMA shape, pipeline stage count, barrier ids; `TensorMut` over `DisjointSlice`; `TmemGuard` typestate; sealed traits | +| cute-layout proofs | cooperative copy coverage and per-thread contiguity, `ldmatrix` source validity, TMA encodability and swizzle phase alignment | +| dialect-cute verifier | tensor provenance, one pipeline init with complete producer and consumer state, cursor provenance, scheduler grids, MMA placement, epilogue phases, no compiler-only type reaches runtime storage | +| Target gate | `sm_100a` or `sm_120a`; `tcgen05` only on `sm_100a`; MXF4 register MMA only on `sm_120a` | +| Launch ABI | cubin argument sizes and offsets match the exported MLIR; `#[grid_constant] &TensorMap` is a 128-byte by-value argument aligned to 64 | + +Any failure stops the build. The importer is closed-world: a `cute_rs::` call it does not +recognise is an error, not a fallback. + +## What you own + +- Thread convergence and which lanes take part in each collective. +- Barrier phases, arrival counts and transaction bytes at run time. +- Asynchronous completion before a buffer or TMEM column is reused. +- Pointer lifetimes behind raw views and descriptors; the tensors a descriptor points + to must stay alive until the GPU is done with them. +- Disjoint writes across threads. + +## Two rules worth remembering + +**SM100 collective TMA loads.** `Sm100ClusterTmaCopy::copy` needs all 32 lanes of the +producer warp with identical arguments; the operation elects a lane internally. An +outer elected-lane branch deadlocks once a stage is reused. TMA stores have one issuing +thread, so this rule does not apply to them. + +**The A collector on SM100.** CUTLASS 4.7's MMA atom cannot choose when to fill or +release the A collector, so the backend emits the `tcgen05` MMA instruction from the +verified plan: + +```text +each K16: first N partition ──► fill the A collector + second N partition ──► last use of the A collector +``` + +Translation owns the descriptor bits, instruction selection, barrier arrivals, +completion multicast and TMEM fences. + +## Grid-constant launch arguments + +```text +#[grid_constant] desc: &TensorMap<...> + +host descriptor ──► 128-byte launch argument ──► device &TensorMap<...> + aligned to 64 bytes same address for every thread +``` + +The descriptor needs no device allocation. The full launch contract is under +[Kernels and Device Functions](../gpu-programming/kernels-and-device-functions.md). diff --git a/cuda-oxide/cuda-oxide-book/cute/cutedsl-comparison.md b/cuda-oxide/cuda-oxide-book/cute/cutedsl-comparison.md new file mode 100644 index 0000000000..091ce1977c --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/cutedsl-comparison.md @@ -0,0 +1,116 @@ +# From Python CuTeDSL to cute-rs + +If you have used NVIDIA's Python **CuTeDSL**, the concepts in cute-rs will feel immediately familiar: layouts, composed swizzles, `zipped_divide`, TMA copy atoms, and tiled MMA. Both compile down to the **exact same CUTLASS `cute` and `cute_nvgpu` MLIR operations**. + +The key difference is **where the static configuration lives** and **when the IR is constructed**: + +```{figure} images/cute-rs-vs-cutedsl.svg +:align: center +:width: 100% + +Python CuTeDSL traces Python objects wrapping live MLIR values at runtime; Rust cute-rs encodes layouts, swizzles, and MMA plans in Rust types and translates them ahead of time. +``` + +--- + +## Why the Kernel Signatures Look Different + +In Python **CuTeDSL**, `@cute.jit` traces a Python function at runtime to build an MLIR module. Because Python objects like `cute.Tensor`, `cute.CopyAtom`, `cute.TiledMma`, and `cute.ComposedLayout` directly wrap MLIR values, the host script constructs all of them and passes them as parameters into `@cute.kernel`: + +```python +# Python CuTeDSL: 18 parameters wrapping MLIR types during @cute.jit tracing +@cute.kernel +def kernel( + self, + tma_atom_a: cute.CopyAtom, mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, mC_mnl: cute.Tensor, + tiled_mma: cute.TiledMma, cta_layout_mnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, b_smem_layout: cute.ComposedLayout, + sfa_smem_layout: cute.Layout, sfb_smem_layout: cute.Layout, + epi_smem_layout: cute.ComposedLayout, tile_sched_params: PersistentTileSchedulerParams, +): ... +``` + +In Rust **cute-rs**, there is no runtime tracer: rustc compiles the kernel ahead of time. Nothing is hidden: +1. **Kernel parameters carry only the hardware launch ABI** (for example, 128-byte hardware TMA descriptors passed by value via `#[grid_constant]`, device pointers, and scalar dimensions). +2. **Layouts, swizzles, operand roles, and MMA configurations live in Rust types** (`ASmem`, `Sm100TiledMma<...>`, `MainloopPipeline`) with zero runtime byte overhead. + +```rust +// Rust cute-rs: static configuration lives in types; arguments carry only runtime TMA descriptors +type ASmem = ComposedLayout, 0, RowMajor<128, 64>>; +type B0Smem = ComposedLayout, 0, RowMajor<96, 64>>; +type B1Smem = ComposedLayout, 0, RowMajor<80, 64>>; +type CSmem = ComposedLayout, 0, RowMajor<128, 32>>; +type TiledMma = Sm100TiledMma; +type MainloopPipeline = PipelineTmaUmma<5, 77_824>; + +#[kernel] +pub fn fp16_gemm_256x352( + #[grid_constant] a: &TensorMap, + #[grid_constant] b0: &TensorMap, + #[grid_constant] b1: &TensorMap, + #[grid_constant] c: &TensorMap, + bias: *const f16, m: u64, n: u64, k: u64, has_bias: u64, +) { ... } +``` + +--- + +## How mir-importer Translates Rust Types to CuTe IR + +When mir-importer translates a cute-rs call, it matches the callee's canonical rustc `DefId`, structurally walks its monomorphized generic type arguments (`static_config.rs`), runs cute-layout static validation proofs, and emits a dialect-cute operation combining those static attributes with the call's runtime SSA operands: + +```{figure} images/rust-types-to-cute-ir.svg +:align: center +:width: 100% + +Compile-time Rust types supply the static IR attributes (shape, stride, swizzle, operand role, MMA partitions), while kernel values supply the runtime SSA operands. +``` + +--- + +## Side-by-Side Concept Map + +| Concept | Python CuTeDSL | Rust cute-rs | +| :------ | :------------- | :----------- | +| **Layout definition** | `cute.make_layout((128, 64), stride=(64, 1))` | `RowMajor<128, 64>` or `Layout<(C<128>, C<64>), (C<64>, C<1>)>` | +| **Swizzled SMEM layout** | `cute.make_composed_layout(Swizzle(3, 3, 3), 0, layout)` | `ComposedLayout, 0, RowMajor<128, 64>>` | +| **Tensor view** | `cute.make_tensor(ptr, layout)` | `Tensor::from_slice(a)`, `TensorMut::from_disjoint_slice(&mut c)`, `Tensor::from_storage(storage, layout)`: constructors are `from_*`, the Rust form of `make_*` | +| **Tiling** | `cute.local_tile(mA, tiler, coord)` | `Tensor::local_tile::(tile)`, equal to `zipped_divide::().slice(tile)`: a one-dimensional tiler and one tile coordinate, no projection | +| **Register fragment** | `cute.make_fragment_like(tensor, dtype)` | no function: the tile or MMA atom type fixes the fragment. `RegisterTile` comes from `Tensor::load()` or `RegisterTile::from_array(..)`; `FragmentA`, `FragmentB` and `FragmentC::ZERO` serve the MMA atoms | +| **TMA descriptor** | `(tma_atom, tma_tensor)` pair passed into `@cute.kernel` | `#[grid_constant] a: &TensorMap` (128B by value) | +| **TMA encoding (host)** | `cpasync.make_tiled_tma_atom(op, mA, smem_layout, cta_tiler)` | `create_tensor_map_tiled::(ptr, rows, cols, ld)` in `cute_rs::copy::tma`, the rank-three SM100 form in `cute_rs::attention::sm100`; `size_in_bytes::()` gives the expected bytes | +| **TMA copy** | `cute.copy(tma_atom, src, dst, tma_bar_ptr=bar)` in both directions | `copy_tma_2d` and `copy_tma_s2g_2d`; `Sm100ClusterTmaCopy::copy`; `TmaCopy::load(desc, &mut dst, row, col, barrier, rank)` and `TmaCopy::store(desc, &src, row, col)` in `cute_rs::attention::sm100` | +| **Shared tensor view** | `storage.sA.get_tensor(a_smem_layout)` | `SharedTensor` or `Sm100SharedTile` | +| **TMA pipeline** | `pipeline.PipelineTmaAsync.create(...)` or `pipeline.PipelineTmaUmma.create(...)` | `PipelineTmaAsync` or `PipelineTmaUmma<5, 77_824>` | +| **Pipeline state** | `make_pipeline_state(PipelineUserType.Producer, stages)`, `state.index`, `state.phase`, `state.advance()` | `PipelineState::::new()`, `index() -> usize`, `phase()`, `advance()`; the SM100 attention ring's `RingCursor::index() -> u32` | +| **Pipeline producer step** | `producer_acquire(state)` waits on the empty barrier and arrives with the expected bytes, under `is_leader_cta` in `PipelineTmaUmma` | `producer_acquire` waits only; the elected lane then calls `producer_expect_tx` on `PipelineTmaAsync`, `PipelineTmaUmma` and `TwoCtaRing` (`producer_expect_tx_cursor` with a cursor) | +| **Persistent scheduler** | `StaticPersistentTileScheduler.create(params, block_idx, grid_dim)`, `get_grid_shape(params, max_active_clusters)`, `work = get_current_work()`, `work.is_valid_tile`, `work.tile_idx`, `advance_to_next_work(advance_count=1)` | `StaticPersistentTileScheduler::::new_1d()`, `get_grid_shape(sm_count, active_ctas_per_sm) -> Option`, `has_work()`, `get_current_work()`, `tile_idx() -> (m, n, l)` with `linear_idx()`, `advance_to_next_work()` | +| **Tiled MMA** | `cute.make_tiled_mma(mma_op, atom_layout, ...)` | `Mxf4TiledMma::::get_slice(lane)` or `Sm100TiledMma` | +| **Tiled copy** | `thr_copy = tiled_copy.get_slice(tidx)`; `cute.copy(tiled_copy, thr_copy.partition_S(gA), thr_copy.partition_D(sA))` | `tiled_copy.get_slice(thread_idx).copy(&g_a, tile_coord, &mut s_a)`: both partitions come from the plan's layouts inside `copy_g2s`. No kernel can reach this surface yet; see the next section | +| **Register vectors** | `TensorSSA` with `a + b`, `a - b`, `a * b`, `a / b`, `a // b`, `-a`, `where`, conversions and reductions | `RegisterTile` with `+`, `-`, `*`, `/` and unary `-`; `map`, `zip_map`, `fold` and `gather` for the rest | +| **Architecture primitives** | `cute.arch.thread_idx()`, `cute.arch.warp_idx()`, `cute.arch.elect_one()`, `cute.arch.cluster_arrive()`, `cute.arch.cluster_dim()`, the primitives-layer `fence_sync_restrict(MemOrder.ACQUIRE)`, ... | `cute_rs::arch::thread_idx()`, `warp_idx()`, `elect_one_sync()`, `cluster_arrive()`, `cluster_dim()`, `fence_sync_restrict::()`, ...: forwarders over the PTX-named ctm-rs layer, the mask passed as an argument | + +--- + +## Where cute-rs Deviates, and Why + +The names follow CuTeDSL 4.7. Where the Rust form differs on purpose, the difference and its reason are listed here, so a CuTeDSL author is not surprised. + +| Concept | CuTeDSL 4.7 | Rust cute-rs | Why | +| :------ | :---------- | :----------- | :-- | +| **Constructors and plans** | `make_layout`, `make_tensor`, `make_tiled_mma` and `make_tiled_copy` build objects while tracing | `new` and `from_*` constructors; a compile-time plan is a type or a const generic (`RowMajor<128, 64>`, `TiledCopy`, `zipped_divide::<8>()`) | the compiler reads the plan from the generic arguments, so there is nothing for a `make_*` call to build | +| **Pipeline producer step** | `producer_acquire` folds the arrive-expect into the wait, under `is_leader_cta` in `PipelineTmaUmma` | `producer_acquire` waits; the elected lane calls `producer_expect_tx`, or `producer_expect_tx_cursor` on `TwoCtaRing` | electing the lane is a convergence obligation the compiler does not prove, so it stays visible in kernel code; the pipeline value stores pointers only | +| **Scheduler** | `WorkTileInfo.is_valid_tile`; `get_grid_shape(params, max_active_clusters)` returns `(m, n, persistent_clusters)`; `advance_to_next_work(advance_count=1)`; `tile_idx` is a `Coord` | `has_work()` on the scheduler; `get_grid_shape(sm_count, active_ctas_per_sm) -> Option`; `advance_to_next_work()`; `tile_idx() -> (m, n, l)` with `linear_idx()` beside it | the grid is one dimensional and has no params object; the bounds check stays outside the coordinate arithmetic; the step is always one grid of blocks; the coordinate is a Rust tuple | +| **Pipeline state** | `PipelineState.index` is an `Int32` | `PipelineState::index() -> usize`; `RingCursor::index() -> u32` | a `usize` scales a shared-buffer offset directly; the SM100 attention ring keeps the `u32` its barriers take | +| **TMA names** | `make_tiled_tma_atom` returns the atom and a TMA tensor; one `cute.copy` loads and stores | `TensorMap` is the one descriptor type; `create_tensor_map_tiled` is the one host encoder, with the rank-three SM100 form in `cute_rs::attention::sm100`; `TmaCopy::load` and `TmaCopy::store` stay two methods; `size_in_bytes::()` is `cute.size_in_bytes`, a cosize, so a padded row stride counts and a swizzle does not | the descriptor is CUDA's `CUtensorMap` typed by element and shared layout; load takes a barrier and a CTA rank and store takes neither, so one signature cannot serve both | +| **Clusters** | `cute.arch.cluster_dim()` is the number of clusters in the grid; `cute.arch.block_in_cluster_dim()` is CTAs per cluster | the same two names with the same meanings | the compiler's own `cluster_dim` vocabulary (the `#[cluster_launch]` attribute, the `cluster_dim` field of `cuda_launch!` and the `cluster_dim_x/y/z` NVVM annotations) means CTAs per cluster; in a kernel, `cluster_dim()` reads as the DSL does and `block_in_cluster_dim()` is the launch value | +| **Named barriers** | `pipeline.NamedBarrier(barrier_id, num_threads)` with `arrive_and_wait()`, its alias `sync()`, `arrive()`, `wait()` and the unaligned forms | `NamedBarrier::::sync()` only; `cute_rs::arch::barrier(id, n)` and `barrier_arrive(id, n)` are the `cute.arch` free functions | both values are const generics, so the id and the thread count (a multiple of 32, at most 1024) are checked at compile time; the split arrive and wait methods are not bound on the type | +| **Sync-restricted fence** | `fence_sync_restrict(order)` in the experimental primitives layer, `order` a `MemOrder` string enum | `cute_rs::arch::fence_sync_restrict::()`, the order a sealed marker type; acquire only | the instruction is chosen when the kernel compiles; the release form waits for its intrinsic to be bound | +| **Tiled copy** | `make_tiled_copy`, `get_slice`, `partition_S`, `partition_D`, `cute.copy` | `TiledCopy`, `get_slice` and `ThrCopy::copy` exist, but no kernel can reach them yet | `GlobalCopyTensor` has no public constructor: `from_carrier` is `pub(crate)` and the only public producers are the block-scaled tensors' `values_for_copy` and `scale_words_for_copy`; the importer's `shared_tensor_overlay` hook accepts only the block-scaled roles `Mkl` and `Nkl`; the verifier rejects a role-typed shared carrier as a `copy_g2s` destination. The kernels that need a block-wide copy call `copy_g2s` directly | +| **Register vectors** | `/` on integer vectors converts both sides to f32 and divides; `//` floors; `-a` is `0 - a` | `/` on integer lanes truncates toward zero; unary `-` negates each lane, so a `+0.0` lane becomes `-0.0` | the operators are Rust's own: floating-point division matches the DSL, integer division and the sign of a negated zero differ | + +ctm-rs keeps PTX instruction names (`cp_async_bulk_tensor_shared_cta_global`, `mbarrier_arrive_expect_tx`, `fence_sync_restrict_acquire`); the `cute.arch` concept names live in `cute_rs::arch` as forwarders over them. Dialect op, type and attribute names are compiler-internal and are not part of this surface. diff --git a/cuda-oxide/cuda-oxide-book/cute/gemm-and-pipelines.md b/cuda-oxide/cuda-oxide-book/cute/gemm-and-pipelines.md new file mode 100644 index 0000000000..f2753b2588 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/gemm-and-pipelines.md @@ -0,0 +1,97 @@ +# TMA Pipelines, Tensor Cores, and ctm-rs Interop + +The Blackwell kernels in this preview combine multi-stage TMA pipelines, swizzled shared-memory views, and tensor-core MMA operations. Depending on the workload, you can express the entire schedule in **cute-rs** (`blockscale_gemm_cute`) or pair **cute-rs** with **ctm-rs** instruction primitives in the same kernel (`fp16_gemm_256x352_cute`, `sm100_fmha_prefill_cute`, `sm100_mla_decode_fp16_cute`). + +--- + +## 1. Block-Scaled MXFP4 GEMM in cute-rs (SM120a) + +In `blockscale_gemm_cute`, one producer warp (warp 8) fills a 3-stage TMA ring buffer in shared memory while eight compute warps (warps `0..7`) load swizzled FP4 fragments and `Float8E8M0FNU` scales into `Mxf4TiledMma`: + +```{figure} images/blockscale-gemm-pipeline.svg +:align: center +:width: 100% + +Data flow of `blockscale_gemm_cute` on SM120a: 4 typed TMA descriptors feed a 3-stage `PipelineTmaAsync`, consumed by 8 compute warps running `Mxf4TiledMma` and an `Sm120Epilogue128x128` TMA store epilogue. +``` + +```rust +// Shared-memory layouts and operand-role views (Mkl for A, Nkl for B) +type TileP = Layout<(C<128>, C<32>), (C<32>, C<1>)>; +type SmemP = ComposedLayout, 0, TileP>; +type ByteSmem = ComposedLayout, 0, RowMajor<128, 64>>; +type SharedA = SharedTensor; +type SharedB = SharedTensor; + +// Compute warps (0..7) wait on the TMA slot, load fragments + scales, and accumulate: +pipeline.consumer_wait(&consumer); +let slot = consumer.index(); +let s_a = unsafe { SharedA::from_raw_parts(smem_base.add(slot * STAGE_BYTES).cast()) }; +let s_b = unsafe { SharedB::from_raw_parts(smem_base.add(slot * STAGE_BYTES + B_OFF).cast()) }; + +let scale_tile = unsafe { tiled_mma.load_scale_atom_128(&s_sfa, &s_sfb, wm, wn) }; +let mut kk = 0usize; +while kk < 2 { + let scales = unsafe { scale_tile.pairs_at_unchecked(kk) }; + let a_frags = unsafe { tiled_mma.load_a_128(&s_a, wm, kk) }; + let b_frags = unsafe { tiled_mma.get_b_tile_k64(&s_b, wn, kk) }; + accumulators.gemm(&tiled_mma, a_frags, b_frags, scales); + kk += 1; +} +pipeline.consumer_release(&consumer); +consumer = consumer.advance(); +``` + +--- + +## 2. Combining cute-rs and ctm-rs on SM100a (2-CTA GEMM) + +On datacenter Blackwell (`sm_100a`), two CTAs cooperate on a `256×352` output tile (`256×192` + `256×160` N-partitions sharing the A collector) and accumulate into Tensor Memory (TMEM). + +`fp16_gemm_256x352_cute` uses **cute-rs** for layout-checked cluster TMA, pipelines, paired MMA, and TMEM epilogue slices (dialect-cute), while using **ctm-rs** for cluster rank, barriers, and collective `TmemGuard` allocation (dialect-nvvm): + +```{figure} images/sm100-two-cta-gemm.svg +:align: center +:width: 100% + +In `fp16_gemm_256x352_cute`, cute-rs (teal) owns the 5-stage cluster TMA pipeline, paired 2-CTA MMA, and TMEM epilogue, while ctm-rs (purple) supplies cluster rank, barriers, and `TmemGuard` lifecycle in the same Pliron module. +``` + +```rust +// Producer warp (warp 5): acquire the stage, let one elected lane post the +// expected bytes, then all 32 lanes issue the cluster TMA copies via cute-rs +pipeline.producer_acquire(stage as u32, phase); +if rank == 0 && warp::is_elected_sync(u32::MAX) { + pipeline.producer_expect_tx(stage as u32); +} +let barrier = pipeline.producer_get_barrier(stage as u32); +Sm100ClusterTmaCopy::copy(a, &mut a_tile, row, kt * 64, barrier, rank); +Sm100ClusterTmaCopy::copy(b0, &mut b0_tile, col, kt * 64, barrier, rank); +Sm100ClusterTmaCopy::copy(b1, &mut b1_tile, col + 192, kt * 64, barrier, rank); + +// MMA warp (warp 4, CTA rank 0): paired 2-CTA MMA into TMEM +pipeline.consumer_wait(stage as u32, phase); +if warp::is_elected_sync(u32::MAX) { + mma.gemm(tmem, &a_tile, &b0_tile, &b1_tile, accumulate); + pipeline.consumer_release(stage as u32); +} + +// Epilogue warps (warps 0..3): allocate TMEM via ctm-rs, drain 11 slices via cute-rs +if warp_id == 0 { + TmemGuard::::from_raw_slot(tmem_token).alloc_collective(); + tcgen05::relinquish_alloc_permit_cg2(); +} +``` + +--- + +## Where to Go Next + +| Workload | Location | Highlights | +| :------- | :------- | :--------- | +| **Elementwise / SwiGLU / RMSNorm** | `cuda-oxide/cute-rs-preview/examples/{elementwise,swiglu,rmsnorm}_cute` | 1D tiling, register transforms, warp reductions, and ragged tails. | +| **FFN & RoPE KV-Cache Attention** | `cuda-oxide/cute-rs-preview/examples/{ffn,rope_cache}_cute` | BF16 tensor-core projections, typed weight roles, and 128-wide GQA. | +| **MXFP4 GEMV & Block-Scaled GEMM** | `cuda-oxide/cute-rs-preview/examples/{nvfp4_gemv,blockscale_gemm}_cute` | Packed FP4 + `Float8E8M0FNU` scales, 3-stage TMA pipelines, and `*_ctm` counterparts. | +| **SM100a 2-CTA GEMM & Attention** | `cuda-oxide/cute-rs-preview/examples/{fp16_gemm_256x352,sm100_fmha_prefill,sm100_mla_decode_fp16}_cute` | Two-CTA cluster TMA, TMEM, A-collector reuse, and FlashAttention / MLA decode. | +| **Qwen3-8B (BF16 & NVFP4)** | `cuda-oxide/cute-rs-preview/models/qwen3_8b_cute` | 36-layer `Qwen/Qwen3-8B` (`H=4096, I=12288, GQA 32:8`) in BF16 and NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) on RTX 5090 (`visual.md`). | +| **End-to-End Qwen3-TTS Model** | `cuda-oxide/cute-rs-preview/models/qwen_decoder_cute` | Complete Talker + Predictor speech synthesis on SM120a (see `visual.md` inside that folder for a visual tour). | diff --git a/cuda-oxide/cuda-oxide-book/cute/images/blockscale-gemm-pipeline.svg b/cuda-oxide/cuda-oxide-book/cute/images/blockscale-gemm-pipeline.svg new file mode 100644 index 0000000000..0f6ee5f35f --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/blockscale-gemm-pipeline.svg @@ -0,0 +1,46 @@ + + + +Block-scaled MXFP4 GEMM pipeline in cute-rs on SM120a + + + +Block-scaled MXFP4 GEMM pipeline (blockscale_gemm_cute, SM120a) +One TMA producer warp feeds a 3-stage shared-memory ring while 8 compute warps accumulate K64 block-scaled MMA tiles. + +GLOBAL MEMORY & TMA +4 Typed TMA Load Descs +A: E2M1 [M, K] +B: E2M1 [N, K] +SFA: UE8M0 scales +SFB: UE8M0 scales +TensorMap<u8, ByteSmem> + + +WARP 8 (PRODUCER) +3-Stage Shared Ring +PipelineTmaAsync<3> +SharedA (Mkl, B64) +SharedB (Nkl, B64) +Sm120ScaleAtom + + +WARPS 0..7 (COMPUTE) +Mxf4TiledMma +load_a_128 / get_b_tile +load_scale_atom_128 +gemm(...) +FP32 Mxf4AccumulatorTile2x8 + + +EPILOGUE & STORE +Sm120Epilogue +FP32 → FP16x2 +stmatrix → smem +copy_tma_s2g_2d +Global C [M, N] FP16 + +Persistent CTA Work Loop — StaticPersistentTileScheduler +Each resident CTA steps across 128×128 output tiles; PipelineState index and phase parity persist across tiles without resetting barriers. +while scheduler.has_work() { let (m, n, _) = scheduler.get_current_work().tile_idx(); ... scheduler.advance_to_next_work(); } + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/ctm-rs-at-a-glance.svg b/cuda-oxide/cuda-oxide-book/cute/images/ctm-rs-at-a-glance.svg new file mode 100644 index 0000000000..967c83672d --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/ctm-rs-at-a-glance.svg @@ -0,0 +1,100 @@ + + + + +ctm-rs at a glance +Typed instruction primitives beneath cute-rs. The types check the local contract; the author owns the schedule. +TYPED RUST API +INSTRUCTION FAMILIES (dialect-nvvm) +COMPILER PROFILES + +SharedRegion<F> +alignment, capacity, address space + +TensorMap<f16, F> +128-byte descriptor, grid-constant ABI + +TmemGuard<State, N, CtaGroup2> +collective alloc, borrowed views + +SmemOperand<Role, F> +operand role and physical format + +InstructionDescriptor<S> +MMA shape, K-major or MN-major + +Mxf4A, Mxf4B, ScalePairA/B +A4 B2 C4 registers, UE8M0 pairs + +CtaMask<G>, ManagedBarrier +nonzero mask, ready barrier + +tcgen05 +sm_100a +alloc, dealloc, F16 two-CTA MMA, 32x32b ld/st, commit, fences + +TMA +both +2D load, two-CTA multicast load, 2D store, prefetch, bulk wait + +mbarrier, cluster +both +init, arrive, expect_tx, parity wait, mapa, remote arrive + +mma.sync +MXF4: sm_120a +m16n8k16 F16 and BF16; m16n8k64 MXF4 with UE8M0 scales + +ldmatrix, stmatrix +both +x4 loads for A and B, x2 FP16 stores + +cp.async +both +4-byte copies, noinc barrier arrive, commit and wait groups + +warp, numerics +both +shuffle, elect, vote, reduce; bf16 convert, f32x2, approx math + +selects + +cutlass-primitives-mlir22 +ctm-rs alone + +no CuTe operations, types or attributes + +the *_ctm GEMM examples + +CUDA_OXIDE_MLIR_PROFILE selects it + +cutlass-full-cute-mlir22 +ctm-rs beside cute-rs + +one Pliron module, two dialects + +the SM100 attention examples + +default profile + +feeds + +targets +sm_100a or sm_120a +tcgen05 needs sm_100a; MXF4 needs sm_120a + +The author still owns, at every unsafe boundary: +election and lane participation · barrier phases and arrival counts · asynchronous completion before reuse · collector fill and release order · pointer lifetimes + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-preview-stack-and-flow.svg b/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-preview-stack-and-flow.svg new file mode 100644 index 0000000000..01c648b2e3 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-preview-stack-and-flow.svg @@ -0,0 +1,66 @@ + + + +The cute-rs-preview stack: cute-rs over ctm-rs, lowered through dialect-cute to CUTLASS MLIR + + + +The cute-rs-preview stack: cute-rs over ctm-rs, lowered through dialect-cute to CUTLASS MLIR +cute-rs encodes tensor layouts and pipelines in Rust types; ctm-rs supplies instruction primitives and the kernel ABI. +1 · KERNEL AUTHORING IN RUST + + +cute-rs — Semantic CuTe abstractions +layout-driven +Tensor / Tile +TMA / Copy +Pipelines +TiledMma / Epilogue + +ctm-rs — Instruction primitives & kernel ABI +over cuda-device +#[kernel] / ABI +thread / cluster +mbarrier +tcgen05 / TMEM + +cute-layout +Shared layout algebra + +IntTuple / Layout +ComposedLayout / Swizzle +Static TMA, copy & +ldmatrix proofs + + + +2 · COMPILE-TIME RECOGNITION & IR IMPORT + +mir-importer (on rustc Stable MIR) +Matches canonical DefId paths · Decodes generic type markers · Emits Pliron ops + +3 · ONE PLIRON MODULE (COEXISTING DIALECTS) + + +dialect-cute +Tensors, TMA, pipelines, MMA + +dialect-mir +Control flow, scalars, pointers + +dialect-nvvm +Coords, barriers, TMEM, math +4 · VERIFY, EXPORT & COMPILE + + +Verify & Prepare +Provenance & protocol checks + + +cuda-oxide-mlir-export +cute + cute_nvgpu + nvvm MLIR + + +CUTLASS 4.7 → Cubin +Validated & embedded in host + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-vs-cutedsl.svg b/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-vs-cutedsl.svg new file mode 100644 index 0000000000..4dfa5bbca6 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/cute-rs-vs-cutedsl.svg @@ -0,0 +1,53 @@ + + + +Python CuTeDSL versus Rust cute-rs: runtime tracing of MLIR wrappers versus ahead-of-time type translation + + + +Python CuTeDSL vs. Rust cute-rs: runtime tracing vs. type translation +Both target the same CUTLASS cute and cute_nvgpu MLIR operations, but encode static configuration differently. + +Python CuTeDSL + +Rust cute-rs (cuda-oxide) +Execution model +when IR is built + +Runtime JIT Tracing (@cute.jit) +Python interpreter executes the function; +each Python object directly wraps an MLIR value + +Ahead-of-Time Rust Compilation +No Python runtime or tracing; layouts, swizzles, +and MMA plans are encoded in Rust types + +Kernel signature +what args carry + +Passes MLIR wrapper objects +tma_atom_a: cute.CopyAtom +mA_mkl: cute.Tensor, mma: cute.TiledMma +a_smem: cute.ComposedLayout, ... + +Passes explicit hardware ABI only +#[grid_constant] a: &TensorMap<f16, ASmem> +128B descriptor by value; CopyAtom, TiledMma +& layouts live in type ASmem / local types + +Lowering to IR +how ops appear + +Trace-time emission +Calling cute.copy(...) or cute.gemm(...) in +Python emits MLIR ops into the active trace + +Type-directed MIR translation +mir-importer decodes Rust generic types into +verified dialect-cute ops, then exports MLIR + + + +Same Target Dialect: CUTLASS cute & cute_nvgpu MLIR → Cubin +Identical layout algebra, TMA descriptors, shared swizzles, and tensor-core instructions + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/host-device-tma-type-agreement.svg b/cuda-oxide/cuda-oxide-book/cute/images/host-device-tma-type-agreement.svg new file mode 100644 index 0000000000..055b98e307 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/host-device-tma-type-agreement.svg @@ -0,0 +1,45 @@ + + + +Host and device TMA agreement through a shared Rust layout type + + + +One Rust layout type binds host TMA encoding to device TMA & MMA +Host and device never drift: both sides reference the same ComposedLayout<Swizzle, Offset, Layout> type definition. + + + +Shared Rust Layout Type (compiled for both host & device) + +pub type ASmem = ComposedLayout<Swizzle<3, 3, 3>, 0, RowMajor<128, 64>>; + + + + + + + +Host (CPU) — Descriptor Encoding & Launch + +create_tensor_map_tiled::<f16, ASmem>(ptr, m, k, ...) +• Reads tile shape (128×64) and 128B swizzle from ASmem +• Validates base alignment & calls cuTensorMapEncodeTiled +• Produces TensorMap<f16, ASmem> (128B, 64B-aligned) +Passed by value in launch params (no extra device alloc) + + + +128B + + + +Device (GPU) — Kernel & Compile-Time Lowering + +#[grid_constant] a: &TensorMap<f16, ASmem> + +Sm100SharedTile<ASmem, OperandA> +• Type system enforces matching ASmem on copy & MMA +• mir-importer proves ASmem swizzle phase alignment +• Lowers to cute_nvgpu TMA & tensor-core MMA ops + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/rust-types-to-cute-ir.svg b/cuda-oxide/cuda-oxide-book/cute/images/rust-types-to-cute-ir.svg new file mode 100644 index 0000000000..1d3e9faaff --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/rust-types-to-cute-ir.svg @@ -0,0 +1,43 @@ + + + +How Rust types and runtime values combine into verified dialect-cute operations + + + +From Rust types and kernel arguments to verified CuTe IR +Static layout and MMA types carry zero runtime bytes; mir-importer pairs them with SSA operands at compile time. + +Compile-time Rust types (zero bytes at runtime) + +ComposedLayout<Swizzle<3, 3, 3>, 0, RowMajor<128, 64>> + +Sm100SharedTile<ASmem, OperandA> + +Sm100TiledMma<ASmem, B0Smem, B1Smem, 192, 160> +Carries: shape, stride, swizzle, operand role, MMA partitions + +Runtime kernel values (SSA operands) + +#[grid_constant] a: &TensorMap<f16, ASmem> + +smem_ptr, tmem_addr, barrier_ptr + +row, kt * 64, stage, phase, accumulate +Carries: 128B tensor map, addresses, coordinates, phases + + + +mir-importer::translator::cute + cute-layout::validation +1. Recognizes canonical DefId (e.g. Sm100ClusterTmaCopy::copy, Sm100TiledMma::gemm) +2. Decodes generic type tree into static Layout / Swizzle attrs & proves alignment and coverage + + +Emitted dialect-cute operation → CUTLASS MLIR + +Static IR Attributes & Types +Swizzle<3, 3, 3>, (128, 64):(64, 1), Role=A + +Runtime SSA Operands +%desc_a, %smem_tile, %row, %col, %bar + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/sm100-two-cta-gemm.svg b/cuda-oxide/cuda-oxide-book/cute/images/sm100-two-cta-gemm.svg new file mode 100644 index 0000000000..87eb948deb --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/sm100-two-cta-gemm.svg @@ -0,0 +1,38 @@ + + + +SM100 two-CTA GEMM combining cute-rs and ctm-rs + + + +Combining cute-rs and ctm-rs in one kernel (fp16_gemm_256x352_cute, SM100a) +Teal boxes use cute-rs (dialect-cute); purple boxes use ctm-rs instruction primitives (dialect-nvvm) in the same Pliron module. + +TMA WARP (WARP 5) +5-Stage Cluster TMA Load +PipelineTmaUmma<5, 77_824> +Sm100ClusterTmaCopy::copy +ASmem (128×64) + B0/B1Smem +All 32 lanes participate per CTA + + +MMA WARP (WARP 4, CTA 0) +Paired 2-CTA MMA → TMEM +Sm100TiledMma<..., 192, 160> +256×192 fills A collector +256×160 reuses A collector +Accumulates in 512 TMEM columns + + +EPILOGUE WARPS (WARPS 0..3) +TMEM → Shared C → TMA +Sm100TmemEpilogue::<CSmem> +11 × (128×32) slices + bias +sm100::PipelineTmaStore<2> +FP32 → FP16 rounded & TMA stored + +ctm-rs Instruction Primitives in the Same Kernel (dialect-nvvm) +cluster_ctarank() / warp_id() +barrier::barrier_cta_sync + fences +TmemGuard::<..., 512, CtaGroup2> + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/support-additions.svg b/cuda-oxide/cuda-oxide-book/cute/images/support-additions.svg new file mode 100644 index 0000000000..4f8672fa3c --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/support-additions.svg @@ -0,0 +1,129 @@ + + + +Feature additions in CuTeDSL 4.4 to 4.7 +What a CuTeDSL user relies on beyond the per-architecture atoms, and what the preview has instead. +FEATURE +CuTeDSL 4.7 +cute-rs +WHAT THE PREVIEW HAS + +Epilogue fusion (EFC) + + +two fixed FP16 epilogues; activations fused by hand; no user hook +Dynamic layouts + + +runtime lengths and one row pitch; tiles, strides and swizzles are types + +JIT and caching + + +ahead of time through rustc and Cargo; nothing to JIT or cache +AOT export + + +cubin embedded in the Rust binary; no C header or .so export + +TVM FFI, PyTorch, JAX, DLPack + + +none; buffers live in cuda-core DeviceBuffer +Autotuning + + +none; tile shapes are const generics, tune in a host harness + +Task Scheduling + + +none; the verifier checks protocol order, not schedules +Inline PTX and cute.ffi + + +none on the CuTe path; hardware through ctm-rs and cuda-device + +printf, print_tensor, IR dumps + + +MLIR, MIR and LLVM dumps; no device printf in CuTe kernels +IKET profiling + + +none; profile the embedded cubin with Nsight + +CUDA graphs, PDL, cooperative launch + + +graphs, the PDL attribute and cooperative launch from the host +Typed shared memory + + +typed static and dynamic arrays; no bump allocator + +TensorSSA register vectors + + +RegisterTile with +, -, *, / and unary -, map, zip_map, fold and gather +block_copy + + +a typed two-CTA TMA copy; no shared-to-TMEM copy + +Grouped GEMM, MoE, ptr-array + + +none; one launch per group +Mixed-input GEMM + + +none; operands share one type + +Distributed kernels + + +none; single GPU +Grouped block-scaled and blockwise GEMM + + +dense block-scaled GEMM and GEMV only + +FMHA backward + + +forward only +Convolution + + +implicit im2col on the BF16 register path; no TMA im2col + +Sub-byte types + + +FP4 and scale carriers; no Int4, no FP6 +Primitives API (4.7, experimental) + + +ctm-rs covers the shipped kernels; no atomics, wgmma, libdevice + +available + +fixed or narrower form only + +not available + +not applicable + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/support-matrix.svg b/cuda-oxide/cuda-oxide-book/cute/images/support-matrix.svg new file mode 100644 index 0000000000..673adccf40 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/support-matrix.svg @@ -0,0 +1,204 @@ + + + +cute-rs preview against CuTeDSL 4.7 +Rust CuTe kernels compiled ahead of time through the pinned CUTLASS 4.7 compiler for two Blackwell targets. Status at 5eaa6cbda (2026-10-22). +TARGETS + +Ampere sm_80/86/87 + +CuTeDSL + +cute-rs + +Ada sm_89 + +CuTeDSL + +cute-rs + +Hopper sm_90a + +CuTeDSL + +cute-rs + +Blackwell sm_100a + +CuTeDSL + +cute-rs + +Blackwell sm_103a, sm_110a + +CuTeDSL + +cute-rs + +GeForce Blackwell sm_120a + +CuTeDSL + +cute-rs + +GeForce Blackwell sm_121a + +CuTeDSL + +cute-rs +FEATURE AREA +CuTeDSL 4.7 +cute-rs +sm_100a +cute-rs +sm_120a +WHAT THE PREVIEW HAS + +Layout algebra, static + + + +layouts are Rust types checked at compile time; no layout ops in kernels +Dynamic layouts, runtime shapes + + + +runtime lengths and one runtime row pitch; strides and tiles are types + +Tensors, partitions, tiled atoms + + + +1D tiles (zipped_divide, slice, 1D local_tile) and fixed 2D block-scaled views; no make_tiled_copy +TMA copies + + + +2D tile load and store, two-CTA load; no im2col, reduce-store, multicast + +cp.async, bulk, DSMEM copies + + + +cooperative cp.async with a coverage proof; no bulk or DSMEM data copies +ldmatrix, stmatrix, atomics + + + +ldmatrix x4 and x2.trans, stmatrix x2, no-allocate stores; no atomics + +tcgen05 MMA kinds + + + +1 of 11 kinds: F16 in, F32 accumulate, one or two CTAs +Warp MMA kinds (mma.sync) + + + +3 of 7 kinds: MXF4 with UE8M0 scales (sm_120a), F16, BF16 + +NVFP4 tensor cores (E4M3 scales) + + + +NVFP4 projections run on CUDA cores through scaled loads and dot +wgmma (Hopper) + + + +no sm_90a target + +TMEM + + + +allocate, 32x32b load and store, FP32 to FP16; no 16xNb, LdRed or S2T +Pipelines and barriers + + + +5 typed pipelines and a two-CTA ring; no CLC, no compiler pipelining + +Schedulers, persistence, PDL + + + +one static 1D persistent scheduler; no grouped, CLC or stream-K +Clusters and multicast + + + +two-CTA clusters validated; no N-CTA multicast plan + +Epilogue fusion (EFC) + + + +two fixed FP16 epilogues, bias on sm_100a; activations fused by hand +Numerics and intrinsics + + + +f16, f32 native; BF16 and FP4 carriers; no inline PTX or libdevice + +Control flow, printf, debugging + + + +full Rust control flow and traps, MLIR dumps; no device printf +Host: JIT, torch, JAX, AOT export + + + +AOT cubin inside the Rust binary; no JIT, no Python interop + +Autotuning, Task Scheduling, IKET + + + +no autotuner, schedule checker or in-kernel tracing +Verification before codegen + + + +whole-module layout, pipeline and role checks; CuTeDSL checks at trace + +Grouped, MoE, mixed-input, FP8 + + + +dense kernels only +Shipped kernel examples + + + +4 sm_100a and 10 sm_120a example crates versus 86 CuTeDSL kernel files + +End-to-end models + + + +Qwen3-TTS 1.7B and Qwen3-8B on sm_120a; CuTeDSL ships none + +available on that target + +fixed or narrower form only + +not available + +not applicable to that target +Every other architecture is rejected before compilation. The support matrix page in the cuda-oxide book has the full table and evidence. + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/support-targets.svg b/cuda-oxide/cuda-oxide-book/cute/images/support-targets.svg new file mode 100644 index 0000000000..497921b17e --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/support-targets.svg @@ -0,0 +1,73 @@ + + + +Targets +CuTeDSL 4.7 compiles for all of these. The preview compiles for two and rejects the rest before any MLIR is built. +TARGET +GPU +CuTeDSL 4.7 +cute-rs preview + +sm_80 +sm_86 +sm_87 +Ampere + + +rejected before compilation +sm_89 +Ada + + +rejected before compilation + +sm_90a +Hopper + + +rejected before compilation +sm_100a +Blackwell datacenter, B200 class + + +yes + +sm_100 +sm_100f +sm_103a +sm_110a +Blackwell datacenter, other variants + + +rejected before compilation +sm_120a +Blackwell GeForce, RTX 50 series + + +yes + +sm_120 +sm_120f +sm_121a +Blackwell GeForce, other variants + + +rejected before compilation + +compiles + +rejected by the profile gate +Within the pair, tcgen05 operations need sm_100a and MXF4 register MMA needs sm_120a. cargo oxide run uses the detected GPU; cargo oxide build needs --arch. + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/three-tiers-comparison.svg b/cuda-oxide/cuda-oxide-book/cute/images/three-tiers-comparison.svg new file mode 100644 index 0000000000..769e7b0ecb --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/three-tiers-comparison.svg @@ -0,0 +1,45 @@ + + + +Three kernel authoring tiers in cuda-oxide: cute-rs, ctm-rs, and cuda-device + + +Choosing the right layer in cuda-oxide +Use cute-rs for tensor layouts, copies, and MMA; drop down to ctm-rs or cuda-device only when you need explicit hardware control. + +PRIMARY TENSOR LAYER +cute-rs + cute-layout + +Tensor / TensorMut / slice +ComposedLayout<Swizzle, RowMajor> +TensorMap / PipelineTmaAsync +Mxf4TiledMma / Sm100TiledMma +Lowers to: dialect-cute +Best for: GEMM, GEMV, FlashAttention, +MLA, RMSNorm, SwiGLU, full models +Examples: *_cute, qwen_decoder_cute + +INSTRUCTION PRIMITIVES +ctm-rs + +#[kernel] & launch contracts +thread / warp / cluster coords +mbarrier & async proxy fences +tcgen05::tmem::TmemGuard +Lowers to: dialect-nvvm +Best for: warp-specialized barriers, +TMEM lifecycle, or CTM-only kernels +Examples: *_ctm + inside SM100 *_cute + +FOUNDATIONAL SIMT LAYER +cuda-device + +DisjointSlice / SharedArray +Scalar & warp shuffle ops +Atomics & raw PTX intrinsics +Arbitrary Rust control flow +Lowers to: dialect-mir + nvvm +Best for: sparse/irregular SIMT, +pointer chasing, custom graph ops +Examples: vecadd, async_mlp, etc. + diff --git a/cuda-oxide/cuda-oxide-book/cute/images/zipped-divide-and-slice.svg b/cuda-oxide/cuda-oxide-book/cute/images/zipped-divide-and-slice.svg new file mode 100644 index 0000000000..080b0333b0 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/images/zipped-divide-and-slice.svg @@ -0,0 +1,56 @@ + + + +1D tensor tiling with zipped_divide and per-thread slicing + + + +How zipped_divide::<4>() and .slice(tid) partition a 1D tensor +A flat 1,003-element slice becomes a 2D logical view: (offset_in_tile, tile_index). Full tiles vectorize; the final ragged tile loops safely. + + + +Tensor::from_slice(a) // 1,003 x f32 in global memory +Flat index: [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, ... , 996, 997, 998, 999, 1000, 1001, 1002 ] + + +.zipped_divide::<4>() → .slice(tid) + + + +Thread 0: .slice(0) +[ 0, 1, 2, 3 ] +is_full() == true (16B) + + +Thread 1: .slice(1) +[ 4, 5, 6, 7 ] +is_full() == true (16B) + + +Threads 2 .. 249 +[ 4*tid .. 4*tid + 3 ] +is_full() == true (16B) + + +Thread 250: .slice(250) +[ 1000, 1001, 1002, -- ] +is_full() == false (3 valid) + + + + + + + +Full-tile fast path (if t_a.is_full()) + +let a_vals = unsafe { t_a.load() }; unsafe { t_c.store(a_vals + b_vals) }; +Emits dialect-cute vector tile load/store → lowers to 128-bit vectorized cute.copy + + +Ragged tail (else) + +t_c.store_at(k, ...) +Plain scalar loop from t_a.base() + diff --git a/cuda-oxide/cuda-oxide-book/cute/layouts-and-tiling.md b/cuda-oxide/cuda-oxide-book/cute/layouts-and-tiling.md new file mode 100644 index 0000000000..0b7885654d --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/layouts-and-tiling.md @@ -0,0 +1,92 @@ +# Layouts, Tiling, and Partitions by Example + +cute-layout and cute-rs let you describe how logical coordinates map to physical memory: whether that is a 1D vector split across threads, a row-wise reduction across warps, or a 2D swizzled shared-memory tile shared by host TMA encoding and device MMA. + +--- + +## 1. Evaluating Layouts with cute-layout + +cute-layout is a `no_std` crate shared by both cute-rs and the compiler. A `Layout` pairs a shape `IntTuple` with a stride `IntTuple` and maps logical coordinates to linear offsets: + +```rust +use cute_layout::{IntTuple, Layout}; + +// Shape (2, 3) with row-major stride (3, 1) +let layout: Layout = "(2,3):(3,1)".parse().unwrap(); +assert_eq!(layout.checked_call(&IntTuple::Leaf(1)), Some(3)); +``` + +In kernel signatures, you write static layouts as Rust type markers (`C`, Rust tuples such as `(A, B)` for nested modes, `Layout`, `RowMajor`, `ColMajor`, `Swizzle`, and `ComposedLayout`). + +--- + +## 2. 1D Tiling and Ragged Tails (zipped_divide and slice) + +`zipped_divide::()` restructures a flat 1D tensor into a 2D view indexed by `(offset_inside_tile, tile_number)`. Calling `.slice(tid)` selects the tile owned by thread `tid`. `local_tile::(tile)` is the one-call form, equal to `zipped_divide::().slice(tile)`: CuTe's `local_tile` for a one-dimensional tiler. + +```{figure} images/zipped-divide-and-slice.svg +:align: center +:width: 100% + +`zipped_divide::<4>()` and `.slice(tid)` partition a 1,003-element vector across 251 threads: threads `0..249` take the vectorized `cute.copy` path, while thread `250` runs a scalar loop over the 3-element tail. +``` + +```rust +let tid = thread::index_1d().get(); +let g_a = cute_rs::Tensor::from_slice(a).zipped_divide::<4>(); +let g_b = cute_rs::Tensor::from_slice(b).zipped_divide::<4>(); +let g_c = cute_rs::TensorMut::from_disjoint_slice(&mut c).zipped_divide::<4>(); + +let t_a = g_a.slice(tid); +let t_b = g_b.slice(tid); +let mut t_c = g_c.slice(tid); + +if t_a.is_full() { + let sum = unsafe { t_a.load() } + unsafe { t_b.load() }; + unsafe { t_c.store(sum) }; +} else { + let mut k = t_a.base(); + while k < a.len() { + unsafe { t_c.store_at(k, a[k] + b[k]) }; + k += 1; + } +} +``` + +--- + +## 3. Shared Host/Device TMA Layout Types + +For 2D and 3D Tensor Memory Accelerator (TMA) transfers, you define the swizzled shared-memory layout type **once** and use it on both the host and the device: + +```{figure} images/host-device-tma-type-agreement.svg +:align: center +:width: 100% + +A single `ComposedLayout` type alias binds host `create_tensor_map_tiled` encoding to device `&TensorMap` and `Sm100SharedTile`. +``` + +```rust +use cute_rs::copy::tma::{TmaEncodeOptions, create_tensor_map_tiled_with_options}; +use cute_rs::{ComposedLayout, RowMajor, Swizzle, TensorMap}; + +// Shared layout type: 128x64 tile with a 128-byte swizzle +pub type ASmem = ComposedLayout, 0, RowMajor<128, 64>>; + +// 1. Host: encode the 128-byte hardware CUtensorMap from `ASmem` +let a_desc: TensorMap = create_tensor_map_tiled_with_options( + a_dev.as_device_ptr(), + m, + k, + k, + TmaEncodeOptions::default(), +)?; + +// 2. Device: receive the descriptor by value via #[grid_constant] +#[kernel] +pub fn gemm_kernel(#[grid_constant] a: &TensorMap) { + // `a` is guaranteed to carry the exact `ASmem` tile shape and swizzle +} +``` + +If a kernel attempts to pass `&TensorMap` into a copy or MMA expecting a different layout or element type, rustc rejects the program at compile time before device codegen even runs. diff --git a/cuda-oxide/cuda-oxide-book/cute/overview.md b/cuda-oxide/cuda-oxide-book/cute/overview.md new file mode 100644 index 0000000000..4fae407277 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/overview.md @@ -0,0 +1,94 @@ +# Overview: cute-rs and ctm-rs (preview) + +:::{admonition} Preview +:class: warning +cute-rs compiles for two targets, `sm_100a` and `sm_120a`, through the pinned CUTLASS 4.7 compiler, and covers a validated subset of CuTeDSL 4.7 on them. The [support matrix](support-matrix.md) lists what is supported, what is partial and what is missing. Expect API changes between releases. +::: + +**cute-rs** brings CUTLASS CuTe's layout algebra, partitions, TMA transfers, pipelines, and tensor-core MMA operations into Rust's type system for two Blackwell targets, without a Python interpreter or runtime JIT tracer. + +```{figure} images/cute-rs-preview-stack-and-flow.svg +:align: center +:width: 100% + +The cute-rs-preview stack: kernels author tensor computation in cute-rs over ctm-rs instruction primitives, import into one Pliron module (dialect-cute, dialect-mir, and dialect-nvvm), pass whole-module verification, and compile through CUTLASS 4.7 MLIR into an embedded cubin. +``` + +--- + +## The Two-Layer Kernel Stack + +Every CuTe kernel in cuda-oxide is built on two complementary device crates plus a shared layout algebra crate: + +| Crate | Layer | What it owns | +| :---- | :---- | :----------- | +| **cute-rs** | Semantic CuTe layer | `Tensor` / `TensorMut`, `zipped_divide`, `slice`, `local_tile`, cooperative/tiled/TMA copies, producer/consumer pipelines, reductions, and tiled MMA (`mma::bf16`, `mma::block_scaled`, `sm100`, `attention::sm100`). | +| **ctm-rs** | Instruction & ABI layer | `#[cuda_module]`, `#[kernel]`, launch contracts, thread/warp/cluster coordinates, `mbarrier` state machines, `tcgen05::tmem::TmemGuard`, and packed register math over cuda-device. | +| **cute-layout** | Shared layout algebra | `IntTuple`, `Layout`, `ComposedLayout`, `Swizzle`, and static compile-time proofs for cooperative copies, `ldmatrix`, and TMA swizzle phase alignment. | + +```{figure} images/three-tiers-comparison.svg +:align: center +:width: 100% + +When to use cute-rs, ctm-rs, and cuda-device. +``` + +--- + +## Your First cute-rs Kernel + +In cute-rs, a `Tensor` pairs storage with a logical coordinate layout. Splitting a vector into 4-element (16-byte) tiles and slicing the tile owned by `tid` takes three lines: + +```rust +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; + +#[cuda_module] +mod kernels { + use super::*; + + #[kernel] + pub fn add_f32(a: &[f32], b: &[f32], mut c: DisjointSlice) { + let tid = thread::index_1d().get(); + let t_a = cute_rs::Tensor::from_slice(a).zipped_divide::<4>().slice(tid); + let t_b = cute_rs::Tensor::from_slice(b).zipped_divide::<4>().slice(tid); + let mut t_c = cute_rs::TensorMut::from_disjoint_slice(&mut c).zipped_divide::<4>().slice(tid); + + if t_a.is_full() { + // Complete 16-byte tile -> lowers to a vectorized cute.copy + unsafe { t_c.store(t_a.load() + t_b.load()) }; + } else { + // Ragged boundary tile -> safe scalar tail loop + let mut k = t_a.base(); + while k < a.len() { + unsafe { t_c.store_at(k, a[k] + b[k]) }; + k += 1; + } + } + } +} +``` + +--- + +## Run and Inspect the Emitted CuTe MLIR + +Install the pinned CUTLASS 4.7 compiler once, then run any `*_cute` example from `cuda-oxide/`: + +```bash +# Install the managed CUTLASS 4.7 compiler +cargo oxide toolchain install cutlass + +# Run the f32 + f16 elementwise example through the CUTLASS MLIR backend +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run elementwise_cute --arch sm_120a + +# Dump the prepared CUTLASS MLIR module to inspect the generated cute ops +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ +CUDA_OXIDE_MLIR_OUTPUT="$PWD/module.mlir" \ + cargo oxide build elementwise_cute --arch sm_120a +``` + +:::{note} +dialect-cute, dialect-mir, and dialect-nvvm coexist inside a **single Pliron module** rather than acting as separate sequential passes. Before exporting MLIR, dialect-cute's whole-module verifier checks tensor origins, pipeline state transitions, and TMA/MMA layout and collector contracts. +::: diff --git a/cuda-oxide/cuda-oxide-book/cute/support-matrix.md b/cuda-oxide/cuda-oxide-book/cute/support-matrix.md new file mode 100644 index 0000000000..63c1f2ae5d --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/cute/support-matrix.md @@ -0,0 +1,155 @@ +# Support Matrix (preview) + +cute-rs is a **preview**. It compiles Rust CuTe kernels ahead of time through the pinned CUTLASS 4.7 compiler for two targets, `sm_100a` and `sm_120a`, and covers a validated subset of what CuTeDSL 4.7 offers on them. This page lists what the subset holds, what it does not, and what cute-rs adds. Status at commit `5eaa6cbda`, 2026-10-22. + +```{figure} images/support-matrix.svg +:align: center +:width: 100% + +One glance: CuTeDSL 4.7 across its architectures against the cute-rs preview on its two targets. +``` + +## Targets + +```{figure} images/support-targets.svg +:align: center +:width: 100% + +The two targets the preview compiles for, and the ones it rejects. +``` + +:::{dropdown} Text version + +| Target | GPU | CuTeDSL 4.7 | cute-rs preview | +| :--- | :--- | :--- | :--- | +| `sm_80`, `sm_86`, `sm_87` | Ampere | yes | **rejected** | +| `sm_89` | Ada | yes | **rejected** | +| `sm_90a` | Hopper | yes | **rejected** | +| `sm_100a` | Blackwell datacenter, B200 class | yes | **yes** | +| `sm_100`, `sm_100f`, `sm_103a`, `sm_110a` | Blackwell datacenter, other variants | yes | **rejected** | +| `sm_120a` | Blackwell GeForce, RTX 50 series | yes | **yes** | +| `sm_120`, `sm_120f`, `sm_121a` | Blackwell GeForce, other variants | yes | **rejected** | + +::: + +The target check runs before any MLIR is built, and the error names the target. Within the pair, `tcgen05` operations need `sm_100a` and MXF4 register MMA needs `sm_120a`. `cargo oxide run` uses the detected GPU, so an RTX 5090 gets `sm_120a`; `cargo oxide build` needs `--arch`. Every CuTe kernel needs `CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir`; the native backend has no CuTe lowering. + +## Feature areas + +The figure at the top of the page is the feature table, with the cute-rs column split by target. Where each row is proven: + +- Layout algebra, static: cute-layout tests, elementwise_cute, blockscale_gemm_cute +- Dynamic layouts, runtime shapes: fmha_prefill_sm120a_cute, fmha_decode_sm120a_cute +- Tensors, partitions, tiled atoms: elementwise_cute, nvfp4_gemv_cute, blockscale_gemm_cute +- TMA copies: blockscale_gemm_cute, fp16_gemm_256x352_cute, sm100_fmha_prefill_cute +- cp.async, bulk, DSMEM copies: fmha_prefill_sm120a_cute +- ldmatrix, stmatrix, atomics: blockscale_gemm_cute, ffn_cute, sm100_mla_decode_fp16_cute +- tcgen05 MMA kinds: fp16_gemm_256x352_cute, sm100_fmha_prefill_cute, sm100_mla_decode_fp16_cute +- Warp MMA kinds (mma.sync): blockscale_gemm_cute, ffn_cute, fmha_prefill_sm120a_cute +- NVFP4 tensor cores (E4M3 scales): qwen3_8b_cute uses the CUDA-core path; no tensor-core producer +- wgmma (Hopper): none +- TMEM: fp16_gemm_256x352_cute, sm100_fmha_prefill_cute +- Pipelines and barriers: blockscale_gemm_cute, fp16_gemm_256x352_cute, sm100_mla_decode_fp16_cute +- Schedulers, persistence, PDL: blockscale_gemm_cute, sm100_fmha_prefill_cute +- Clusters and multicast: fp16_gemm_256x352_cute, sm100_mla_decode_fp16_cute +- Epilogue fusion (EFC): blockscale_gemm_cute, fp16_gemm_256x352_cute, swiglu_cute +- Numerics and intrinsics: rmsnorm_cute, swiglu_cute, qwen_decoder_cute +- Control flow, printf, debugging: every example +- Host: JIT, torch, JAX, AOT export: every example; fp16_gemm_256x352_cute for CUDA graphs +- Autotuning, Task Scheduling, IKET: none +- Verification before codegen: dialect-cute verify tests, 48 cute-rs and 30 ctm-rs compile-fail cases +- Grouped, MoE, mixed-input, FP8: none +- Shipped kernel examples: the 14 crates under cute-rs-preview/examples +- End-to-end models: qwen_decoder_cute, qwen3_8b_cute + +:::{dropdown} Text version of the feature table + +"partial" means a fixed or narrower form of the feature exists; "n/a" means the feature does not apply to those targets. + +| Feature area | CuTeDSL 4.7 | cute-rs preview | What the preview has | +| :--- | :--- | :--- | :--- | +| Layout algebra, static | yes | partial | layouts are Rust types checked at compile time; no layout ops in kernels | +| Dynamic layouts, runtime shapes | yes | partial | runtime lengths and one runtime row pitch; strides and tiles are types | +| Tensors, partitions, tiled atoms | yes | partial | 1D tiles (zipped_divide, slice, 1D local_tile) and fixed 2D block-scaled views; TiledCopy, get_slice and ThrCopy exist but no kernel can reach them yet, so block-wide copies call copy_g2s | +| TMA copies | yes | partial | 2D tile load and store, two-CTA load; no im2col, reduce-store, multicast | +| cp.async, bulk, DSMEM copies | yes | partial | cooperative cp.async with a coverage proof; no bulk or DSMEM data copies | +| ldmatrix, stmatrix, atomics | yes | partial | ldmatrix x4 and x2.trans, stmatrix x2, no-allocate stores; no atomics | +| tcgen05 MMA kinds | yes | partial, sm_100a only | 1 of 11 kinds: F16 in, F32 accumulate, one or two CTAs | +| Warp MMA kinds (mma.sync) | yes | partial | 3 of 7 kinds: MXF4 with UE8M0 scales (sm_120a), F16, BF16 | +| NVFP4 tensor cores (E4M3 scales) | yes | no | NVFP4 projections run on CUDA cores through scaled loads and dot | +| wgmma (Hopper) | yes | n/a | no sm_90a target | +| TMEM | yes | partial, sm_100a only | allocate, 32x32b load and store, FP32 to FP16; no 16xNb, LdRed or S2T | +| Pipelines and barriers | yes | partial | 5 typed pipelines and a two-CTA ring; no CLC, no compiler pipelining | +| Schedulers, persistence, PDL | yes | partial | one static 1D persistent scheduler; no grouped, CLC or stream-K | +| Clusters and multicast | yes | partial, sm_100a only | two-CTA clusters validated; no N-CTA multicast plan | +| Epilogue fusion (EFC) | yes | partial | two fixed FP16 epilogues, bias on sm_100a; activations fused by hand | +| Numerics and intrinsics | yes | partial | f16, f32 native; BF16 and FP4 carriers; no inline PTX or libdevice | +| Control flow, printf, debugging | yes | partial | full Rust control flow and traps, MLIR dumps; no device printf | +| Host: JIT, torch, JAX, AOT export | yes | partial | AOT cubin inside the Rust binary; no JIT, no Python interop | +| Autotuning, Task Scheduling, IKET | yes | no | no autotuner, schedule checker or in-kernel tracing | +| Verification before codegen | yes | yes | whole-module layout, pipeline and role checks; CuTeDSL checks at trace | +| Grouped, MoE, mixed-input, FP8 | yes | no | dense kernels only | +| Shipped kernel examples | yes | partial | 4 sm_100a and 10 sm_120a example crates versus 86 CuTeDSL kernel files | +| End-to-end models | no | sm_100a no, sm_120a yes | Qwen3-TTS 1.7B and Qwen3-8B on sm_120a; CuTeDSL ships none | + +::: + +## Feature additions in CuTeDSL 4.4 to 4.7 + +```{figure} images/support-additions.svg +:align: center +:width: 100% + +What a CuTeDSL user relies on beyond the per-architecture atoms, and what the preview has instead. Epilogue fusion is the row most often asked about. +``` + +:::{dropdown} Text version, with the full comparison + +| Feature | CuTeDSL 4.7 | cute-rs preview | Status | +| :--- | :--- | :--- | :--- | +| Epilogue fusion (EFC) | A Python epilogue function over the accumulator and extra tensors; the GEMM stages its inputs through TMA, runs it per sub-tile with built-in activations, alpha, beta and bias, and stores every output. | Two fixed epilogues: FP32 to FP16 with a row bias on sm_100a, FP32 to FP16 through stmatrix on sm_120a. Activations are fused by hand in register code (swiglu_cute, ffn_cute). No user hook. | partial | +| Dynamic layouts | Fully dynamic layouts; one compiled kernel serves any shape. | Runtime lengths and one runtime row pitch; tile sizes, strides and swizzles are types. One binary covers any M, N, K the static tile divides. | partial | +| JIT and caching | Trace and compile at call time with an on-disk cache. | Ahead of time through rustc and Cargo; nothing to JIT. | n/a | +| AOT export | Export a kernel as a .so plus C header; aarch64 cross build. | The cubin is embedded in the Rust binary. No C export; write an extern C wrapper by hand. | partial | +| TVM FFI, PyTorch, JAX, DLPack | Framework tensors flow in through DLPack; TVM FFI modules for C++, PyTorch and JAX. | None. Buffers live in cuda-core DeviceBuffer; Python appears only in reference harnesses. | no | +| Autotuning | autotune_jit and benchmark helpers search a configuration space and cache the winner. | None. Tile shapes are const generics; tune by instantiating variants in a host harness. | no | +| Task Scheduling | Static analysis of warp-specialized schedules; compilation stops on known concurrency issues (4.7, Primitives-only kernels). | None. The whole-module verifier checks pipeline and epilogue protocol order, not schedules or races. | no | +| Inline PTX and cute.ffi | llvm.inline_asm in kernels; external LLVM bitcode functions. | None on the CuTe path; hardware access goes through the ctm-rs and cuda-device intrinsics. | no | +| printf, print_tensor, IR dumps | Device printf and tensor printing, LaTeX layouts, IR, PTX and SASS dumps. | MLIR, MIR and LLVM dumps and host-side layout Display. No device printf in CuTe kernels. | partial | +| IKET profiling | In-kernel event markers collected into timelines. | None. Use Nsight on the embedded cubin. | no | +| CUDA graphs, PDL, cooperative launch | Graph capture, programmatic dependent launch, cooperative launch, launch events. | Graphs, the PDL launch attribute and cooperative launch from the Rust host. No device-side PDL helpers in kernels. | partial | +| Typed shared memory | SmemAllocator bump allocator, cute.struct, auto-deduced smem size. | Typed static and dynamic shared arrays and layout-typed tiles; carve dynamic shared memory by hand. | partial | +| TensorSSA register vectors | Immutable register vectors with arithmetic, where, conversions and reductions. | RegisterTile with elementwise `+`, `-`, `*`, `/` and unary `-`, map, zip_map, fold and gather; conversions as closures. | partial | +| block_copy | One call issues TMA (multicast hidden) or shared-to-TMEM copies. | A typed two-CTA TMA copy. No shared-to-TMEM copy. | partial | +| Grouped GEMM, MoE, ptr-array | Grouped and MoE persistent schedulers, device-only problem shapes. | None. One launch per group. | no | +| Mixed-input GEMM | Narrow int4 and int8 weights with bf16 or fp8 activations, dequant in the mainloop. | None. Operands share one type. | no | +| Distributed kernels | NVSHMEM all-reduce and GEMM plus collective kernels. | None. Single GPU. | no | +| Grouped block-scaled and blockwise GEMM | Grouped MXFP4 and NVFP4 GEMM, 128x128 blockwise scaling. | Dense block-scaled GEMM and GEMV only. | no | +| FMHA backward | dQ, dK, dV in one SM100 kernel. | Forward only. | no | +| Convolution | TMA im2col copy atom; the 4.7 example tree ships no convolution kernel. | Implicit im2col on the BF16 mma.sync register path in the Qwen codec kernels; no TMA im2col. | partial | +| Sub-byte types | Int4, FP4, two FP6 formats, packed carriers, MXF8F6F4 mixes. | FP4 and scale carriers for block-scaled GEMM and GEMV. No Int4, no FP6. | partial | +| Primitives API (4.7, experimental) | A thin NVVM wrapper beneath CuTe for Tensor Core programming through SIMT. | ctm-rs: typed tcgen05, TMEM, TMA, mbarrier, named barrier, fence, cluster, cp.async and warp primitives for the shipped kernels. No atomics, reductions, wgmma or libdevice. | partial | + +::: + +## What cute-rs adds + +- Ahead-of-time compilation with no Python runtime: one Rust binary carries the cubin. +- Layouts, swizzles, operand roles and pipeline states are Rust types, decoded by the compiler and checked at compile time. +- A whole-module verifier runs before any MLIR is built: tensor provenance, pipeline lifecycle, operand roles, TMA encodability and swizzle alignment. +- Two end-to-end models, Qwen3-TTS 1.7B and Qwen3-8B on sm_120a, with reference comparisons. CuTeDSL 4.7 ships component kernels only. + +## What the compiler does not prove + +The verifier checks what it can see in types and in the module. These stay with the author at every `unsafe` boundary: + +- Thread convergence and lane participation at collective operations. +- Barrier phases and asynchronous completion before a buffer is reused. +- Pointer lifetimes behind raw views and descriptors. +- Disjoint writes across threads. + +## Sources + +- CuTeDSL 4.7 at tag `v4.7.0` of the public CUTLASS repository: the `python/CuTeDSL/cutlass` package, the `examples/python/CuTeDSL` tree and the Python DSL documentation. CuTeDSL calls itself a public beta. +- The target gate: `crates/cuda-oxide-mlir-export/src/profile.rs`. The pinned compiler: `cargo oxide toolchain install cutlass` installs CUTLASS 4.7.0 and checks its digest. +- The figure is generated from the same row data as this table; regenerate both when the compiler surface changes. diff --git a/cuda-oxide/cuda-oxide-book/getting-started/hello-gpu.md b/cuda-oxide/cuda-oxide-book/getting-started/hello-gpu.md index 3a4ff4ec5e..b0a9233ef6 100644 --- a/cuda-oxide/cuda-oxide-book/getting-started/hello-gpu.md +++ b/cuda-oxide/cuda-oxide-book/getting-started/hello-gpu.md @@ -6,10 +6,10 @@ This section walks through installing cuda-oxide, creating a project, writing a ## Install cargo-oxide -If you haven't already, install the build tool with the pinned nightly toolchain: +If you haven't already, install the build tool with the pinned toolchain: ```bash -cargo +nightly-2026-10-11 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide +cargo +1.99.1 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide ``` Verify that your environment is set up correctly: @@ -35,10 +35,12 @@ This generates a ready-to-run project: ```text my_first_kernel/ +├── .cargo/ +│ └── config.toml # sets RUSTC_BOOTSTRAP = "1" for cargo check / rust-analyzer ├── Cargo.toml # dependencies on cuda-device, cuda-host, cuda-core ├── README.md # doctor + run instructions ├── .gitignore # target/ and generated device artifacts -├── rust-toolchain.toml # pins the required nightly toolchain +├── rust-toolchain.toml # pins the required Rust toolchain (1.99.1) └── src/ └── main.rs # kernel + host code in one file ``` @@ -211,9 +213,9 @@ The `--async` flag generates a project with `tokio` and `cuda-async` dependencie Here's the generated async vecadd template (with minor formatting edits for readability): ```rust -use cuda_async::simt::device_context::init_device_contexts; -use cuda_async::simt::device_operation::DeviceOperation; -use cuda_core::simt::LaunchConfig; +use cuda_async::device_context::init_device_contexts; +use cuda_async::device_operation::DeviceOperation; +use cuda_core::LaunchConfig; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; @@ -233,8 +235,8 @@ mod kernels { #[tokio::main] async fn main() -> Result<(), Box> { - use cuda_async::simt::device_box::DeviceBox; - use cuda_core::simt::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; + use cuda_async::device_box::DeviceBox; + use cuda_core::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async}; use std::mem; // 1. Initialize the device context map (default device 0, 1 device). @@ -250,7 +252,7 @@ async fn main() -> Result<(), Box> { // 3. Allocate device memory and copy host data. let (a_dev, b_dev, mut c_dev) = - cuda_async::simt::device_context::with_cuda_context(0, |ctx| { + cuda_async::device_context::with_cuda_context(0, |ctx| { let stream = ctx.default_stream(); let num_bytes = N * mem::size_of::(); unsafe { @@ -261,9 +263,9 @@ async fn main() -> Result<(), Box> { memcpy_htod_async(b, b_host.as_ptr(), num_bytes, stream.cu_stream()).unwrap(); stream.synchronize().unwrap(); ( - DeviceBox::<[f32]>::from_raw_parts(a, N, 0), - DeviceBox::<[f32]>::from_raw_parts(b, N, 0), - DeviceBox::<[f32]>::from_raw_parts(c, N, 0), + DeviceBox::<[f32]>::from_raw_parts(a, N, stream.clone()), + DeviceBox::<[f32]>::from_raw_parts(b, N, stream.clone()), + DeviceBox::<[f32]>::from_raw_parts(c, N, stream.clone()), ) } })?; @@ -283,7 +285,7 @@ async fn main() -> Result<(), Box> { // 5. Copy results back to host. let mut c_host = vec![0.0f32; N]; - cuda_async::simt::device_context::with_cuda_context(0, |ctx| { + cuda_async::device_context::with_cuda_context(0, |ctx| { let stream = ctx.default_stream(); unsafe { memcpy_dtoh_async( diff --git a/cuda-oxide/cuda-oxide-book/getting-started/images/backend-install-by-rename.svg b/cuda-oxide/cuda-oxide-book/getting-started/images/backend-install-by-rename.svg index 22f5473993..27a7d1e396 100644 --- a/cuda-oxide/cuda-oxide-book/getting-started/images/backend-install-by-rename.svg +++ b/cuda-oxide/cuda-oxide-book/getting-started/images/backend-install-by-rename.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/getting-started/installation.md b/cuda-oxide/cuda-oxide-book/getting-started/installation.md index 4a39d2ea3a..7557eb8e13 100644 --- a/cuda-oxide/cuda-oxide-book/getting-started/installation.md +++ b/cuda-oxide/cuda-oxide-book/getting-started/installation.md @@ -20,7 +20,7 @@ without a driver and fails on that call with `CUDA_ERROR_NOT_INITIALIZED`; the error message names the library files the loader tried. A driver whose CUDA major version is older than the toolkit's (a 12.x driver with a 13.x build) fails the same way with a "driver too old" message. -| **Rust** | Nightly (pinned) | Pinned in `cuda-oxide/rust-toolchain.toml` | +| **Rust** | 1.99.1 (pinned) | Pinned in `cuda-oxide/rust-toolchain.toml` | :::{note} cuda-oxide currently targets **Linux only**. Windows is not supported. @@ -32,7 +32,7 @@ cuda-oxide currently targets **Linux only**. Windows is not supported. The repository includes a cuda-oxide devcontainer in `cuda-oxide/.devcontainer/`. Using it is the quickest way to get a reproducible development environment with -CUDA Toolkit 13.0, LLVM 21, Clang 21, and the pinned Rust nightly already +CUDA Toolkit 13.0, LLVM 21, Clang 21, and the pinned Rust toolchain already installed. The host does not need the CUDA Toolkit installed. It does need: @@ -76,7 +76,7 @@ Rust setup sections below. ## Nix / flake.nix The repository also ships `cuda-oxide/flake.nix`, providing a reproducible dev shell -(CUDA 13, LLVM 22, Clang, pinned Rust nightly). Requires +(CUDA 13, Clang/libclang, and the pinned Rust toolchain with LLVM 23 `llvm-tools`). Requires [Nix](https://nixos.org/download/) with flakes enabled, an NVIDIA driver on the host, and Linux (x86\_64 or aarch64). @@ -178,6 +178,30 @@ for the supported driver ranges and PTX limitation. --- +## Official CUTLASS compiler (optional) + +The high-level CuTe translation backend uses NVIDIA's official CUTLASS 4.7 +compiler library. Install the repository-pinned archive with: + +```bash +cargo oxide toolchain install cutlass +``` + +The installer verifies the archive and compiler-library SHA-256 digests and +stores the toolchain in Cargo's managed data directory. Select it per build: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build elementwise_cute --arch sm_120a +``` + +Set `CUDA_OXIDE_CUTLASS_COMPILER` to an absolute +`libCutlassCompiler.so` path only when intentionally overriding the managed +installation. The CuTe chapter's [Build with the CUTLASS Compiler](../cute/build-with-cutlass.md) +page lists the two supported targets, the options and the failure rules. + +--- + ## LLVM 21+ (optional) cuda-oxide uses LLVM's NVPTX backend to lower LLVM IR to PTX. @@ -255,15 +279,15 @@ Validated on Asus GX10 / NVIDIA DGX Spark. ## Rust toolchain The CUDA Oxide workspace ships `cuda-oxide/rust-toolchain.toml`, which pins the -exact nightly version and required components. When you first run a `cargo` +exact stable release and required components. When you first run a `cargo` command inside `cuda-oxide/`, rustup installs the correct toolchain -automatically. Commands at the repository root and in `cutile-rs/` use stable. +automatically. If you need to install it manually: ```bash -rustup toolchain install nightly-2026-10-11 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-10-11 +rustup toolchain install 1.99.1 +rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain 1.99.1 ``` Three of these components are what the codegen backend and doctor need: @@ -282,10 +306,10 @@ The other three are not needed to build: `rust-analyzer` powers IDE support, `cl **Inside the cuda-oxide repo**, it works out of the box via a workspace alias -- no extra install step. -**For use outside the repo** (your own projects), install it with the pinned nightly toolchain: +**For use outside the repo** (your own projects), install it with the pinned toolchain: ```bash -cargo +nightly-2026-10-11 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide +cargo +1.99.1 install --git https://github.com/NVIDIA/cuda-rust.git cargo-oxide ``` On first run, `cargo-oxide` will automatically fetch and build the codegen backend. Subsequent runs reuse the cached build. @@ -346,6 +370,7 @@ The full set, as `cargo oxide --help` reports it: | `doctor` | Check that your environment is set up correctly | | `setup` | Build and cache the codegen backend | | `update` | Refresh the cached codegen backend (or run `setup` inside the workspace) | +| `toolchain` | Install compiler toolchains managed by cargo-oxide | Four of these are worth calling out, because they do something `cargo` itself cannot: @@ -405,8 +430,8 @@ If everything is configured correctly, this compiles a Rust kernel to PTX, launc :::{tip} **Common issues:** -- `No working llc found` -- prefer `rustup component add llvm-tools --toolchain nightly-2026-10-11`, or install LLVM 21+ (`sudo apt install llvm-21`), add `/usr/lib/llvm-21/bin` to your `PATH`, or set `CUDA_OXIDE_LLC=/usr/bin/llc-21`. +- `No working llc found` -- prefer `rustup component add llvm-tools --toolchain 1.99.1`, or install LLVM 21+ (`sudo apt install llvm-21`), add `/usr/lib/llvm-21/bin` to your `PATH`, or set `CUDA_OXIDE_LLC=/usr/bin/llc-21`. - `'stddef.h' file not found` when building host `cuda-bindings` -- install clang dev headers: `sudo apt install clang-21` (or `libclang-common-21-dev`). - `cuda.h not found` -- Set `CUDA_TOOLKIT_PATH` to your CUDA install root, or ensure `/usr/local/cuda/include/cuda.h` exists. -- `rust-src` / `llvm-tools` component missing -- Run `rustup component add rust-src llvm-tools --toolchain nightly-2026-10-11`. +- `rust-src` / `llvm-tools` component missing -- Run `rustup component add rust-src llvm-tools --toolchain 1.99.1`. ::: diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/closures-and-generics.md b/cuda-oxide/cuda-oxide-book/gpu-programming/closures-and-generics.md index 7055e8100c..e63121abb7 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/closures-and-generics.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/closures-and-generics.md @@ -63,7 +63,7 @@ forces the concrete instantiation and lets the loader look up the matching PTX entry point: ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; // SAFETY: config is 1D and both buffers cover every launched index. unsafe { diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/error-handling-and-debugging.md b/cuda-oxide/cuda-oxide-book/gpu-programming/error-handling-and-debugging.md index 56038c8a42..0c34bc5cf1 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/error-handling-and-debugging.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/error-handling-and-debugging.md @@ -158,7 +158,7 @@ The async path (`{kernel}_async` / `DeviceOperation`) uses `DeviceError`, which wraps driver errors alongside context and scheduling failures: ```rust -use cuda_async::simt::error::DeviceError; +use cuda_async::error::DeviceError; let result: Result, DeviceError> = operation.sync(); ``` @@ -488,19 +488,19 @@ cargo oxide doctor Doctor checks: -| Check | What it verifies | -|:----------------|:-----------------------------------------------| -| Rust toolchain | Nightly compiler with required components | -| Codegen backend | `librustc_codegen_cuda.so` built | -| CUDA headers | `cuda.h` present under the toolkit root | -| CUDA toolkit | `nvcc` found and version compatible | -| libNVVM | `libnvvm.so` loadable (libdevice math kernels) | -| nvJitLink | `libnvJitLink.so` loadable (same) | -| libdevice | `libdevice.10.bc` discoverable (same) | -| LLVM | `llc` (21+) available for PTX generation | -| Driver / GPU | `nvidia-smi` reports a GPU and its compute cap | -| cuda-gdb | Optional; only needed for `cargo oxide debug` | -| compute-sanitizer | Optional; only needed for `cargo oxide sanitize` | +| Check | What it verifies | +|:----------------|:----------------------------------------------------| +| Rust toolchain | Pinned compiler (`1.99.1`) with required components | +| Codegen backend | `librustc_codegen_cuda.so` built | +| CUDA headers | `cuda.h` present under the toolkit root | +| CUDA toolkit | `nvcc` found and version compatible | +| libNVVM | `libnvvm.so` loadable (libdevice math kernels) | +| nvJitLink | `libnvJitLink.so` loadable (same) | +| libdevice | `libdevice.10.bc` discoverable (same) | +| LLVM | `llc` (21+) available for PTX generation | +| Driver / GPU | `nvidia-smi` reports a GPU and its compute cap | +| cuda-gdb | Optional; only needed for `cargo oxide debug` | +| compute-sanitizer | Optional; only needed for `cargo oxide sanitize` | The libNVVM / nvJitLink / libdevice checks fire only when a kernel calls CUDA libdevice math (`sin`, `cos`, `exp`, `pow`, `sqrt`, ...). If your diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/execution-model.md b/cuda-oxide/cuda-oxide-book/gpu-programming/execution-model.md index 7694d271f5..476b5bb6bb 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/execution-model.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/execution-model.md @@ -160,7 +160,7 @@ for architecture-specific SM resource limits and occupancy calculations. On the host side, `LaunchConfig` tells the runtime how to shape the grid: ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; // Quick 1D launch: 256 threads per block, enough blocks to cover N elements let cfg = LaunchConfig::for_num_elems(N as u32); diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg b/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg index 9da0b71355..9a0932eb90 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/images/inline-intent-mapping.svg b/cuda-oxide/cuda-oxide-book/gpu-programming/images/inline-intent-mapping.svg index ecbfc22b84..cd048c5329 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/images/inline-intent-mapping.svg +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/images/inline-intent-mapping.svg @@ -8,9 +8,9 @@ --mono: 'IBM Plex Mono', ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; } @media (prefers-color-scheme: dark) { :root, svg { - --bg: #0f151b; --fg: #e4ebf2; --muted: #9fb0c0; --line: #4d5b6a; --panel: #171f28; - --meta: #5fe3d2; --meta-fill: #0f3a37; --code: #f6b544; --code-fill: #3b2a0c; - --rust: #c9a0f0; --rust-fill: #2b1b3f; + --bg: #1b1d23; --fg: #f2f4f7; --muted: #bcc4d0; --line: #5c6475; --panel: #282c36; + --meta: #7ed8bb; --meta-fill: #346357; --code: #f4c474; --code-fill: #735a3c; + --rust: #d8d4fb; --rust-fill: #524e7d; } } diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md b/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md index 8a6b94844f..8d424c5daa 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md @@ -53,7 +53,7 @@ PreparedLaunch -> safe launch of exactly K GPU execution can still overlap the host until you synchronize. ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, cuda_module, kernel, thread}; @@ -532,7 +532,7 @@ explains why this can differ from the normal LLVM-to-PTX path. `LaunchConfig` specifies the grid shape: ```rust -use cuda_core::simt::LaunchConfig; +use cuda_core::LaunchConfig; let config = LaunchConfig { grid_dim: (num_blocks, 1, 1), @@ -592,8 +592,8 @@ instead of enqueuing immediately. No stream is specified at launch time -- the scheduling policy chooses one when the operation is executed: ```rust -use cuda_async::simt::device_context::init_device_contexts; -use cuda_async::simt::device_operation::DeviceOperation; +use cuda_async::device_context::init_device_contexts; +use cuda_async::device_operation::DeviceOperation; init_device_contexts(0, 1)?; let module = kernels::load_async(0)?; diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/memory-and-data-movement.md b/cuda-oxide/cuda-oxide-book/gpu-programming/memory-and-data-movement.md index 282fcf5890..e85dfd126e 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/memory-and-data-movement.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/memory-and-data-movement.md @@ -115,7 +115,7 @@ let results = c_dev.to_host_vec(&stream).unwrap(); | Method | Direction | Description | |:--------------------------------- |:--------------|:-------------------------------| | `from_host(&stream, &[T])` | Host → Device | Allocate + blocking copy | -| `zeroed(&stream, len)` | -- | Allocate + zero-fill | +| `zeroed(&stream, len)` | -- | Allocate + zero-fill + synchronize | | `to_host_vec(&stream)` | Device → Host | Async copy + return `Vec` | | `copy_to_host(&stream, &mut [T])` | Device → Host | Copy into existing slice | | `cu_deviceptr()` | -- | Raw `CUdeviceptr` for FFI | @@ -148,8 +148,8 @@ and executes only after all preceding work on that stream completes. Critically, it does **not** synchronize the device: ```rust -use cuda_async::simt::device_box::DeviceBox; -use cuda_async::simt::device_context::init_device_contexts; +use cuda_async::device_box::DeviceBox; +use cuda_async::device_context::init_device_contexts; init_device_contexts(0, 1)?; // Initialize device context map (default device 0) diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/virtual-memory-and-peer-access.md b/cuda-oxide/cuda-oxide-book/gpu-programming/virtual-memory-and-peer-access.md index d7bbc9a06f..26ea9c6faf 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/virtual-memory-and-peer-access.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/virtual-memory-and-peer-access.md @@ -26,12 +26,12 @@ This chapter covers: :::{seealso} The complete integration test is available in -[`cuda-core/tests/simt_vmm_p2p.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/tests/simt_vmm_p2p.rs). +[`cuda-core/tests/vmm_p2p.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/tests/vmm_p2p.rs). The corresponding implementations are: -* [`cuda-core/src/simt/vmm.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/src/simt/vmm.rs) -* [`cuda-core/src/simt/peer.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/src/simt/peer.rs) +* [`cuda-core/src/vmm.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/src/vmm.rs) +* [`cuda-core/src/peer.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/src/peer.rs) ::: ## VMM and P2P solve different problems @@ -83,7 +83,7 @@ to access each other directly. Always query support before enabling peer access: ```rust -use cuda_core::simt::peer; +use cuda_core::peer; use cuda_core::{CudaContext, DriverError}; fn check_peer_access() -> Result<(), DriverError> { @@ -114,7 +114,7 @@ Use `peer::enable_peer_access` to allow one context to access memory associated with another device: ```rust -use cuda_core::simt::peer; +use cuda_core::peer; use cuda_core::{CudaContext, DriverError}; fn enable_bidirectional_peer_access() -> Result<(), DriverError> { @@ -183,7 +183,7 @@ of the device's allocation granularity. Use `allocation_granularity` and `align_size` before allocating: ```rust -use cuda_core::simt::vmm; +use cuda_core::vmm; use cuda_core::{CudaContext, DriverError}; fn aligned_allocation_size(requested: usize) -> Result { @@ -203,7 +203,7 @@ The following example allocates physical memory, maps it into a virtual address range, writes data through the mapping, and reads it back. ```rust -use cuda_core::simt::vmm; +use cuda_core::vmm; use cuda_core::{CudaContext, DriverError}; fn single_gpu_vmm_roundtrip() -> Result<(), DriverError> { @@ -239,7 +239,7 @@ fn single_gpu_vmm_roundtrip() -> Result<(), DriverError> { let input: Vec = (0..64).map(|i| (i * 3 + 7) as u8).collect(); unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( reservation.base(), input.as_ptr(), input.len(), @@ -252,7 +252,7 @@ fn single_gpu_vmm_roundtrip() -> Result<(), DriverError> { let mut output = vec![0u8; input.len()]; unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( output.as_mut_ptr(), reservation.base(), output.len(), @@ -292,8 +292,8 @@ The following pattern: * Reads the same physical memory through GPU 1's mapping. ```rust -use cuda_core::simt::{peer, vmm}; use cuda_core::{CudaContext, DriverError, IntoResult}; +use cuda_core::{peer, vmm}; use std::mem::MaybeUninit; fn gpu_count() -> Result { @@ -380,7 +380,7 @@ fn cross_gpu_vmm_roundtrip() -> Result<(), DriverError> { ctx0.bind_to_thread()?; unsafe { - cuda_core::simt::memory::memcpy_htod_async( + cuda_core::memory::memcpy_htod_async( reservation0.base(), input.as_ptr(), byte_length, @@ -395,7 +395,7 @@ fn cross_gpu_vmm_roundtrip() -> Result<(), DriverError> { ctx1.bind_to_thread()?; unsafe { - cuda_core::simt::memory::memcpy_dtoh_async( + cuda_core::memory::memcpy_dtoh_async( output.as_mut_ptr(), reservation1.base(), byte_length, diff --git a/cuda-oxide/cuda-oxide-book/index.md b/cuda-oxide/cuda-oxide-book/index.md index 5c15827081..3922b14645 100644 --- a/cuda-oxide/cuda-oxide-book/index.md +++ b/cuda-oxide/cuda-oxide-book/index.md @@ -113,6 +113,10 @@ Compose GPU work as lazy `DeviceOperation` graphs. Schedule across stream pools. Await results with `.await`. ::: +:::{grid-item-card} 🧩 Typed CuTe Kernels (preview) +Encode layouts, swizzles, TMA pipelines, and tensor-core MMA in Rust's type system with [cute-rs](cute/overview.md), compiled through CUTLASS 4.7 for `sm_100a` and `sm_120a`. See the [support matrix](cute/support-matrix.md). +::: + :::: ```{toctree} @@ -163,11 +167,25 @@ async-programming/overlapping-transfers-and-compute ```{toctree} :hidden: :maxdepth: 2 -:caption: Building a Real Application +:caption: Async Host Pipeline Walkthrough projects/async-mlp-pipeline ``` +```{toctree} +:hidden: +:maxdepth: 2 +:caption: CuTe in Rust (preview) + +cute/overview +cute/support-matrix +cute/build-with-cutlass +cute/cutedsl-comparison +cute/layouts-and-tiling +cute/gemm-and-pipelines +cute/contracts +``` + ```{toctree} :hidden: :maxdepth: 2 diff --git a/cuda-oxide/cuda-oxide-book/projects/async-mlp-pipeline.md b/cuda-oxide/cuda-oxide-book/projects/async-mlp-pipeline.md index 7c12483443..4bc1e48ae8 100644 --- a/cuda-oxide/cuda-oxide-book/projects/async-mlp-pipeline.md +++ b/cuda-oxide/cuda-oxide-book/projects/async-mlp-pipeline.md @@ -76,8 +76,8 @@ The `Cargo.toml` pulls in exactly the crates we need: ```toml [dependencies] cuda-device = { path = "../../../cuda-device" } # #[kernel], DisjointSlice, thread::* -cuda-core = "0.3.1" # shared with cutile-rs: CudaModule, simt::LaunchConfig -cuda-async = "0.3.1" # shared with cutile-rs: simt::DeviceOperation, zip!, and_then +cuda-core = "0.3.1" # shared with cutile-rs: CudaModule, LaunchConfig +cuda-async = "0.3.1" # shared with cutile-rs: DeviceOp, zip!, and_then tokio = { version = "1", features = ["rt", "macros"] } ``` @@ -207,81 +207,45 @@ A few patterns that recur across all three kernels: ## Bridging host and device -Three helper functions wrap raw CUDA driver calls into `DeviceOperation` -values. They use `with_context` to receive the scheduler's stream at -execution time — not at the call site. This is the key insight: you build -the recipe now, the scheduler picks the stream later. - -### h2d — host to device +The shared `DeviceBuffer` and `DeviceBox` wrappers transfer ownership without +copying device memory. A host-to-device source must be retained independently +of the operation's output; host combinators can discard outputs before GPU +completion. An ordinary host `Vec` is returned only after its copy completes. ```rust fn h2d(host_data: Vec) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let n = host_data.len(); - let num_bytes = n * mem::size_of::(); - unsafe { - let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memcpy_htod_async(dptr, host_data.as_ptr(), num_bytes, stream.cu_stream()) - .unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) - } + let host_data = Arc::new(host_data); + // Retain before DMA: dropping or forgetting the operation must not + // release the source while the device still reads it. + ctx.retain(host_data.clone()).unwrap(); + let buffer = unsafe { + DeviceBuffer::from_host_async_unchecked(ctx.get_cuda_stream(), &host_data) + }.unwrap(); + value(DeviceBox::from_buffer(buffer)) }) } -``` - -The `host_data` vector is captured by `move`. The closure runs when the -scheduler executes the operation — at that point it has a real CUDA stream. -`malloc_async` and `memcpy_htod_async` are stream-ordered, so the allocation -and copy are serialized on the scheduler's chosen stream. -### zeros — zero-initialized device buffer - -```rust +/// Zero-initialization and later kernels share the same stream and storage. fn zeros(n: usize) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let num_bytes = n * mem::size_of::(); - unsafe { - let dptr = malloc_async(stream.cu_stream(), num_bytes).unwrap(); - memset_d8_async(dptr, 0, num_bytes, stream.cu_stream()).unwrap(); - value(DeviceBox::from_raw_parts(dptr, n, ctx.get_device_id())) - } + value(DeviceBox::from_buffer(DeviceBuffer::zeroed(ctx.get_cuda_stream(), n).unwrap())) }) } -``` -Same pattern as `h2d`, but `memset_d8_async` instead of `memcpy`. The GEMM -kernel uses `beta = 0.0` so the initial contents do not matter, but zeroing -is defensive. - -### d2h — device to host - -```rust +/// Exposing an ordinary Vec to host combinators requires completion before +/// returning it. The shared transfer helper also handles partial-error cleanup. fn d2h(dev: DeviceBox<[f32]>) -> impl DeviceOperation> { device_operation::with_context(move |ctx| { - let stream = ctx.get_cuda_stream(); - let n = dev.len(); - let num_bytes = n * mem::size_of::(); - let mut host = vec![0.0f32; n]; - unsafe { - memcpy_dtoh_async( - host.as_mut_ptr(), dev.cu_deviceptr(), - num_bytes, stream.cu_stream(), - ).unwrap(); - } - let _ = &dev; - value(host) + value(dev.into_buffer().to_host_vec(ctx.get_cuda_stream()).unwrap()) }) } -``` -The `let _ = &dev;` line looks like a no-op, but it keeps `dev` alive in the -closure until the stream synchronizes. Without it, `dev` would be dropped -after the `memcpy_dtoh_async` call but before the stream finishes the copy — -a classic use-after-free on the device side. +``` ---- +These helpers use the execution context's stream, so their GPU work remains +ordered with subsequent Tile or SIMT launches. The `DeviceBox` conversions +preserve the allocation owner, CUDA context and access tracker. ## Composing the pipeline @@ -339,7 +303,7 @@ copies. This is where the magic lives. For each batch, we build a four-stage chain: ```rust - use cuda_async::simt::launch::{AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; + use cuda_async::launch::{AsyncKernelLaunchBuilder, OwnedAsyncKernelLaunch}; let pipeline = zip!(h2d(batch_data), zeros(DIM * DIM), zeros(DIM)) .and_then(move |(input, hidden, output)| { diff --git a/cuda-oxide/cute-rs-preview/.gitignore b/cuda-oxide/cute-rs-preview/.gitignore new file mode 100644 index 0000000000..7abaa9b3ee --- /dev/null +++ b/cuda-oxide/cute-rs-preview/.gitignore @@ -0,0 +1,12 @@ +# Local pipeline diagnostics +tmp_pipeline_cutlass_mlir_sm120a* + +# Local comparison experiment +/examples/fp16_gemm_256x352_cute/test_compare.py + +# Run outputs and measurement records stay out of the tree; the backup branch +# backup/cute-rs-preview-2026-10-21 on the fork holds the records removed in 2026-10. +__pycache__/ +*.pyc +results/ +/models/qwen3_8b_cute/reference/cutedsl_comparison.json diff --git a/cuda-oxide/cute-rs-preview/AGENTS.md b/cuda-oxide/cute-rs-preview/AGENTS.md new file mode 100644 index 0000000000..5d5b23682e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/AGENTS.md @@ -0,0 +1,123 @@ +# cute-rs agent guide + +cute-rs exists to show that Rust kernels written as CuTe operations compile +through the CUTLASS compiler into CuTe kernels. Everything in this guide +follows from that: the kernels must be CuTe, the names must be CuTeDSL's, and +the target is parity with CuTeDSL rather than a second design. + +## Kernel architecture + +Kernels intended to demonstrate or exercise cute-rs (including the Qwen +acceptance workloads and every `*_cute` example) must express tensor +computation through cute-rs tensor views, layouts, partitions, register tiles, +copies, reductions and MMA operations. This is a hard requirement, alongside +numerical correctness and competitive measured performance. + +- A `*_cute` kernel imports `cute_rs` and, for instruction-level steps, + `ctm_rs`. It does not import `cuda_device`. Kernel entry attributes, thread + identifiers and shared-storage declarations come through the `ctm_rs` + re-exports. If a kernel needs something neither crate offers, add it to + cute-rs (or to ctm-rs when it is one instruction) as a typed operation; do + not reach past them. +- The deliverable is the exported MLIR, not a kernel that happens to run. The + tensor computation must reach the CUTLASS compiler as `cute.*` operations; + check with `CUDA_OXIDE_MLIR_OUTPUT=` that the loads, copies and MMA + steps appear as CuTe operations rather than scalar code. A kernel rewritten + around a compiler gap or a missing operation is a workaround, and the + repository rule applies: fix the compiler or add the operation. +- Keep instruction-level CTM operations underneath reusable cute-rs + operations. CTM kernel declarations, launch contracts, thread identifiers and + ABI storage types may remain at kernel entry points. Explicit `*_ctm` + examples continue to demonstrate the lower-level instruction layer. +- Do not hand-code hardware fragment packing, lane-to-value maps, or + shared-memory swizzle formulas in cute-rs workload kernels. Describe those + mappings using typed layouts and derive accesses from them. +- When an abstraction is missing, extend cute-rs with a reusable, typed + operation and test its invariants. Moving a workload-specific implementation + behind a new name, or wrapping a raw instruction while leaving its layout + logic in the kernel, does not satisfy this requirement. +- Keep operand roles, element formats, dimensions, storage bounds and mutable + ownership explicit. State collective-participation and aliasing obligations + at unsafe boundaries; do not imply the type system proves convergence. +- Preserve arithmetic order and observable rounding unless the task explicitly + authorizes changing them. Test layout coverage and tails independently of + device code, then validate GPU numerics, sanitizers and generated + instructions. +- Benchmark representative end-to-end workloads against the preserved + baseline. Better abstractions must not silently sacrifice performance. + Record remaining performance gaps honestly; abstraction migration alone is + not a speedup claim. +- Existing legacy violations are migration work, not a precedent for new code. + Report any remaining violations explicitly when delivering a migration. + +## Names follow CuTeDSL + +A CuTe concept carries the name CuTeDSL 4.7 gives it, so that a CuTeDSL author +reads a cute-rs kernel without a glossary. + +- Look the name up before writing it: the CuTeDSL sources + (`python/CuTeDSL/cutlass/cute`, `cutlass/pipeline`, `cutlass/utils` and + `cutlass/experimental/primitives` in the CUTLASS repository) and the + [comparison page](../cuda-oxide-book/cute/cutedsl-comparison.md). Functions + stay snake_case, types CamelCase; where rustc's `non_snake_case` lint forbids + a CuTeDSL method name such as `partition_A`, the Rust form is `partition_a`. +- Rust mechanics keep Rust idioms: constructors are `new` or `from_*`, `len` + pairs with `is_empty`, `fold` takes an initial value, conversions are + `from_*` and `to_*`, primitive types are `i32`, `f32` and `f16`. A + compile-time plan is a type or a const generic (`RowMajor<128, 64>`, + `Swizzle<3, 3, 3>`, `zipped_divide::<8>()`), because the compiler reads it + from the generic arguments; that is the one place the `make_*` vocabulary + does not carry over. +- ctm-rs uses PTX instruction names (`mbarrier_arrive_expect_tx`, + `barrier_cluster_arrive`, `cp_async_bulk_tensor_shared_cta_global`). The + cute.arch concept names (`thread_idx`, `warp_idx`, `elect_one_sync`, + `cluster_arrive`, `mbarrier_init_fence`) live in `cute_rs::arch` as + forwarders, mirroring cute.arch over primitives in CuTeDSL. +- One name per concept across cute-rs, ctm-rs and the compiler crates. No + abbreviations the DSL does not use (`Fragment`, not `Frag`; `Mxf4`, the + form of the PTX kind and the DSL op), no second name for something that + already has one, and no CUDA C++ or PTX name where cute.arch has a concept + name. +- The compiler recognises cute-rs definitions by path (the importer's + registry and static-configuration constants), some author-facing methods by + name (the importer's classify arms) and types by item name (the codegen + collector). An alias or a forwarding function is free; a rename moves those + constants and their tests in the same change. Callable intrinsics are + forwarded as functions, never re-exported with `pub use`. + +## Parity with CuTeDSL, not reinvention + +The fundamentals have landed. The work now is bringing cute-rs to feature +parity with CuTeDSL, and every addition is measured against it. + +- Before adding an operation, type, attribute, pipeline, scheduler or helper, + find the CuTeDSL counterpart: what it is called, what it accepts, what it + lowers to, and which CuTeDSL example uses it. Implement that concept under + that name. +- The implementation may differ where Rust does better: compile-time plans as + types, typestate for collective and asynchronous operations, sealed roles + for operands and CTA groups, `Result` instead of exceptions. The surface + vocabulary and the semantics still follow CuTeDSL, and a deliberate + deviation is written down in the [comparison page](../cuda-oxide-book/cute/cutedsl-comparison.md). +- The [support matrix](../cuda-oxide-book/cute/support-matrix.md) is the + parity ledger. When a gap closes, update its row and the figure it is drawn + from; when a new operation is only partly covered, say which forms. +- Do not design a cute-rs-only feature where CuTeDSL has one that solves the + same problem; port theirs. Where cute-rs genuinely adds something CuTeDSL + lacks (ahead-of-time compilation, compile-time layout and protocol checks), + document it as an addition, not as the CuTe surface. + +## Example and model documentation + +- Keep each entry README approachable, following `examples/elementwise_cute`: + explain the computation, show the CuTe data flow, give a runnable command, + and describe the expected result. +- Document public and private APIs, including helpers, types, fields and + constants. Explain shapes, units, ownership, readiness, rounding and error + behavior where relevant. Unsafe APIs must state caller obligations; comments + must not claim that types prove thread convergence or asynchronous completion. +- Put numerical contracts, comparison harnesses and reproduction commands in + the workload's `reference/` folder, linked from its README. Do not commit + measurement records: result JSON, traces, profiles, logs and SASS dumps + belong in a run's output directory outside the tree. Keep essential + prerequisites and known user-visible limitations in the entry README. diff --git a/cuda-oxide/cute-rs-preview/README.md b/cuda-oxide/cute-rs-preview/README.md new file mode 100644 index 0000000000..abdf1cfe5b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/README.md @@ -0,0 +1,80 @@ +# cute-rs: CuTe kernels in Rust (preview) + +> **Preview.** `cute-rs` compiles Rust CuTe kernels ahead of time through the pinned +> CUTLASS 4.7 compiler for two targets, `sm_100a` (B200 class) and `sm_120a` +> (RTX 50 series). Every other architecture is rejected before compilation. It covers +> a validated subset of CuTeDSL 4.7, not feature parity; the +> [support matrix](../cuda-oxide-book/cute/support-matrix.md) lists what is in and +> what is not. APIs can change between releases. + +Layouts, swizzles, operand roles, TMA pipelines and tensor-core MMA are Rust types. +The compiler decodes them, verifies the whole module, and hands CUTLASS MLIR to the +pinned compiler. No Python interpreter, no runtime tracer. + +![The cute-rs stack](../cuda-oxide-book/cute/images/cute-rs-preview-stack-and-flow.svg) + +| Crate | Role | +| --- | --- | +| [`cute-rs`](cute-rs/README.md) | tensors, tiles, copies, pipelines and tiled MMA; the kernel-facing API | +| [`ctm-rs`](ctm-rs/README.md) | typed instruction primitives beneath it: TMA, mbarrier, tcgen05, TMEM, mma.sync | +| [`cute-layout`](cute-layout/README.md) | layout algebra and the compile-time proofs shared with the compiler | +| [`dialect-cute`](dialect-cute/README.md) | the Pliron dialect and the whole-module verifier | + +## Run + +```bash +cargo oxide toolchain install cutlass # once +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide run elementwise_cute --arch sm_120a +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide run blockscale_gemm_cute --arch sm_120a +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build fp16_gemm_256x352_cute --arch sm_100a +``` + +`elementwise_cute` checks every `f32` and `f16` bit pattern against a host reference. +The GEMM examples validate against a CPU reference and print their timings. Add +`CUDA_OXIDE_MLIR_OUTPUT="$PWD/module.mlir"` to a `build` to keep the MLIR text. +[Build with the CUTLASS compiler](../cuda-oxide-book/cute/build-with-cutlass.md) has +the options and the failure rules. + +## Examples + +| Example | Target | Shows | +| --- | --- | --- | +| [`elementwise_cute`](examples/elementwise_cute) | `sm_120a`, `sm_100a` | 1D tiles, vectorized copy, scalar tail | +| [`swiglu_cute`](examples/swiglu_cute) | `sm_120a` | register transforms, FP32 math, FP16 output | +| [`rmsnorm_cute`](examples/rmsnorm_cute) | `sm_120a` | warp reductions, explicit BF16 rounding | +| [`ffn_cute`](examples/ffn_cute) | `sm_120a` | BF16 `mma.sync` projections with typed weight roles | +| [`rope_cache_cute`](examples/rope_cache_cute) | `sm_120a` | RoPE, owned KV cache, GQA attention | +| [`nvfp4_gemv_cute`](examples/nvfp4_gemv_cute) | `sm_120a` | packed FP4 with block scales on CUDA cores | +| [`blockscale_gemm_cute`](examples/blockscale_gemm_cute) | `sm_120a` | MXFP4 GEMM: 3-stage TMA ring, tensor cores, TMA-store epilogue | +| [`blockscale_gemm_ctm`](examples/blockscale_gemm_ctm) | `sm_120a` | the same GEMM on `ctm-rs` alone | +| [`fmha_prefill_sm120a_cute`](examples/fmha_prefill_sm120a_cute) | `sm_120a` | FP16 D64 prefill, causal or dense | +| [`fmha_decode_sm120a_cute`](examples/fmha_decode_sm120a_cute) | `sm_120a` | FP16 D64 decode with split-KV reduction | +| [`fp16_gemm_256x352_cute`](examples/fp16_gemm_256x352_cute) | `sm_100a` | two-CTA GEMM: cluster TMA, TMEM, A-collector reuse | +| [`fp16_gemm_256x352_ctm`](examples/fp16_gemm_256x352_ctm) | `sm_100a` | the same GEMM on `ctm-rs` alone | +| [`sm100_fmha_prefill_cute`](examples/sm100_fmha_prefill_cute) | `sm_100a` | FP16 D128 prefill, MHA or GQA, warp specialized | +| [`sm100_mla_decode_fp16_cute`](examples/sm100_mla_decode_fp16_cute) | `sm_100a` | MLA decode with a paged latent cache and two-CTA clusters | + +| Model | Target | Scope | +| --- | --- | --- | +| [`qwen_decoder_cute`](models/qwen_decoder_cute) | `sm_120a` | Qwen3-TTS 1.7B end to end, text in and waveform out ([visual guide](models/qwen_decoder_cute/visual.md)) | +| [`qwen3_8b_cute`](models/qwen3_8b_cute) | `sm_120a` | 36-layer Qwen3-8B decoder in BF16 and NVFP4 ([visual guide](models/qwen3_8b_cute/visual.md)) | + +`cargo oxide list` prints both tables. Each README gives the device requirement, one +command and the expected result; the `reference/` folders hold the comparison +harnesses, the pinned reference sources and the validation commands. + +## Read next + +- [Overview](../cuda-oxide-book/cute/overview.md): the two-layer stack and a first kernel. +- [From CuTeDSL to cute-rs](../cuda-oxide-book/cute/cutedsl-comparison.md): why the kernel signatures differ and how Rust types become CuTe IR. +- [Layouts and tiling](../cuda-oxide-book/cute/layouts-and-tiling.md) and [GEMM and pipelines](../cuda-oxide-book/cute/gemm-and-pipelines.md). +- [Checks and obligations](../cuda-oxide-book/cute/contracts.md): what the compiler proves and what you own. + +## Compiler boundary + +The importer matches `cute-rs` types and methods by their rustc definition path, not by +re-export. Moving a recognised type or inherent `impl` changes its compiler identity; +update the importer decoders and the collector paths together, and keep compiler hooks +in `__compiler` modules. Host tests do not exercise the importer: verify a boundary +change with the identity tests, the compile-fail tests and a device build. +[AGENTS.md](AGENTS.md) states the kernel architecture rules. diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/Cargo.toml b/cuda-oxide/cute-rs-preview/ctm-rs/Cargo.toml new file mode 100644 index 0000000000..05b32e88dc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "ctm-rs" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Typed instruction-level CUDA primitives for cuda-oxide kernels" + +[dependencies] +cuda-device = { workspace = true } +cuda-core = { workspace = true, optional = true } +thiserror = { workspace = true, optional = true } + +[features] +default = [] +# Host-side TMA descriptor encoding is kept out of kernel dependency graphs. +host = ["dep:cuda-core", "dep:thiserror"] + +[dev-dependencies] +# Kernel macro tests also emit host metadata when workspace features unify. +cuda-host = { workspace = true } +trybuild = "1" diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/README.md b/cuda-oxide/cute-rs-preview/ctm-rs/README.md new file mode 100644 index 0000000000..bc0d05e732 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/README.md @@ -0,0 +1,90 @@ +# ctm-rs: typed instruction primitives (preview) + +`ctm-rs` is the instruction layer beneath `cute-rs`: typed MMA operands, physical +shared-memory formats, tensor maps, explicit barriers and collective TMEM +allocation over `cuda-device` intrinsics. It knows nothing about CuTe layouts. + +```text +cute-rs tensors, tiles, pipelines, tiled MMA dialect-cute +ctm-rs instructions with typed operands dialect-nvvm +cuda-device raw intrinsics +``` + +![ctm-rs at a glance](../../cuda-oxide-book/cute/images/ctm-rs-at-a-glance.svg) + +## What it covers + +| Family | Operations | Targets | +| --- | --- | --- | +| Shared memory | `SharedRegion` with checked alignment and capacity, typed CTA and cluster addresses | both | +| TMA | 2D load, two-CTA multicast load, 2D store, descriptor prefetch, bulk commit and wait; host `create_tensor_map_tiled` behind the `host` feature | both; multicast on `sm_100a` | +| Barriers, fences and clusters | mbarrier init, arrive, expect-tx, parity wait; named CTA barrier sync and arrive (`bar` and `barrier.cta`, plain and aligned); `fence.mbarrier_init`, `fence.proxy.async` for `shared::cta` and its two `sync_restrict` forms, `fence.acquire.sync_restrict`; cluster arrive and wait, `mapa`, remote arrive, `CtaMask` | both | +| tcgen05 | `TmemGuard` collective alloc and dealloc, F16 two-CTA MMA through `InstructionDescriptor`, 32x32b TMEM load and store, commit, fences | `sm_100a` | +| Register MMA | m16n8k16 F16 and BF16 with F32 accumulate; m16n8k64 MXF4 with UE8M0 scales; ldmatrix x4, stmatrix x2 | both; MXF4 on `sm_120a` | +| Async copy | 4-byte `cp.async`, noinc barrier arrive, commit and wait groups | both | +| Warp | shuffle (full-warp and masked), elect, vote (all, any, ballot), lane id, active mask, sum reduce, `setmaxnreg` | both | +| Numerics | BF16 round-to-nearest convert, packed f32x2 add, mul and fma, approximate ex2, lg2, rcp, rsqrt and sin, trap | both | + +Not covered: other tcgen05 kinds and TMEM shapes, atomics and reductions, wgmma, +libdevice, arbitrary physical layouts. An unsupported selector fails at export, +not at run time. + +## What the types check, and what you own + +| Checked at compile time | Yours at every `unsafe` boundary | +| --- | --- | +| Operand roles, element formats, MMA shapes | Which lanes take part in a collective | +| Descriptor fields and packed address ranges | Barrier phases and arrival counts | +| Borrowed TMEM views and shared regions | Asynchronous completion before reuse | +| Alignment and capacity of adopted storage | Collector fill and release order, pointer lifetimes | + +Destructors neither synchronize nor free hardware resources. + +## Use + +```rust,ignore +use ctm_rs::{groups::CtaGroup2, tcgen05::{descriptors::*, mma::{self, collector_a}}}; + +type Tile = DenseSharedF16; +let descriptor = InstructionDescriptor::::new(); +// dst, a and b carry matching group, format and role types. +mma::mma::(&dst, &a.k16::<0>(), &b.k16::<0>(), &descriptor, false); +// Commit, wait and hand off to the consumers before reuse or readback. +``` + +[`fp16_gemm_256x352_ctm`](../examples/fp16_gemm_256x352_ctm) (`sm_100a`) and +[`blockscale_gemm_ctm`](../examples/blockscale_gemm_ctm) (`sm_120a`) use ctm-rs alone +through the primitives profile, `CUDA_OXIDE_MLIR_PROFILE=cutlass-primitives-mlir22`. +The SM100 attention examples use it beside `cute-rs` in one module under the default +profile. + +The crate is `no_std`. Kernel macros and canonical device types are reexported; +crates that use them still depend on `cuda-device` directly because the macro +expansions name it. + +## Modules + +| Modules | Responsibility | +| --- | --- | +| `memory`, `mma` | storage formats and typed MMA operands | +| `tcgen05` | descriptors, MMA issue, TMEM ownership and completion | +| `async_copy`, `tma` | shared-memory copies and tensor-map operations | +| `mbarrier`, `barrier`, `fence`, `cluster`, `groups` | mbarrier readiness, named CTA barriers, memory fences, cluster operations, CTA-group types | +| `warp`, `thread` | collectives, coordinates, register limits | +| `convert`, `float`, `f32x2`, `debug` | conversion, arithmetic and debugging | +| `host::tma` | host tensor-map encoding behind the `host` feature | + +## Check + +From `cuda-oxide/` with its pinned toolchain: + +```sh +cargo test -p ctm-rs --all-targets +cargo test -p ctm-rs --all-targets --features host +python3 scripts/check-primitives-api-compile.py --example fp16 --output-dir /tmp/ctm-fp16 +``` + +The tests cover formats, carriers, storage checks and compile-time rejection of +invalid contracts. The compile check builds the fixtures through the CUTLASS +consumer and records the MLIR and a validated cubin; it launches nothing. GPU +correctness and sanitizer commands live in each example's README. diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/async_copy.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/async_copy.rs new file mode 100644 index 0000000000..c4d17cc6c8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/async_copy.rs @@ -0,0 +1,105 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit completion controls for ordinary global-to-shared `cp.async`. +//! These operate on the executing thread's copy groups, separately from TMA +//! bulk-copy groups. Kernel authors supply CTA synchronization before sharing +//! completed destinations with other threads. + +use crate::mbarrier::{GeneralBarrier, ManagedBarrier, Ready}; +use crate::memory::SharedAddress; + +/// Copy one aligned 32-bit global word into the calling CTA's shared memory. +/// +/// # Safety +/// `source` must refer to global memory and remain unchanged until this copy +/// completes, even after its Rust borrow ends. Both operands must remain live +/// through completion. The destination is four-byte aligned shared storage +/// with no conflicting access. The issuing thread must supply matching copy +/// completion; synchronize threads before another thread reads or reuses the +/// destination. Custom completion assembly must include a compiler memory clobber. +/// +/// ```compile_fail +/// # let source = 0u32; +/// unsafe { ctm_rs::async_copy::cp_async_ca_4(core::ptr::null_mut(), &source) }; +/// ``` +#[inline(always)] +pub unsafe fn cp_async_ca_4(destination: SharedAddress<'_, u32>, source: &u32) { + unsafe { cuda_device::async_copy::cp_async_ca_4(destination.as_raw(), source) } +} + +/// Signal this thread's outstanding copies through a ready local barrier. +/// The instruction consumes an already expected arrival without increasing it. +/// +/// # Safety +/// This thread must be an expected participant in this live barrier phase. +/// Keep the barrier and all associated copy buffers live until completion; +/// callers coordinate every participant and phase explicitly. +/// +/// ```compile_fail +/// use ctm_rs::mbarrier::{GeneralBarrier, ManagedBarrier, Uninit}; +/// # fn rejected(barrier: &ManagedBarrier) { +/// unsafe { ctm_rs::async_copy::cp_async_mbarrier_arrive_noinc(barrier) }; +/// # } +/// ``` +#[inline(always)] +pub unsafe fn cp_async_mbarrier_arrive_noinc( + barrier: &ManagedBarrier, +) { + unsafe { cuda_device::async_copy::cp_async_mbarrier_arrive_noinc(barrier.as_ptr().cast_mut()) } +} + +/// Commit this thread's uncommitted asynchronous copies as one group. +/// +/// # Safety +/// Copies must satisfy their source/destination lifetime and alignment +/// contracts. Commit does not wait for completion or synchronize threads. +#[inline(always)] +pub unsafe fn cp_async_commit_group() { + unsafe { cuda_device::async_copy::cp_async_commit_group() } +} + +/// Wait until at most `N` of this thread's recent committed groups are pending. +/// +/// `N` is an instruction immediate in `0..=7`. This inserts neither a CTA +/// barrier nor a wait for other threads' copies. +/// +/// # Safety +/// Access only destinations whose copy groups this wait completes. Synchronize +/// threads before another thread accesses a completed destination. All copy +/// source and destination allocations must remain live until completion. +/// +/// ```compile_fail +/// unsafe { ctm_rs::async_copy::cp_async_wait_group::<8>() }; +/// ``` +#[inline(always)] +pub unsafe fn cp_async_wait_group() { + const { assert!(N <= 7, "cp.async wait count must be in 0..=7") } + unsafe { cuda_device::async_copy::cp_async_wait_group(N) } +} + +/// Issue one four-byte global-to-shared copy with the cache-all policy (sm_80+). +/// No commit, wait or thread synchronization is inserted. +/// +/// # Safety +/// Both pointers must be four-byte aligned and cover four live bytes in their +/// respective address spaces. Keep both allocations live until completion; +/// coordinate destination access and finish the copy before reading it. +#[inline(always)] +pub unsafe fn cp_async_ca_4_raw(shared_dst: *mut u32, global_src: *const u32) { + unsafe { cuda_device::async_copy::cp_async_ca_4(shared_dst, global_src) } +} + +/// Attach this thread's preceding `cp.async` completion to a barrier (sm_80+). +/// This does not increment the barrier's expected arrival count or wait. +/// +/// # Safety +/// `barrier` must name a live, aligned, initialized CTA-shared mbarrier. Its +/// phase must already expect this completion arrival. Keep the barrier and +/// copy allocations live until completion and coordinate phase transitions. +#[inline(always)] +pub unsafe fn cp_async_mbarrier_arrive_noinc_raw(barrier: *mut crate::mbarrier::Barrier) { + unsafe { cuda_device::async_copy::cp_async_mbarrier_arrive_noinc(barrier) } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/barrier.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/barrier.rs new file mode 100644 index 0000000000..58b7c9dcb7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/barrier.rs @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! The `bar.*` and `barrier.cta.*` instruction family: named CTA barriers. +//! +//! Each function forwards one instruction with the caller's barrier ID and +//! thread count; the aligned forms require every non-exited thread of a warp +//! to execute with identical control flow. `bar.warp.sync` is a warp +//! collective and stays in [`crate::warp::bar_warp_sync`]; mbarrier objects +//! live in [`crate::mbarrier`] and memory fences in [`crate::fence`]. + +/// Signal arrival at a numbered CTA barrier without waiting. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +#[inline(always)] +pub unsafe fn barrier_cta_arrive(barrier_id: u32, thread_count: u32) { + unsafe { cuda_device::barrier::barrier_cta_arrive(barrier_id, thread_count) } +} + +/// Signal arrival at a numbered CTA barrier in aligned mode without waiting. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +/// Every non-exited thread in the warp must execute with identical control flow. +#[inline(always)] +pub unsafe fn barrier_cta_arrive_aligned(barrier_id: u32, thread_count: u32) { + unsafe { cuda_device::barrier::barrier_cta_arrive_aligned(barrier_id, thread_count) } +} + +/// Synchronize the requested number of threads at a numbered CTA barrier. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +#[inline(always)] +pub unsafe fn barrier_cta_sync(barrier_id: u32, thread_count: u32) { + unsafe { cuda_device::barrier::barrier_cta_sync(barrier_id, thread_count) } +} + +/// Synchronize at a numbered CTA barrier in aligned mode. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +/// Every non-exited thread in the warp must execute with identical control flow. +#[inline(always)] +pub unsafe fn barrier_cta_sync_aligned(barrier_id: u32, thread_count: u32) { + unsafe { cuda_device::barrier::barrier_cta_sync_aligned(barrier_id, thread_count) } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/cluster.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/cluster.rs new file mode 100644 index 0000000000..b26900b9a7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/cluster.rs @@ -0,0 +1,200 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Explicit cluster membership, shared-address mapping and remote arrival. +use super::{ + groups::CtaGroup, + memory::{ClusterSharedAddress, SharedAddress}, +}; +use core::marker::PhantomData; +use cuda_device::barrier::Barrier; + +/// Generate forwarding functions for cluster special registers. +macro_rules! coordinate { + ($($name:ident),* $(,)?) => {$( + #[doc = concat!("Read the cluster coordinate through [`cuda_device::cluster::", stringify!($name), "`].")] + #[allow(non_snake_case)] + #[inline(always)] + pub fn $name() -> u32 { + cuda_device::cluster::$name() + } + )*}; +} +coordinate!( + cluster_ctaidX, + cluster_ctaidY, + cluster_ctaidZ, + cluster_nctaidX, + cluster_nctaidY, + cluster_nctaidZ, + cluster_idx, + cluster_idx_x, + cluster_idx_y, + cluster_idx_z, + num_clusters, + cluster_grid_dim_x, + cluster_grid_dim_y, + cluster_grid_dim_z, + cluster_size, +); + +/// Read this CTA's linear rank in its cluster, the value PTX `%cluster_ctarank` holds. +/// +/// Forwards [`cuda_device::cluster::block_rank`] under the CuTeDSL primitives name. +#[inline(always)] +pub fn cluster_ctarank() -> u32 { + cuda_device::cluster::block_rank() +} + +/// Arrive at the cluster barrier with release ordering. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. +#[inline(always)] +pub unsafe fn barrier_cluster_arrive() { + unsafe { cuda_device::cluster::barrier_cluster_arrive() } +} + +/// Arrive at the cluster barrier in aligned mode with release ordering. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. Every non-exited thread in the warp must +/// execute with identical control flow. +#[inline(always)] +pub unsafe fn barrier_cluster_arrive_aligned() { + unsafe { cuda_device::cluster::barrier_cluster_arrive_aligned() } +} + +/// Arrive at the cluster barrier without a release guarantee. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. Publish earlier memory accesses with the +/// required cluster-scope fence first. +#[inline(always)] +pub unsafe fn barrier_cluster_arrive_relaxed() { + unsafe { cuda_device::cluster::barrier_cluster_arrive_relaxed() } +} + +/// Arrive at the cluster barrier in aligned mode without a release guarantee. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. Every non-exited thread in the warp must +/// execute with identical control flow. Publish earlier memory accesses with +/// the required cluster-scope fence first. +#[inline(always)] +pub unsafe fn barrier_cluster_arrive_relaxed_aligned() { + unsafe { cuda_device::cluster::barrier_cluster_arrive_relaxed_aligned() } +} + +/// Wait at the cluster barrier with acquire ordering. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute this matching wait. +#[inline(always)] +pub unsafe fn barrier_cluster_wait() { + unsafe { cuda_device::cluster::barrier_cluster_wait() } +} + +/// Wait at the cluster barrier in aligned mode with acquire ordering. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute this matching wait. Every non-exited thread in the warp must +/// execute with identical control flow. +#[inline(always)] +pub unsafe fn barrier_cluster_wait_aligned() { + unsafe { cuda_device::cluster::barrier_cluster_wait_aligned() } +} + +/// A nonempty multicast mask. Group mode does not restrict cluster membership. +pub struct CtaMask { + /// Nonzero set of destination CTA ranks, one bit per rank. + bits: u16, + /// Instruction group selector; does not limit which ranks the mask can name. + group: PhantomData, +} +impl Copy for CtaMask {} +impl Clone for CtaMask { + fn clone(&self) -> Self { + *self + } +} +/// Why a CTA multicast mask cannot be issued. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MaskError { + /// No destination CTA is selected. + Empty, +} +impl CtaMask { + /// Adopt a nonzero set of destination-rank bits; group mode does not cap membership. + #[inline(always)] + pub const fn try_from_bits(bits: u16) -> Result { + if bits == 0 { + Err(MaskError::Empty) + } else { + Ok(Self { + bits, + group: PhantomData, + }) + } + } + /// Return the mask bits consumed by the multicast instruction. + #[inline(always)] + pub const fn bits(self) -> u16 { + self.bits + } +} +/// Map local shared storage to the same offset in a cluster peer (PTX `mapa`). +/// +/// # Safety +/// `rank` must name an active CTA in this cluster with matching live storage. +/// Remote initialization, readiness and lifetime must be coordinated explicitly. +#[inline(always)] +pub unsafe fn mapa<'s, T>(address: SharedAddress<'s, T>, rank: u32) -> ClusterSharedAddress<'s, T> { + unsafe { + ClusterSharedAddress::from_raw(cuda_device::cluster::map_shared_rank_mut( + address.as_raw(), + rank, + )) + } +} +/// Issue one remote arrival without waiting or changing a software phase. +/// +/// # Safety +/// The remote barrier must be initialized and expecting this arrival. All +/// participants must agree on the phase and keep its storage live. +#[inline(always)] +pub unsafe fn arrive_remote(address: ClusterSharedAddress<'_, Barrier>) { + unsafe { + let offset = cuda_device::shared::cvta_generic_to_cluster_shared_offset( + address.as_raw().cast::(), + ); + cuda_device::barrier::mbarrier_arrive_cluster(offset) + } +} + +/// Arrive at a remote barrier without publishing preceding memory accesses. +/// +/// The barrier count and execution dependency retain cluster scope. Unlike +/// [`arrive_remote`], this operation has no release memory ordering. +/// +/// # Safety +/// The remote barrier must be initialized and expecting this arrival. All +/// participants must agree on the phase and keep its storage live. The caller +/// must supply the proxy or TMEM fences required by its data-handoff protocol; +/// this relaxed arrival alone does not make preceding data accesses visible. +#[inline(always)] +pub unsafe fn arrive_remote_relaxed(address: ClusterSharedAddress<'_, Barrier>) { + unsafe { + let offset = cuda_device::shared::cvta_generic_to_cluster_shared_offset( + address.as_raw().cast::(), + ); + cuda_device::barrier::mbarrier_arrive_cluster_relaxed(offset) + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/convert.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/convert.rs new file mode 100644 index 0000000000..d07aa1132b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/convert.rs @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit instruction packing and numeric conversion. + +/// Convert two FP32 values to packed FP16 using nearest-even rounding. +/// +/// `lo` occupies bits `0..16`, and `hi` occupies bits `16..32`. This preserves +/// the argument order of `cuda_device::convert::cvt_f16x2_f32` and does not +/// introduce flush-to-zero mode or an additional conversion. +#[must_use] +#[inline(always)] +pub fn cvt_f16x2_f32(lo: f32, hi: f32) -> u32 { + cuda_device::convert::cvt_f16x2_f32(lo, hi) +} + +/// Round FP32 to raw BF16 storage bits with round-to-nearest, ties-to-even. +/// Signed zeros and infinities are preserved; NaNs remain NaNs. +#[must_use] +#[inline(always)] +pub fn f32_to_bf16_rne(value: f32) -> u16 { + cuda_device::convert::f32_to_bf16_rne(value) +} + +/// Widen raw BF16 storage bits to FP32 exactly, retaining NaN payload bits. +#[must_use] +#[inline(always)] +pub fn bf16_to_f32(bits: u16) -> f32 { + cuda_device::convert::bf16_to_f32(bits) +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/debug.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/debug.rs new file mode 100644 index 0000000000..0ff028dd7c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/debug.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit device execution failure for violated kernel contracts. + +/// Abort device execution and report an interrupt to the host. +#[inline(always)] +pub fn trap() -> ! { + cuda_device::debug::trap() +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/f32x2.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/f32x2.rs new file mode 100644 index 0000000000..89864fd861 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/f32x2.rs @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! SM100 packed FP32 register arithmetic, distinct from integer addresses. + +/// Two FP32 lanes occupying one 64-bit register, without memory-vector alignment. +#[repr(transparent)] +#[derive(Clone, Copy, Debug)] +pub struct PackedF32x2(u64); + +/// Pack scalar lanes without changing NaN payloads or signed zero. +#[must_use] +#[inline(always)] +pub fn pack_f32x2(lo: f32, hi: f32) -> PackedF32x2 { + PackedF32x2(cuda_device::f32x2::pack_f32x2(lo, hi)) +} + +/// Recover the low lane without arithmetic conversion. +#[must_use] +#[inline(always)] +pub fn unpack_f32x2_lo(value: PackedF32x2) -> f32 { + cuda_device::f32x2::unpack_f32x2_lo(value.0) +} + +/// Recover the high lane without arithmetic conversion. +#[must_use] +#[inline(always)] +pub fn unpack_f32x2_hi(value: PackedF32x2) -> f32 { + cuda_device::f32x2::unpack_f32x2_hi(value.0) +} + +/// Add matching FP32 lanes with round-to-nearest-even on SM100+. +/// +/// ```compile_fail +/// ctm_rs::f32x2::add_f32x2(0_u64, 0_u64); +/// ``` +#[must_use] +#[inline(always)] +pub fn add_f32x2(a: PackedF32x2, b: PackedF32x2) -> PackedF32x2 { + PackedF32x2(cuda_device::f32x2::add_f32x2(a.0, b.0)) +} + +/// Multiply matching FP32 lanes with round-to-nearest-even on SM100+. +#[must_use] +#[inline(always)] +pub fn mul_f32x2(a: PackedF32x2, b: PackedF32x2) -> PackedF32x2 { + PackedF32x2(cuda_device::f32x2::mul_f32x2(a.0, b.0)) +} + +/// Fused multiply-add each FP32 lane, rounding once per lane, on SM100+. +#[must_use] +#[inline(always)] +pub fn fma_f32x2(a: PackedF32x2, b: PackedF32x2, c: PackedF32x2) -> PackedF32x2 { + PackedF32x2(cuda_device::f32x2::fma_f32x2(a.0, b.0, c.0)) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Swapped lanes, numeric casts and NaN canonicalization are observable. + #[test] + fn packed_lanes_preserve_storage_bits() { + let patterns = [ + 0_u32, + 0x8000_0000, + 1, + 0x007f_ffff, + 0x0080_0000, + 0x3f80_0000, + 0x7f80_0000, + 0xff80_0000, + 0x7fc1_2345, + 0x7f81_2345, + ]; + for low in patterns { + for high in patterns { + let packed = pack_f32x2(f32::from_bits(low), f32::from_bits(high)); + assert_eq!(unpack_f32x2_lo(packed).to_bits(), low); + assert_eq!(unpack_f32x2_hi(packed).to_bits(), high); + } + } + assert_eq!(core::mem::size_of::(), 8); + assert_eq!( + core::mem::align_of::(), + core::mem::align_of::() + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/fence.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/fence.rs new file mode 100644 index 0000000000..0f0b8125ab --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/fence.rs @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! The `fence.*` instruction family: `membar`/`fence` memory-consistency +//! instructions, including `fence.mbarrier_init` and the `fence.proxy.async` +//! forms. +//! +//! A fence orders memory accesses between threads, scopes or proxies; none of +//! these waits for asynchronous completion or substitutes for a barrier. + +/// Publish preceding mbarrier initialization at cluster scope. +/// +/// PTX `fence.mbarrier_init.release.cluster` is the only form of this fence, +/// so the name carries no ordering or scope suffix. +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract of +/// [`cuda_device::barrier::fence_mbarrier_init_release_cluster`]. +#[inline(always)] +pub unsafe fn fence_mbarrier_init() { + unsafe { cuda_device::barrier::fence_mbarrier_init_release_cluster() } +} + +/// Make preceding CTA-shared generic-proxy writes visible to the async proxy. +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract of +/// [`cuda_device::barrier::fence_proxy_async_shared_cta`]. +#[inline(always)] +pub unsafe fn fence_proxy_async_shared_cta() { + unsafe { cuda_device::barrier::fence_proxy_async_shared_cta() } +} + +/// Acquire cluster-shared async-proxy writes through the generic proxy (sm_90+). +/// +/// PTX `fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster`. +/// +/// # Safety +/// Pair this fence with the compatible cluster-scoped release and barrier +/// protocol. It neither waits for asynchronous operations nor orders global memory. +#[inline(always)] +pub unsafe fn fence_proxy_async_acquire_sync_restrict() { + unsafe { cuda_device::barrier::fence_proxy_async_generic_acquire_shared_cluster_cluster() } +} + +/// Release CTA-shared generic writes to the async proxy at cluster scope (sm_90+). +/// +/// PTX `fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster`. +/// +/// # Safety +/// Pair this fence with the compatible cluster-scoped acquire and barrier +/// protocol, and keep the shared storage live through its asynchronous uses. +#[inline(always)] +pub unsafe fn fence_proxy_async_release_sync_restrict() { + unsafe { cuda_device::barrier::fence_proxy_async_generic_release_shared_cta_cluster() } +} + +/// Acquire cluster-shared memory at cluster scope (sm_90+), excluding global memory. +/// +/// PTX `fence.acquire.sync_restrict::shared::cluster.cluster`; CuTeDSL's +/// `fence_sync_restrict(order)` covers acquire and release, and only the +/// acquire form has an intrinsic here. +/// +/// # Safety +/// Pair this with a compatible cluster-scoped synchronization protocol. This +/// does not substitute for a proxy fence or asynchronous completion wait. +#[inline(always)] +pub unsafe fn fence_sync_restrict_acquire() { + unsafe { cuda_device::barrier::fence_acquire_shared_cluster_cluster() } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/float.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/float.rs new file mode 100644 index 0000000000..720bffc529 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/float.rs @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Instruction-specific approximate FP32 math. +//! The flush-to-zero variants remain distinct so kernels retain their chosen +//! treatment of subnormal values and hardware approximation behavior. + +/// Compute `2^value` with `ex2.approx.f32`, without flush-to-zero mode. +#[must_use] +#[inline(always)] +pub fn ex2_approx_f32(value: f32) -> f32 { + cuda_device::float::ex2_approx_f32(value) +} + +/// Compute `2^value` with `ex2.approx.ftz.f32`, flushing subnormal values to zero. +#[must_use] +#[inline(always)] +pub fn ex2_approx_ftz_f32(value: f32) -> f32 { + cuda_device::float::ex2_approx_ftz_f32(value) +} + +/// Compute a reciprocal with `rcp.approx.ftz.f32`, flushing subnormal values to zero. +#[must_use] +#[inline(always)] +pub fn rcp_approx_ftz_f32(value: f32) -> f32 { + cuda_device::float::rcp_approx_ftz_f32(value) +} + +/// Compute reciprocal square root with `rsqrt.approx.f32`, without flush-to-zero mode. +#[must_use] +#[inline(always)] +pub fn rsqrt_approx_f32(value: f32) -> f32 { + cuda_device::float::rsqrt_approx_f32(value) +} + +/// Compute base-two logarithm with `lg2.approx.f32`, without flush-to-zero mode. +#[must_use] +#[inline(always)] +pub fn lg2_approx_f32(value: f32) -> f32 { + cuda_device::float::lg2_approx_f32(value) +} + +/// Compute sine with `sin.approx.f32`, without flush-to-zero mode. +/// +/// This exposes the PTX approximation, not a full-range libdevice sine. PTX +/// bounds absolute error by 2^-20.5 on [-2π,2π] and 2^-14.7 on [-100π,100π]; +/// outside the latter interval its accuracy is best effort. Reduce the angle +/// before this instruction when an application needs the smaller error bound. +#[must_use] +#[inline(always)] +pub fn sin_approx_f32(value: f32) -> f32 { + cuda_device::float::sin_approx_f32(value) +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/groups.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/groups.rs new file mode 100644 index 0000000000..91ca3259ee --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/groups.rs @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Static CTA-group selectors. They do not elect issuing threads. +/// Limits allocation dispatch to the supported one- and two-CTA modes. +mod sealed { + /// Implemented only by the supported CTA-group selectors. + pub trait Group {} +} +/// Instruction selector for one participating CTA. +pub enum CtaGroup1 {} +/// Instruction selector for a cooperating pair of CTAs. +pub enum CtaGroup2 {} +/// Reviewed hardware CTA-group modes. +pub trait CtaGroup: sealed::Group { + #[doc(hidden)] + /// # Safety + /// A complete issuing warp in each required CTA must participate uniformly. + unsafe fn alloc(slot: *mut u32, columns: u32); + #[doc(hidden)] + /// # Safety + /// All accesses must be complete and the required warps must participate. + unsafe fn dealloc(address: u32, columns: u32); +} +/// Bind a group selector to its literal allocation and deallocation intrinsics. +macro_rules! group { + ($g:ty,$alloc:ident,$dealloc:ident) => { + impl sealed::Group for $g {} + impl CtaGroup for $g { + #[inline(always)] + unsafe fn alloc(slot: *mut u32, columns: u32) { + unsafe { cuda_device::tcgen05::$alloc(slot, columns) } + } + #[inline(always)] + unsafe fn dealloc(address: u32, columns: u32) { + unsafe { cuda_device::tcgen05::$dealloc(address, columns) } + } + } + }; +} +group!(CtaGroup1, tcgen05_alloc, tcgen05_dealloc); +group!(CtaGroup2, tcgen05_alloc_cg2, tcgen05_dealloc_cg2); diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/host/mod.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/host/mod.rs new file mode 100644 index 0000000000..00b1d3ceec --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/host/mod.rs @@ -0,0 +1,8 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Host utilities for preparing typed tensor and memory operations. + +pub mod tma; diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/host/tma.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/host/tma.rs new file mode 100644 index 0000000000..58a5babdad --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/host/tma.rs @@ -0,0 +1,357 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Host encoding for the explicit Rust primitives tensor-map API. + +use crate::{memory::PhysicalFormat, tma::TensorMap}; +use core::{ffi::c_void, mem::MaybeUninit}; +use cuda_core::sys as driver; +use cuda_device::TmaDescriptor; + +/// Restricts driver element codes to formats supported by this encoder. +mod sealed { + /// Implemented only for the storage elements with reviewed CUDA driver codes. + pub trait Element {} +} +/// Element types supported by the reviewed tensor-map encoder. +pub trait TmaElement: sealed::Element { + /// CUDA tensor-map element code used by the host driver. + const DRIVER_TYPE: driver::CUtensorMapDataType; +} +impl sealed::Element for f16 {} +impl TmaElement for f16 { + const DRIVER_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16; +} +impl sealed::Element for u8 {} +impl TmaElement for u8 { + const DRIVER_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT8; +} +impl sealed::Element for u16 {} +impl TmaElement for u16 { + const DRIVER_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT16; +} + +/// L2 cache-fetch promotion requested when encoding the descriptor. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum L2Promotion { + /// Use the driver’s default fetch behavior without promotion. + #[default] + None, + /// Request a 128-byte L2 fetch. + B128, +} +/// Optional cache policy; tensor shape and swizzle come from the physical format. +#[derive(Clone, Copy, Debug, Default)] +pub struct EncodeOptions { + /// L2 promotion passed to cuTensorMapEncodeTiled. + pub l2_promotion: L2Promotion, +} +/// Invalid tensor-map arguments or a CUDA driver encoding failure. +#[derive(Debug, thiserror::Error)] +pub enum EncodeError { + /// The base is null or not aligned to 16 bytes. + #[error("tensor-map base must be non-null and 16-byte aligned")] + InvalidBase, + /// A global extent is zero or exceeds 2^32 elements. + #[error("tensor-map extents must be between 1 and 2^32 inclusive")] + InvalidExtent, + /// Row pitch is too short, misaligned, overflowing, or outside the driver range. + #[error("row stride must cover the row, be 16-byte aligned and be less than 2^40 bytes")] + InvalidStride, + /// The selected tile shape and swizzle cannot be encoded by this API. + #[error("unsupported physical tile format")] + InvalidFormat, + /// The CUDA driver rejected descriptor encoding. + #[error("cuTensorMapEncodeTiled failed with CUDA status {0}")] + Driver(#[doc = "CUDA driver status returned by the encoder."] driver::CUresult), +} + +/// Validated arguments for one rank-two CUDA tensor-map encoding. +struct EncodeParams { + /// CUDA driver code for one storage element. + data_type: driver::CUtensorMapDataType, + /// Global extents in elements, ordered as columns then rows. + dimensions: [u64; 2], + /// Byte distance between adjacent global rows. + strides: [u64; 1], + /// Shared tile extents in elements, ordered as columns then rows. + box_dimensions: [u32; 2], + /// Sampling step within each logical dimension; both are one. + element_strides: [u32; 2], + /// Shared-memory byte permutation selected by the physical format. + swizzle: driver::CUtensorMapSwizzle, + /// Requested L2 fetch promotion passed to the driver. + l2: driver::CUtensorMapL2promotion, +} + +/// Check address, extent, pitch and physical-format limits before calling the driver. +fn validate_2d>( + base: usize, + rows: u64, + cols: u64, + ld_elements: u64, + options: EncodeOptions, +) -> Result { + if base == 0 || !base.is_multiple_of(16) { + return Err(EncodeError::InvalidBase); + } + if rows == 0 || cols == 0 || rows > 1u64 << 32 || cols > 1u64 << 32 { + return Err(EncodeError::InvalidExtent); + } + let stride = ld_elements + .checked_mul(size_of::() as u64) + .ok_or(EncodeError::InvalidStride)?; + if ld_elements < cols || stride % 16 != 0 || stride >= 1u64 << 40 { + return Err(EncodeError::InvalidStride); + } + if F::ROWS == 0 || F::ROWS > 256 || F::COLS == 0 || F::COLS > 256 { + return Err(EncodeError::InvalidFormat); + } + let row_bytes = F::COLS + .checked_mul(size_of::()) + .ok_or(EncodeError::InvalidFormat)?; + let swizzle = match F::SWIZZLE_BYTES { + 0 if row_bytes % 16 == 0 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_NONE, + 64 if row_bytes == 64 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_64B, + 128 if row_bytes == 128 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, + _ => return Err(EncodeError::InvalidFormat), + }; + Ok(EncodeParams { + data_type: T::DRIVER_TYPE, + dimensions: [cols, rows], + strides: [stride], + box_dimensions: [F::COLS as u32, F::ROWS as u32], + element_strides: [1, 1], + swizzle, + l2: match options.l2_promotion { + L2Promotion::None => { + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE + } + L2Promotion::B128 => { + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B + } + }, + }) +} + +// Only a successful encoder may initialize and publish these opaque words. +/// Publish opaque descriptor bytes only after the supplied driver encoder succeeds. +fn encode_with>( + encode: impl FnOnce(*mut driver::CUtensorMap) -> driver::CUresult, +) -> Result, EncodeError> { + const { + assert!(size_of::() == 128); + assert!(align_of::() >= 128); + } + let mut output = MaybeUninit::::uninit(); + let result = encode(output.as_mut_ptr()); + if result != driver::cudaError_enum_CUDA_SUCCESS { + return Err(EncodeError::Driver(result)); + } + let output = unsafe { output.assume_init() }; + Ok(unsafe { TensorMap::from_encoded(TmaDescriptor::from_raw_words(output.opaque)) }) +} + +/// Encode a row-major global tensor for a reviewed physical shared format. +/// +/// This is the CuTeDSL host name `create_tensor_map_tiled`; the data type, box +/// and swizzle come from `T` and `F` instead of explicit arguments. +/// Coordinates use the inner (column) dimension first. Encoding does not +/// retain the allocation: it must cover the specified extents and remain live +/// until every asynchronous device operation using the returned map completes. +// The driver encodes this device address; it does not access the tensor data. +// Actual tensor-map accesses are unsafe device operations with lifetime contracts. +#[allow(clippy::not_unsafe_ptr_arg_deref)] +pub fn create_tensor_map_tiled>( + global_base: *mut c_void, + rows: u64, + cols: u64, + ld_elements: u64, + options: EncodeOptions, +) -> Result, EncodeError> { + let p = validate_2d::(global_base as usize, rows, cols, ld_elements, options)?; + encode_with(|output| unsafe { + driver::cuTensorMapEncodeTiled( + output, + p.data_type, + 2, + global_base, + p.dimensions.as_ptr(), + p.strides.as_ptr(), + p.box_dimensions.as_ptr(), + p.element_strides.as_ptr(), + driver::CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, + p.swizzle, + p.l2, + driver::CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + }) +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::{F16Tile128, F16TileB80, F16TileB96, F16TileC}; + + /// Host encoding must keep packed values, scale words and output elements in their correct driver formats. + #[test] + fn blockscaled_tma_arguments_preserve_packed_bytes_scales_and_output() { + use crate::memory::{Mxf4ScaleTile, Mxf4Tile128}; + + let packed = + validate_2d::(0x1000, 256, 128, 128, EncodeOptions::default()) + .unwrap(); + assert_eq!( + packed.data_type, + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT8 + ); + assert_eq!(packed.dimensions, [128, 256]); + assert_eq!(packed.strides, [128]); + assert_eq!(packed.box_dimensions, [64, 128]); + assert_eq!(packed.element_strides, [1, 1]); + assert_eq!( + packed.swizzle, + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_64B + ); + + let scales = + validate_2d::(0x1000, 2, 512, 512, EncodeOptions::default()) + .unwrap(); + assert_eq!( + scales.data_type, + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT16 + ); + assert_eq!(scales.dimensions, [512, 2]); + assert_eq!(scales.strides, [1024]); + assert_eq!(scales.box_dimensions, [256, 1]); + assert_eq!(scales.element_strides, [1, 1]); + assert_eq!( + scales.swizzle, + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_NONE + ); + + let output = + validate_2d::(0x1000, 256, 256, 256, EncodeOptions::default()) + .unwrap(); + assert_eq!( + output.data_type, + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16 + ); + assert_eq!(output.dimensions, [256, 256]); + assert_eq!(output.strides, [512]); + assert_eq!(output.box_dimensions, [64, 128]); + assert_eq!( + output.swizzle, + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B + ); + } + + /// Row pitches must be checked after converting packed storage elements to bytes. + #[test] + fn blockscaled_tma_strides_are_validated_in_bytes() { + use crate::memory::{Mxf4ScaleTile, Mxf4Tile128}; + + for pitch in [63, 65, 1u64 << 40, u64::MAX] { + assert!(matches!( + validate_2d::(0x1000, 128, 64, pitch, EncodeOptions::default()), + Err(EncodeError::InvalidStride) + )); + } + for pitch in [255, 257, 1u64 << 39, u64::MAX] { + assert!(matches!( + validate_2d::(0x1000, 1, 256, pitch, EncodeOptions::default()), + Err(EncodeError::InvalidStride) + )); + } + } + + /// A failing driver callback must not publish its uninitialized descriptor output. + #[test] + fn tma_driver_failure_does_not_adopt_uninitialized_output() { + let result = encode_with::(|output| { + assert_eq!((output as usize) % 128, 0); + driver::cudaError_enum_CUDA_ERROR_INVALID_VALUE + }); + assert!(matches!( + result, + Err(EncodeError::Driver( + driver::cudaError_enum_CUDA_ERROR_INVALID_VALUE + )) + )); + } + + /// Validated encoder parameters must preserve tile extents, swizzle spans and cache options. + #[test] + fn tma_arguments_preserve_physical_shapes_and_options() { + let p = + validate_2d::(0x1000, 256, 64, 64, EncodeOptions::default()).unwrap(); + assert_eq!(p.dimensions, [64, 256]); + assert_eq!(p.strides, [128]); + assert_eq!(p.box_dimensions, [64, 128]); + assert_eq!(p.element_strides, [1, 1]); + assert_eq!( + p.swizzle, + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B + ); + let b = validate_2d::( + 0x1000, + 192, + 64, + 64, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .unwrap(); + assert_eq!(b.box_dimensions, [64, 96]); + assert_eq!( + b.l2, + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B + ); + assert_eq!( + validate_2d::(0x1000, 160, 64, 64, EncodeOptions::default()) + .unwrap() + .box_dimensions, + [64, 80] + ); + let c = + validate_2d::(0x1000, 128, 352, 352, EncodeOptions::default()).unwrap(); + assert_eq!(c.box_dimensions, [32, 128]); + assert_eq!( + c.swizzle, + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_64B + ); + } + + /// Invalid bases, extents and pitches must fail validation before invoking the CUDA driver. + #[test] + fn tma_rejects_invalid_dynamic_layout_before_driver_call() { + for (base, rows, cols, pitch) in [ + (0, 256, 64, 64), + (0x1001, 256, 64, 64), + (0x1000, 0, 64, 64), + (0x1000, 256, 0, 64), + (0x1000, (1 << 32) + 1, 64, 64), + (0x1000, 256, (1 << 32) + 1, 1 << 33), + (0x1000, 256, 64, 63), + (0x1000, 256, 64, 65), + (0x1000, 256, 64, 1 << 39), + (0x1000, 256, 64, u64::MAX), + ] { + assert!( + validate_2d::(base, rows, cols, pitch, EncodeOptions::default()) + .is_err() + ); + } + assert!( + validate_2d::(0x1000, 1 << 32, 64, 64, EncodeOptions::default()) + .is_ok() + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/lib.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/lib.rs new file mode 100644 index 0000000000..d1811db881 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/lib.rs @@ -0,0 +1,167 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Instruction-level CUDA primitives with statically checked operands. +//! +//! Physical formats and instruction selectors are Rust types. Election, +//! synchronization, instruction order and asynchronous completion are explicit +//! responsibilities of the kernel author, documented on unsafe operations. +//! +//! Register carriers can be constructed without issuing an instruction: +//! +//! ``` +//! use ctm_rs::mma::AccumulatorTile16x8; +//! let tile = AccumulatorTile16x8::<2, 8>::zero(); // 2 × 8 fragments per lane. +//! assert_eq!(tile.fragment::<1, 7>().registers(), [0.0; 4]); +//! ``` +//! +//! Setup, coverage and the safety contract are in +//! `cuda-oxide/cute-rs-preview/ctm-rs/README.md`. +//! Enable `host` for `host::tma` descriptor encoding. +//! Users of the reexported kernel macros must also depend directly on +//! `cuda-device`: their expansions refer to `::cuda_device`. + +#![feature(f16)] +#![no_std] + +#[cfg(feature = "host")] +extern crate std; + +#[cfg(feature = "host")] +pub mod host; + +/// Kernel macros using cuda-device's launch ABI; require a direct `cuda-device` dependency. +pub use cuda_device::{cluster_launch, cuda_module, kernel, launch_bounds, launch_contract}; + +pub mod async_copy; +pub mod barrier; +pub mod cluster; +pub mod convert; +pub mod debug; +pub mod f32x2; +pub mod fence; +pub mod float; +pub mod groups; +pub mod mbarrier; +pub mod memory; +pub mod mma; +pub mod tcgen05; +pub mod tma; +pub mod warp; + +// Keep callable APIs as forwarding functions instead of public function +// reexports. rustc may print a reexport's path as an intrinsic's name; the +// device compiler recognizes raw operations by their cuda-device paths. The +// scoped-index and unroll compiler markers below are intentional exceptions: +// their canonical definitions must survive macro and importer recognition. + +/// Hardware thread/block coordinates, without inferred thread mapping. +pub mod thread { + /// Canonical index witness and compiler markers, preserving their defining crate. + /// + /// `index_1d` calls must appear inside a kernel/device body for the macro to + /// attach its lifetime scope. `__unroll_config` is recognized at its loop + /// call site. Wrapping either in an ordinary forwarding body is incorrect. + pub use cuda_device::thread::{__unroll_config, ThreadIndex, index_1d}; + + // Scope rewriting preserves the user's `thread::` prefix. Reexport the + // canonical helpers so that generated calls retain the original witness. + #[doc(hidden)] + pub use cuda_device::thread::__internal; + + /// Generate named forwarding functions while preserving the intrinsic’s defining path. + macro_rules! coordinate { + ($($name:ident),* $(,)?) => {$( + #[doc = concat!("Read the hardware coordinate through [`cuda_device::thread::", stringify!($name), "`].")] + #[allow(non_snake_case)] + #[inline(always)] + pub fn $name() -> u32 { + cuda_device::thread::$name() + } + )*}; + } + coordinate!( + blockDim_x, + blockDim_y, + blockDim_z, + blockIdx_x, + blockIdx_y, + blockIdx_z, + gridDim_x, + gridDim_y, + gridDim_z, + threadIdx_x, + threadIdx_y, + threadIdx_z, + ); + + /// Restricts register redistribution counts to instruction immediates. + mod sealed { + /// Implemented only for multiples of eight in 24..=256. + pub trait RegisterCount {} + } + /// Sealed proof that `N` is a multiple of eight in `24..=256`. + pub trait ValidRegisterCount: sealed::RegisterCount {} + /// Admit the register-budget immediates supported by setmaxnreg. + macro_rules! register_counts { + ($($n:literal),* $(,)?) => {$( + impl sealed::RegisterCount<$n> for () {} + impl ValidRegisterCount<$n> for () {} + )*}; + } + register_counts!( + 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, + 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, + ); + + /// Increase a warpgroup's register limit to an immediate count. + /// `N` must be a multiple of eight in `24..=256`. + /// + /// # Safety + /// Every thread in the warpgroup must execute the same operation and count + /// on supported hardware. The current register limit must be at most `N`. + /// Synchronize the warpgroup before another register-limit adjustment. + /// Coordinate releases by other warpgroups to avoid waiting indefinitely + /// for unavailable registers. + /// + /// ```compile_fail,E0277 + /// unsafe { ctm_rs::thread::setmaxnreg_inc::<16>() }; + /// ``` + /// + /// ```compile_fail,E0277 + /// unsafe { ctm_rs::thread::setmaxnreg_inc::<33>() }; + /// ``` + #[inline(always)] + pub unsafe fn setmaxnreg_inc() + where + (): ValidRegisterCount, + { + unsafe { cuda_device::thread::setmaxnreg_inc::() } + } + + /// Decrease a warpgroup's register limit to an immediate count. + /// `N` must be a multiple of eight in `24..=256`. + /// + /// # Safety + /// Every thread in the warpgroup must execute the same operation and count + /// on supported hardware. The current register limit must be at least `N`. + /// Synchronize the warpgroup before another register-limit adjustment. + /// Subsequent code must respect the selected limit. + /// + /// ```compile_fail,E0277 + /// unsafe { ctm_rs::thread::setmaxnreg_dec::<264>() }; + /// ``` + /// + /// ```compile_fail,E0133 + /// ctm_rs::thread::setmaxnreg_dec::<32>(); + /// ``` + #[inline(always)] + pub unsafe fn setmaxnreg_dec() + where + (): ValidRegisterCount, + { + unsafe { cuda_device::thread::setmaxnreg_dec::() } + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/mbarrier.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/mbarrier.rs new file mode 100644 index 0000000000..4a1c21d8a6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/mbarrier.rs @@ -0,0 +1,160 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! The `mbarrier.*` instruction family: canonical barrier states with explicit +//! instruction-only lifecycle methods, plus the raw init, arrive, expect-tx +//! and parity-test forms on CTA-shared barrier objects. +//! +//! Named CTA barriers (`bar.*`, `barrier.cta.*`) live in [`crate::barrier`]; +//! memory fences (`fence.*`) in [`crate::fence`]. +use crate::memory::SharedAddress; +pub use cuda_device::barrier::{ + Barrier, BarrierToken, GeneralBarrier, Invalidated, ManagedBarrier, Ready, TmaBarrier, Uninit, +}; + +/// Borrow typed shared addresses from initialized barrier handles. +pub trait BarrierAddressExt { + /// Borrow a current-CTA shared address without extending this handle's lifetime. + fn shared_address(&self) -> SharedAddress<'_, Barrier>; +} + +impl BarrierAddressExt for ManagedBarrier { + #[inline(always)] + fn shared_address(&self) -> SharedAddress<'_, Barrier> { + unsafe { SharedAddress::from_raw(self.as_ptr().cast_mut()) } + } +} + +/// Issue one arrival, returning the hardware phase token. +/// +/// # Safety +/// This participant must be expected in the current live barrier phase. +#[inline(always)] +pub unsafe fn arrive(bar: &ManagedBarrier) -> BarrierToken { + bar.arrive() +} +/// Issue one arrival and add the expected TMA byte count. Arrival decrement is one. +/// +/// # Safety +/// The issuing participant must be expected, and `bytes` must match the +/// asynchronous transfers completing against this barrier phase. +#[inline(always)] +pub unsafe fn arrive_expect_tx( + bar: &ManagedBarrier, + bytes: u32, +) -> BarrierToken { + bar.arrive_expect_tx(bytes) +} +/// Try the supplied phase parity once; no loop or phase update is inserted. +/// +/// # Safety +/// Parity must be 0 or 1 and agree with the participant's intended barrier phase. +#[inline(always)] +pub unsafe fn try_wait_parity( + bar: &ManagedBarrier, + parity: u32, +) -> bool { + bar.try_wait_parity(parity) +} + +/// Arrive once and add expected TMA bytes with relaxed cluster scope. +/// +/// # Safety +/// One expected participant issues this for the current phase. `bytes` must +/// cover the asynchronous transfers completing at this live local barrier; +/// initialization must already be published to every cluster participant. +#[inline(always)] +pub unsafe fn arrive_expect_tx_cluster( + bar: &ManagedBarrier, + bytes: u32, +) -> BarrierToken { + BarrierToken::from_raw(unsafe { + cuda_device::barrier::mbarrier_arrive_expect_tx_cluster(bar.as_ptr(), 1, bytes) + }) +} + +/// Test parity once with cluster-scope acquire ordering; no software wait loop. +/// +/// # Safety +/// The live local barrier must track this participant's intended phase. +/// `parity` must be zero or one; coordinate remote participants explicitly. +#[inline(always)] +pub unsafe fn try_wait_parity_cluster( + bar: &ManagedBarrier, + parity: u32, +) -> bool { + unsafe { cuda_device::barrier::mbarrier_try_wait_parity_cluster(bar.as_ptr(), parity) } +} + +/// Initialize one raw CTA-shared barrier, without election or publication (sm_80+). +/// +/// # Safety +/// `bar` must name aligned live CTA-shared storage. Exactly one thread must +/// initialize it with a nonzero expected count before any participant uses it. +/// Publish initialization and synchronize participants explicitly. +#[inline(always)] +pub unsafe fn mbarrier_init(bar: *mut Barrier, expected_count: u32) { + unsafe { cuda_device::barrier::mbarrier_init(bar, expected_count) } +} + +/// Issue one raw CTA-scoped release arrival and return its phase token (sm_80+). +/// +/// # Safety +/// `bar` must name a live initialized CTA-shared barrier whose current phase +/// expects this arrival. Coordinate phase changes and keep it live through use. +#[inline(always)] +pub unsafe fn mbarrier_arrive(bar: *const Barrier) -> u64 { + unsafe { cuda_device::barrier::mbarrier_arrive(bar) } +} + +/// Issue one release arrival and add expected transaction bytes (sm_90+). +/// `tx_count` retains the canonical ABI; the instruction decrements arrivals by one. +/// +/// # Safety +/// `bar` must name a live initialized CTA-shared barrier expecting this arrival. +/// Pass `tx_count = 1`; `bytes` must match the asynchronous transfers completing +/// at this phase. Keep all resources live until those transfers complete. +#[inline(always)] +pub unsafe fn mbarrier_arrive_expect_tx(bar: *const Barrier, tx_count: u32, bytes: u32) -> u64 { + unsafe { cuda_device::barrier::mbarrier_arrive_expect_tx(bar, tx_count, bytes) } +} + +/// Issue one relaxed cluster-scope arrival with expected transaction bytes (sm_90+). +/// `tx_count` retains the canonical ABI; the instruction decrements arrivals by one. +/// +/// # Safety +/// `bar` must name a live initialized local CTA-shared barrier expecting this +/// arrival. Pass `tx_count = 1` and match `bytes` to this phase's transfers. +/// Publish initialization and provide required data/proxy ordering explicitly; +/// this relaxed operation establishes no data-access ordering. +#[inline(always)] +pub unsafe fn mbarrier_arrive_expect_tx_cluster( + bar: *const Barrier, + tx_count: u32, + bytes: u32, +) -> u64 { + unsafe { cuda_device::barrier::mbarrier_arrive_expect_tx_cluster(bar, tx_count, bytes) } +} + +/// Test raw CTA-shared barrier parity once; no wait loop is inserted (sm_90+). +/// +/// # Safety +/// `bar` must name a live initialized local barrier and `parity` must be zero +/// or one for the intended phase. Coordinate phase transitions explicitly. +#[inline(always)] +pub unsafe fn mbarrier_try_wait_parity(bar: *const Barrier, parity: u32) -> bool { + unsafe { cuda_device::barrier::mbarrier_try_wait_parity(bar, parity) } +} + +/// Test local barrier parity once with relaxed cluster ordering (sm_90+). +/// +/// # Safety +/// `bar` must name a live initialized local CTA-shared barrier; `parity` must +/// be zero or one for the intended phase. Publish initialization explicitly. +/// A successful test orders execution/counts only; provide the required data, +/// proxy and TMEM ordering before accessing the completed phase's resources. +#[inline(always)] +pub unsafe fn mbarrier_try_wait_parity_cluster_relaxed(bar: *const Barrier, parity: u32) -> bool { + unsafe { cuda_device::barrier::mbarrier_try_wait_parity_cluster_relaxed(bar, parity) } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/memory.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/memory.rs new file mode 100644 index 0000000000..62784714a5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/memory.rs @@ -0,0 +1,285 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Instruction-visible storage formats, without tensor layout algebra. + +use core::{cell::UnsafeCell, marker::PhantomData}; +use cuda_device::SharedArray; + +// Keep the canonical storage types and defining crate used by the importer +// and launch-contract macros. Reexports do not introduce new storage markers. +/// Canonical sixteen-byte aligned register-word carrier for vector memory accesses. +pub use cuda_device::vector::U32x4; +pub use cuda_device::{CuSimd, DisjointSlice, DynamicSharedArray}; + +/// Read one aligned 32-bit storage word without interpreting its packed lanes. +/// +/// This ordinary memory read preserves every payload bit. It does not establish +/// synchronization, change address spaces, or perform numeric conversion. +/// +/// # Safety +/// +/// `address` must be four-byte aligned, point to four initialized readable bytes +/// within one live allocation, and obey Rust's aliasing rules. The caller must +/// establish producer readiness and prevent concurrent mutation for the read. +#[inline(always)] +pub unsafe fn load_aligned_u32(address: *const u32) -> u32 { + // SAFETY: alignment, allocation extent, initialization and readiness are + // explicit caller obligations; u32 accepts every possible storage pattern. + unsafe { address.read() } +} + +/// Restricts physical storage formats to the instruction layouts reviewed here. +mod sealed { + /// Implemented only for supported element, tile and swizzle combinations. + pub trait Format {} +} + +/// Physical storage contract implemented only for reviewed instruction formats. +/// Extents count storage elements in one CTA; byte sizes and alignment are explicit. +pub trait PhysicalFormat: sealed::Format { + /// Storage element moved by the instruction; may pack narrower logical values. + type Element; + /// Rows in the physical tile held by one CTA. + const ROWS: usize; + /// Storage elements in the tile’s contiguous dimension. + const COLS: usize; + /// Bytes that must remain accessible for this complete tile. + const FOOTPRINT_BYTES: usize; + /// Required shared allocation alignment in bytes. + const BASE_ALIGNMENT: usize; + /// Swizzle span in bytes, or zero for unswizzled storage. + const SWIZZLE_BYTES: usize; +} + +/// Select the hardware 128-byte shared swizzle. +pub enum Swizzle128 {} +/// Select the hardware 64-byte shared swizzle. +pub enum Swizzle64 {} +/// A per-CTA FP16 tile. Only the four aliases below implement [`PhysicalFormat`]. +/// +/// | Alias | Rows × columns | Swizzle span | Base alignment | +/// | --- | --- | --- | --- | +/// | [`F16Tile128`] | 128 × 64 | 128 bytes | 1024 bytes | +/// | [`F16TileB96`] | 96 × 64 | 128 bytes | 1024 bytes | +/// | [`F16TileB80`] | 80 × 64 | 128 bytes | 1024 bytes | +/// | [`F16TileC`] | 128 × 32 | 64 bytes | 512 bytes | +pub struct Swizzled2D( + #[doc = "Carries the swizzle choice without runtime storage."] PhantomData, +); +/// Local A tile for group-two M256 MMA; see [`Swizzled2D`]. +pub type F16Tile128 = Swizzled2D<128, 64, Swizzle128>; +/// Local B tile for group-two N192 MMA; see [`Swizzled2D`]. +pub type F16TileB96 = Swizzled2D<96, 64, Swizzle128>; +/// Local B tile for group-two N160 MMA; see [`Swizzled2D`]. +pub type F16TileB80 = Swizzled2D<80, 64, Swizzle128>; +/// Local FP16 output tile for a 32-column epilogue slice; see [`Swizzled2D`]. +pub type F16TileC = Swizzled2D<128, 32, Swizzle64>; + +/// Define one FP16 tile footprint and its byte-alignment requirement. +macro_rules! format { + ($rows:literal, $cols:literal, $swizzle:ty, $bytes:literal, $alignment:literal) => { + impl sealed::Format for Swizzled2D<$rows, $cols, $swizzle> {} + impl PhysicalFormat for Swizzled2D<$rows, $cols, $swizzle> { + type Element = f16; + const ROWS: usize = $rows; + const COLS: usize = $cols; + const FOOTPRINT_BYTES: usize = $rows * $cols * 2; + const BASE_ALIGNMENT: usize = $alignment; + const SWIZZLE_BYTES: usize = $bytes; + } + }; +} +format!(128, 64, Swizzle128, 128, 1024); +format!(96, 64, Swizzle128, 128, 1024); +format!(80, 64, Swizzle128, 128, 1024); +format!(128, 32, Swizzle64, 64, 512); + +/// A 128-by-64 byte tile holding two packed FP4 values per byte. +/// +/// The inner TMA coordinate and row width count bytes, not FP4 values. +pub enum Mxf4Tile128 {} +impl sealed::Format for Mxf4Tile128 {} +impl PhysicalFormat for Mxf4Tile128 { + type Element = u8; + const ROWS: usize = 128; + const COLS: usize = 64; + const FOOTPRINT_BYTES: usize = 8192; + const BASE_ALIGNMENT: usize = 512; + const SWIZZLE_BYTES: usize = 64; +} + +/// One unswizzled 512-byte scale tile, transported as 256 packed `u16` elements. +/// +/// The kernel is responsible for packing scale values in instruction order. +pub enum Mxf4ScaleTile {} +impl sealed::Format for Mxf4ScaleTile {} +impl PhysicalFormat for Mxf4ScaleTile { + type Element = u16; + const ROWS: usize = 1; + const COLS: usize = 256; + const FOOTPRINT_BYTES: usize = 512; + const BASE_ALIGNMENT: usize = 128; + const SWIZZLE_BYTES: usize = 0; +} + +/// Byte displacement used when encoding a shared-memory instruction address. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ByteOffset( + #[doc = "Unsigned byte displacement; validated by the descriptor constructor."] usize, +); +impl ByteOffset { + /// Wrap a byte offset without checking instruction-specific alignment. + pub const fn new(bytes: usize) -> Self { + Self(bytes) + } + /// Return the displacement in bytes. + pub const fn bytes(self) -> usize { + self.0 + } +} + +/// Define an address wrapper that records its shared-memory domain and borrow. +macro_rules! address { + ($name:ident, $domain:literal) => { + #[doc = $domain] + #[repr(transparent)] + pub struct $name<'s, T> { + /// Raw address in the domain named by this wrapper. + pointer: *mut T, + /// Retains the shared allocation's lifetime without exclusive ownership. + storage: PhantomData<&'s UnsafeCell>, + } + impl Copy for $name<'_, T> {} + impl Clone for $name<'_, T> { + fn clone(&self) -> Self { + *self + } + } + impl<'s, T> $name<'s, T> { + /// Adopt an address without issuing a hardware operation. + /// + /// # Safety + /// The pointer must belong to this address domain and refer to live, + /// suitably aligned storage for `'s`. This does not establish + /// initialization, exclusive access or synchronization. + #[inline(always)] + pub const unsafe fn from_raw(pointer: *mut T) -> Self { + Self { + pointer, + storage: PhantomData, + } + } + /// Return the raw address without extending its allocation lifetime. + #[inline(always)] + pub const fn as_raw(self) -> *mut T { + self.pointer + } + } + }; +} +address!( + SharedAddress, + "An address in the current CTA's shared memory." +); +address!( + ClusterSharedAddress, + "A shared address mapped to a CTA within the cluster." +); + +/// Why a shared allocation cannot satisfy a physical storage format. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MemoryError { + /// The supplied shared allocation base is null. + NullAddress, + /// The base does not meet the format’s byte alignment. + Misaligned, + /// The allocation is smaller than the complete physical tile. + InsufficientCapacity, + /// The storage element count cannot be converted to a byte capacity. + SizeOverflow, +} + +/// Non-owning physical shared storage. No exclusive slice is manufactured. +pub struct SharedRegion<'s, F: PhysicalFormat> { + /// Local shared address whose allocation satisfies this format. + address: SharedAddress<'s, F::Element>, + /// Carries the checked format without storing additional device data. + format: PhantomData, +} +impl<'s, F: PhysicalFormat> SharedRegion<'s, F> { + /// Validate a region's numeric layout and adopt its storage. + /// + /// # Safety + /// If validation succeeds, `base` must point to live shared memory for `'s` + /// and cover `capacity_bytes`. Concurrent accesses and completion of + /// asynchronous operations must be coordinated by the kernel author. + #[inline(always)] + pub unsafe fn from_raw_parts( + base: *mut F::Element, + capacity_bytes: usize, + ) -> Result { + if base.is_null() { + return Err(MemoryError::NullAddress); + } + if !(base as usize).is_multiple_of(F::BASE_ALIGNMENT) { + return Err(MemoryError::Misaligned); + } + if capacity_bytes < F::FOOTPRINT_BYTES { + return Err(MemoryError::InsufficientCapacity); + } + Ok(Self { + address: unsafe { SharedAddress::from_raw(base) }, + format: PhantomData, + }) + } + + /// Validate shared storage and abort the kernel if validation fails. + /// + /// Checks the same null address, alignment and capacity conditions as + /// [`Self::from_raw_parts`]. Use this when the kernel cannot recover from + /// invalid storage. The pointer remains a pointer throughout construction, + /// without an intermediate pointer-bearing `Result` storage carrier. + /// + /// # Safety + /// If validation succeeds, `base` must point to live shared memory for `'s` + /// and cover `capacity_bytes`. Concurrent accesses and completion of + /// asynchronous operations must be coordinated by the kernel author. + #[inline(always)] + pub unsafe fn from_raw_parts_or_trap(base: *mut F::Element, capacity_bytes: usize) -> Self { + if base.is_null() + || !(base as usize).is_multiple_of(F::BASE_ALIGNMENT) + || capacity_bytes < F::FOOTPRINT_BYTES + { + cuda_device::debug::trap(); + } + Self { + address: unsafe { SharedAddress::from_raw(base) }, + format: PhantomData, + } + } + + /// Bind a format to a canonical static shared-storage declaration. + /// + /// # Safety + /// `storage` must name the actual static shared declaration, not a copied + /// marker. The raw-storage and asynchronous-access contracts still apply. + #[inline(always)] + pub unsafe fn from_array( + storage: &'s SharedArray, + ) -> Result { + let capacity = N + .checked_mul(size_of::()) + .ok_or(MemoryError::SizeOverflow)?; + let pointer = + unsafe { SharedArray::as_raw_mut_ptr(core::ptr::from_ref(storage).cast_mut()) }; + unsafe { Self::from_raw_parts(pointer, capacity) } + } + /// Borrow the checked local shared address without issuing a memory access. + #[inline(always)] + pub fn address(&self) -> SharedAddress<'s, F::Element> { + self.address + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/mma.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/mma.rs new file mode 100644 index 0000000000..b178cff455 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/mma.rs @@ -0,0 +1,318 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Warp-level FP16/BF16 and SM120 MXFP4 instructions with role-specific registers. +//! Addresses, lane mapping, synchronization and issue order remain explicit. +use super::memory::SharedAddress; + +/// Map one lane's accumulator register to its row/column in an m16n8 tile. +/// +/// Registers 0/1 are adjacent columns in the lower row; registers 2/3 hold +/// the row eight positions later. As in the canonical mapping, lane and +/// register indices are reduced modulo 32 and 4 respectively. +#[must_use] +#[inline(always)] +pub const fn acc_coords(lane: usize, register: usize) -> (usize, usize) { + cuda_device::mma_frag::acc_coords(lane, register) +} + +/// Expand a body once per listed index, with a const-generic-compatible index. +/// +/// Each expansion has its own block scope and executes in list order. Nest +/// invocations for multiple axes. This is source expansion, not a runtime loop +/// or a compiler unroll hint; it does not imply a particular register allocation. +/// +/// ``` +/// use ctm_rs::mma::{self, AccumulatorTile16x8}; +/// let tile = AccumulatorTile16x8::<2, 8>::zero(); +/// mma::static_for!(ROW in [0, 1] { +/// mma::static_for!(COL in [0, 1, 2, 3, 4, 5, 6, 7] { +/// assert_eq!(tile.fragment::().registers(), [0.0; 4]); +/// }); +/// }); +/// ``` +#[macro_export] +macro_rules! static_for { + ($index:ident in [$($value:literal),* $(,)?] $body:block) => { + $({ + /// Compile-time coordinate for this expansion in the requested index order. + const $index: usize = $value; + $body + })* + }; +} +pub use crate::static_for; + +/// Define an operand-role register carrier and its by-value accessors. +macro_rules! registers { + ($name:ident, $element:ty, $count:literal, $description:literal) => { + #[doc = $description] + #[repr(transparent)] + #[derive(Clone, Copy)] + pub struct $name( + #[doc = "Per-lane register values in the instruction's operand order."] + [$element; $count], + ); + impl $name { + /// Adopt registers without changing their bits or lane ownership. + #[inline(always)] + pub const fn from_registers(registers: [$element; $count]) -> Self { + Self(registers) + } + /// Return the registers by value in their original instruction order. + #[inline(always)] + pub const fn registers(self) -> [$element; $count] { + self.0 + } + } + }; +} +registers!( + F16A, + u32, + 4, + "One lane's packed FP16 A fragment for m16n8k16, in row-layout PTX register order." +); +registers!( + F16B, + u32, + 2, + "One lane's packed FP16 B fragment for m16n8k16, in column-layout PTX register order." +); +registers!( + Bf16A, + u32, + 4, + "One lane's packed BF16 A fragment for m16n8k16, in row-layout PTX register order." +); +registers!( + Bf16B, + u32, + 2, + "One lane's packed BF16 B fragment for m16n8k16, in column-layout PTX register order." +); +registers!( + Mxf4A, + u32, + 4, + "One lane's packed E2M1 A fragment for m16n8k64, in PTX register order." +); +registers!( + Mxf4B, + u32, + 2, + "One lane's packed E2M1 B fragment for m16n8k64, in PTX register order." +); +registers!( + Accumulator16x8, + f32, + 4, + "One lane's four FP32 accumulators for a warp-wide 16x8 output." +); +impl Accumulator16x8 { + /// Initialize this lane’s four accumulator registers to FP32 zero. + #[inline(always)] + pub const fn zero() -> Self { + Self([0.0; 4]) + } +} + +/// Issue one m16n8k16 row/column FP16 MMA, accumulating in FP32. +/// +/// The four A registers and two B registers each pack two FP16 values. The +/// four accumulator registers retain PTX result order. This forwards exactly +/// `mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32`, without loads, +/// synchronization, election or rearrangement of registers. +/// +/// # Safety +/// All 32 lanes must execute uniformly on SM80 or newer. A, B and C must +/// follow the PTX lane/register mappings for the same warp-wide product. +#[must_use = "use the returned accumulator; mma does not modify c in place"] +#[inline(always)] +pub unsafe fn mma_f16(c: Accumulator16x8, a: F16A, b: F16B) -> Accumulator16x8 { + Accumulator16x8(unsafe { cuda_device::wmma::mma_m16n8k16_f32_f16(c.0, a.0, b.0) }) +} + +/// Issue one m16n8k16 row/column BF16 MMA, accumulating in FP32. +/// +/// Each A/B register packs two BF16 values, with the first in its low sixteen +/// bits. Register and lane order follow +/// [`cuda_device::wmma::mma_m16n8k16_f32_bf16`]. The four FP32 result registers +/// retain PTX accumulator order. No conversion, load, synchronization or +/// rearrangement is inserted. +/// +/// # Safety +/// All 32 lanes must execute uniformly on SM80 or newer. A, B and C must +/// follow the PTX lane/register mappings for the same warp-wide product. +#[must_use = "use the returned accumulator; mma does not modify c in place"] +#[inline(always)] +pub unsafe fn mma_bf16(c: Accumulator16x8, a: Bf16A, b: Bf16B) -> Accumulator16x8 { + Accumulator16x8(unsafe { cuda_device::wmma::mma_m16n8k16_f32_bf16(c.0, a.0, b.0) }) +} + +/// One lane's rectangular collection of warp-wide 16x8 accumulator fragments. +/// +/// `ROWS` and `COLS` count fragments, not scalar elements. Physical placement is +/// the caller's responsibility: adjacent indices need not be adjacent in memory. +/// Access is const-indexed so callers can express a static instruction schedule +/// without runtime indexing into the register collection. +#[repr(transparent)] +pub struct AccumulatorTile16x8 { + /// Per-lane accumulator fragments indexed by logical fragment row and column. + fragments: [[Accumulator16x8; COLS]; ROWS], +} + +impl AccumulatorTile16x8 { + /// Initialize every fragment's four FP32 accumulators to zero. + #[inline(always)] + pub const fn zero() -> Self { + Self { + fragments: [[Accumulator16x8::zero(); COLS]; ROWS], + } + } + + /// Read one fragment; an out-of-range coordinate fails at compile time. + #[inline(always)] + pub const fn fragment(&self) -> Accumulator16x8 { + const { + assert!( + ROW < ROWS && COL < COLS, + "accumulator fragment out of bounds" + ) + } + self.fragments[ROW][COL] + } + + /// Mutably borrow one fragment at a compile-time coordinate. + #[inline(always)] + pub fn fragment_mut(&mut self) -> &mut Accumulator16x8 { + const { + assert!( + ROW < ROWS && COL < COLS, + "accumulator fragment out of bounds" + ) + } + &mut self.fragments[ROW][COL] + } +} + +impl AccumulatorTile16x8<2, COLS> { + /// Accumulate one B pair into both fragment rows using exactly four MMAs. + /// + /// Issue order is `(row 0, left)`, `(row 0, right)`, `(row 1, left)`, then + /// `(row 1, right)`, where `left = 2 * PAIR`. No loads or synchronization + /// are inserted. The caller can load the next B pair after this operation. + /// + /// # Safety + /// The collective, architecture and register-mapping requirements of + /// [`mma_mxf4`] apply to every fragment and scale pair. + #[inline(always)] + pub unsafe fn mma_mxf4_pair( + &mut self, + a: [Mxf4A; 2], + b: [Mxf4B; 2], + scale_a: [ScalePairA; 2], + scale_b: [ScalePairB; 2], + ) { + const { assert!(PAIR < COLS / 2, "accumulator column pair out of bounds") } + static_for!(ROW in [0, 1] { + static_for!(COL in [0, 1] { + let c = &mut self.fragments[ROW][const { 2 * PAIR + COL }]; + *c = unsafe { mma_mxf4(*c, a[ROW], b[COL], scale_a[ROW], scale_b[COL]) }; + }); + }); + } +} + +/// Define role-specific packed scale pairs with K32 scales in low/high bytes. +macro_rules! scales { + ($name:ident, $role:literal) => { + #[doc = $role] + #[repr(transparent)] + #[derive(Clone, Copy)] + pub struct $name(#[doc = "Two UE8M0 scales, low byte first along K."] u16); + impl $name { + /// Low byte scales K[0..32], high byte scales K[32..64]. + #[inline(always)] + pub const fn from_bits(bits: u16) -> Self { + Self(bits) + } + /// Return both packed scale bytes unchanged. + #[inline(always)] + pub const fn bits(self) -> u16 { + self.0 + } + } + }; +} +scales!(ScalePairA, "Packed UE8M0 scales for operand A."); +scales!(ScalePairB, "Packed UE8M0 scales for operand B."); + +/// Load four A registers using one row-layout ldmatrix.x4 instruction. +/// +/// # Safety +/// All 32 lanes must participate with live, aligned shared row addresses in the +/// PTX m8n8.x4 mapping. Order writes/reads before and after this weak memory op. +#[inline(always)] +pub unsafe fn load_mxf4_a(address: SharedAddress<'_, u32>) -> Mxf4A { + Mxf4A(unsafe { cuda_device::wmma::ldmatrix_x4(address.as_raw()) }) +} +/// Load one B pair with ldmatrix.x4; separate the even/odd row submatrices. +/// +/// # Safety +/// Same collective/address requirements as [`load_mxf4_a`]. Addresses must +/// select two neighboring N8 fragments of an E2M1 K64 tile. +#[inline(always)] +pub unsafe fn load_mxf4_b_pair(address: SharedAddress<'_, u32>) -> (Mxf4B, Mxf4B) { + let r = unsafe { cuda_device::wmma::ldmatrix_x4(address.as_raw()) }; + (Mxf4B([r[0], r[2]]), Mxf4B([r[1], r[3]])) +} +/// Issue one SM120 m16n8k64 E2M1/UE8M0 block-scaled MMA with selectors zero. +/// No loads, synchronization, election or completion waits are inserted. +/// +/// # Safety +/// All 32 lanes must execute uniformly on SM120a. Fragments and per-lane scale +/// providers must follow the PTX register mapping for the same K64 tile. +/// +/// ```compile_fail +/// use ctm_rs::mma::*; +/// unsafe { +/// mma_mxf4(Accumulator16x8::zero(), Mxf4A::from_registers([0;4]), +/// Mxf4B::from_registers([0;2]), ScalePairB::from_bits(0), ScalePairA::from_bits(0)); +/// } +/// ``` +#[inline(always)] +pub unsafe fn mma_mxf4( + c: Accumulator16x8, + a: Mxf4A, + b: Mxf4B, + scale_a: ScalePairA, + scale_b: ScalePairB, +) -> Accumulator16x8 { + Accumulator16x8(unsafe { + cuda_device::wmma::mma_m16n8k64_mxf4_f32( + c.0, + a.0, + b.0, + scale_a.0 as u32, + 0, + 0, + scale_b.0 as u32, + 0, + 0, + ) + }) +} +/// Convert four accumulators to FP16 and issue one row-layout stmatrix.x2. +/// +/// # Safety +/// All 32 lanes must participate. Shared row addresses must satisfy PTX x2 +/// mapping and alignment, and refer to live output storage. Explicitly publish +/// the writes to the async proxy before any dependent TMA transfer. +#[inline(always)] +pub unsafe fn store_f16(address: SharedAddress<'_, u8>, c: Accumulator16x8) { + let top = crate::convert::cvt_f16x2_f32(c.0[0], c.0[1]); + let bottom = crate::convert::cvt_f16x2_f32(c.0[2], c.0[3]); + unsafe { cuda_device::tcgen05::stmatrix_m8n8_x2(address.as_raw(), top, bottom) } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/descriptors.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/descriptors.rs new file mode 100644 index 0000000000..6a26a4305d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/descriptors.rs @@ -0,0 +1,263 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Typed dense FP16 descriptors using the canonical instruction encoders. +use super::tmem::{TmemF32, TmemFormat}; +use crate::{ + groups::{CtaGroup, CtaGroup2}, + memory::*, +}; +use core::{cell::UnsafeCell, marker::PhantomData}; +use cuda_device::tcgen05::{ + Tcgen05InstructionDescriptor, Tcgen05SmemDescriptor, Tcgen05SwizzleMode, +}; + +/// Restricts MMA specifications, storage formats, operand roles and K slices. +mod sealed { + /// Implemented only for supported dense FP16 instruction configurations. + pub trait Spec {} + /// Marks physical FP16 layouts accepted as shared MMA inputs. + pub trait Input {} + /// Marks the distinct A and B instruction operand roles. + pub trait Role {} + /// Admits one compile-time K16 slice index within a K64 stage. + pub trait Step {} +} +/// The shared operand's K coordinate is contiguous before swizzling. +/// +/// This descriptor marker is distinct from cute-rs's `KMajor` +/// (`cute_rs::block_scaled::KMajor`), a block-scaled operand layout that also +/// carries row and K extents. +pub enum KMajor {} +/// Left matrix input; cannot be substituted for [`OperandB`]. +/// +/// cute-rs reexports this marker as `cute_rs::sm100::OperandA` (also at its +/// crate root), so shared tiles there and instruction operands here share one +/// role vocabulary. +pub enum OperandA {} +/// Right matrix input; cannot be substituted for [`OperandA`]. +/// +/// cute-rs reexports this marker as `cute_rs::sm100::OperandB` (also at its +/// crate root), so shared tiles there and instruction operands here share one +/// role vocabulary. +pub enum OperandB {} +impl sealed::Role for OperandA {} +impl sealed::Role for OperandB {} +/// Sealed distinction between left and right shared MMA operands. +pub trait OperandRole: sealed::Role {} +impl OperandRole for OperandA {} +impl OperandRole for OperandB {} +/// Supported SM100 group-two FP16 MMA configuration. +/// +/// Implemented only for `G = CtaGroup2`, `M = 256`, `N = 192` or `160`, and +/// both majors [`KMajor`]. M/N count the collective output, while A and B +/// occupy `(128, 64)` and `(N / 2, 64)` FP16 elements per CTA. Each issue +/// selects one K16 slice of those K64 shared stages with [`SmemOperand::k16`]. +/// +/// ``` +/// use ctm_rs::{groups::CtaGroup2, tcgen05::descriptors::*}; +/// type Tile = DenseSharedF16; +/// let descriptor = InstructionDescriptor::::new(); +/// ``` +pub struct DenseSharedF16( + #[doc = "Records CTA group and operand-major choices without runtime data."] + PhantomData<(G, AMajor, BMajor)>, +); +/// A supported shared FP16 tile usable by the descriptor encoder. +pub trait MmaInputFormat: PhysicalFormat + sealed::Input {} +/// Register shared FP16 formats that the descriptor encoder can represent. +macro_rules! input { ($($f:ty),*) => { $(impl sealed::Input for $f {} impl MmaInputFormat for $f {})* }; } +input!(F16Tile128, F16TileB96, F16TileB80); +/// One supported combination of instruction bits and typed operand formats. +pub trait MmaSpec: sealed::Spec { + /// CTA group used by every operation touching the accumulator. + type Group: CtaGroup; + /// Physical shared layout of one CTA’s A operand. + type AFormat: MmaInputFormat; + /// Physical shared layout of one CTA’s B operand. + type BFormat: MmaInputFormat; + /// FP32 tensor-memory shape of the collective result. + type DstFormat: TmemFormat; + /// Canonical packed hardware instruction bits for this configuration. + const DESCRIPTOR: Tcgen05InstructionDescriptor; +} +/// Define a supported collective N shape and its local B storage format. +macro_rules! spec { + ($n:literal,$b:ty) => { + impl sealed::Spec for DenseSharedF16 {} + impl MmaSpec for DenseSharedF16 { + type Group = CtaGroup2; + type AFormat = F16Tile128; + type BFormat = $b; + type DstFormat = TmemF32<256, $n>; + const DESCRIPTOR: Tcgen05InstructionDescriptor = + Tcgen05InstructionDescriptor::builder().m(256).n($n).build(); + } + }; +} +spec!(192, F16TileB96); +spec!(160, F16TileB80); + +/// Instruction bits for one supported [`DenseSharedF16`] configuration. +pub struct InstructionDescriptor { + /// Canonical packed instruction descriptor for S. + raw: Tcgen05InstructionDescriptor, + /// Preserves the group, shape and operand-format specification in the type. + spec: PhantomData, +} +impl InstructionDescriptor { + /// Use the static instruction descriptor selected by S. + #[inline(always)] + pub const fn new() -> Self { + Self { + raw: S::DESCRIPTOR, + spec: PhantomData, + } + } + /// Return the instruction’s packed 32-bit descriptor. + #[inline(always)] + pub const fn raw(&self) -> u32 { + self.raw.raw() + } +} +impl Default for InstructionDescriptor { + fn default() -> Self { + Self::new() + } +} + +/// Why a shared address cannot be represented by this operand descriptor. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DescriptorError { + /// The address violates descriptor or physical-format alignment. + Misaligned, + /// The shared byte offset exceeds the descriptor address field. + AddressOutOfRange, +} +#[inline(always)] +/// Pack a 16-byte-aligned shared offset into the fixed B128 K-major descriptor. +fn encode_smem_offset(offset: ByteOffset) -> Result { + if !offset.bytes().is_multiple_of(16) { + return Err(DescriptorError::Misaligned); + } + if offset.bytes() > 0x3fff0 { + return Err(DescriptorError::AddressOutOfRange); + } + Ok(Tcgen05SmemDescriptor::builder() + .address(offset.bytes() as u64) + .leading_dim_bytes(16) + .stride_bytes(1024) + .swizzle(Tcgen05SwizzleMode::Swizzle128B) + .build()) +} +/// Sealed proof that a K16 slice index is in 0..4. +pub trait ValidKStep: sealed::Step {} +/// Admit only the four K16 slice indices in a K64 shared stage. +macro_rules! step { ($($n:literal),*) => { $(impl sealed::Step<$n> for () {} impl ValidKStep<$n> for () {})* }; } +step!(0, 1, 2, 3); + +/// A role-specific descriptor borrowing physical shared storage. +pub struct SmemOperand<'s, R: OperandRole, F: MmaInputFormat> { + /// Original K64-stage descriptor, retained across repeated slice selection. + stage: Tcgen05SmemDescriptor, + /// Descriptor for the currently selected K16 slice. + selected: Tcgen05SmemDescriptor, + /// Keeps the shared allocation borrowed through asynchronous operand use. + storage: PhantomData<&'s UnsafeCell>, + /// Records A/B role and physical format without runtime data. + marker: PhantomData<(R, F)>, +} +impl<'s, R: OperandRole, F: MmaInputFormat> SmemOperand<'s, R, F> { + /// Borrow a K64 stage after checking its shared address and base alignment. + /// Does not publish writes or wait for copies; the MMA issuer coordinates those. + #[inline(always)] + pub fn from_region(region: &SharedRegion<'s, F>) -> Result { + let offset = unsafe { + cuda_device::shared::cvta_generic_to_shared_offset(region.address().as_raw().cast()) + }; + if offset > 0x3ffff { + return Err(DescriptorError::AddressOutOfRange); + } + if offset % F::BASE_ALIGNMENT as u64 != 0 { + return Err(DescriptorError::Misaligned); + } + let stage = encode_smem_offset(ByteOffset::new(offset as usize))?; + Ok(Self { + stage, + selected: stage, + storage: PhantomData, + marker: PhantomData, + }) + } + /// Select one K16 slice relative to the original K64 stage, even on a slice. + #[inline(always)] + pub fn k16(&self) -> Self + where + (): ValidKStep, + { + // A 1024-aligned, representable stage leaves room for all four slices + // within the address field. Only that field changes, in 16-byte units. + Self { + stage: self.stage, + selected: Tcgen05SmemDescriptor::from_raw(self.stage.raw() + STEP as u64 * 2), + storage: PhantomData, + marker: PhantomData, + } + } + #[inline(always)] + /// Return the selected slice descriptor in the hardware instruction encoding. + pub(crate) fn raw(&self) -> u64 { + self.selected.raw() + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + use crate::groups::CtaGroup2; + /// Supported instruction and shared-address descriptors must match known hardware bit encodings. + #[test] + fn instruction_and_smem_descriptors_match_canonical_bits() { + /// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. + type L = DenseSharedF16; + /// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. + type R = DenseSharedF16; + assert_eq!(InstructionDescriptor::::new().raw(), 0x1030_0010); + assert_eq!(InstructionDescriptor::::new().raw(), 0x1028_0010); + assert_eq!( + encode_smem_offset(ByteOffset::new(0)).unwrap().raw(), + 0x4000_4040_0001_0000 + ); + assert_eq!( + encode_smem_offset(ByteOffset::new(0x1000)).unwrap().raw(), + 0x4000_4040_0001_0100 + ); + for step in 0..4 { + assert_eq!( + encode_smem_offset(ByteOffset::new(0x1000 + step * 32)) + .unwrap() + .raw(), + 0x4000_4040_0001_0100 + step as u64 * 2 + ); + } + assert!(encode_smem_offset(ByteOffset::new(15)).is_err()); + assert!(encode_smem_offset(ByteOffset::new(0x40000)).is_err()); + assert!(encode_smem_offset(ByteOffset::new(usize::MAX)).is_err()); + assert!(encode_smem_offset(ByteOffset::new(0x3fff0)).is_ok()); + } + + /// Selecting a K16 slice from another slice must still use the original K64 stage base. + #[test] + fn k16_selection_is_relative_to_original_stage() { + let stage = encode_smem_offset(ByteOffset::new(0x1000)).unwrap(); + let a: SmemOperand<'_, OperandA, F16Tile128> = SmemOperand { + stage, + selected: stage, + storage: PhantomData, + marker: PhantomData, + }; + assert_eq!(a.k16::<3>().k16::<1>().raw(), 0x4000_4040_0001_0102); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mma.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mma.rs new file mode 100644 index 0000000000..6ceeebff53 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mma.rs @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Explicit MMA issue, collector selection and TMEM register readback. +use super::{ + descriptors::*, + tmem::{TmemOperand, TmemReadF32}, +}; +use crate::groups::CtaGroup2; +use cuda_device::cusimd::CuSimd; + +/// Explicit A-collector reuse selectors; kernel code must order compatible issues. +pub mod collector_a { + /// Read A without retaining it for a subsequent issue. + pub enum Discard {} + /// Use retained A for its final issue and release the collector. + pub enum LastUse {} + /// Read A and retain it for subsequent collector uses. + pub enum Fill {} + /// Reuse retained A while keeping the collector available for later issues. + pub enum Use {} +} +/// Restricts MMA issue to supported specification and collector-selector pairs. +mod sealed { + /// Dispatch a typed MMA call to a literal hardware collector selector. + pub trait MmaDispatch { + /// Issue shared/shared FP16 MMA; caller supplies election, live operands and collector ordering. + unsafe fn issue(dst: u32, a: u64, b: u64, descriptor: u32, enable_d: bool); + } +} +/// Implement one collector selector using the existing group-two MMA intrinsic. +macro_rules! dispatch { + ($usage:ty,$selector:literal) => { + impl sealed::MmaDispatch for (S, $usage) { + #[inline(always)] + unsafe fn issue(dst: u32, a: u64, b: u64, descriptor: u32, enable_d: bool) { + unsafe { + cuda_device::tcgen05::tcgen05_mma_shared::<0, 2, $selector>( + dst, a, b, descriptor, enable_d, + ) + } + } + } + }; +} +dispatch!(collector_a::Discard, 0); +dispatch!(collector_a::LastUse, 1); +dispatch!(collector_a::Fill, 2); +dispatch!(collector_a::Use, 3); + +/// Issue one MMA. The collector selector does not infer a sequence or pipeline. +/// +/// cute-rs exposes the `enable_d` flag (PTX `enable-input-d`) as `accumulate`, +/// CuTeDSL's `Field.ACCUMULATE`. +/// +/// # Safety +/// The elected issuing thread must satisfy group-2 participation and collector +/// sequencing rules. Shared operands must be published and remain live until +/// completion; the accumulator must be initialized when `enable_d` is true. +/// Complete dependent asynchronous work explicitly before reading/reusing it. +#[inline(always)] +pub unsafe fn mma( + dst: &impl TmemOperand, + a: &SmemOperand<'_, OperandA, S::AFormat>, + b: &SmemOperand<'_, OperandB, S::BFormat>, + descriptor: &InstructionDescriptor, + enable_d: bool, +) where + S: MmaSpec, + (S, U): sealed::MmaDispatch, +{ + unsafe { + <(S, U) as sealed::MmaDispatch>::issue( + dst.raw_address(), + a.raw(), + b.raw(), + descriptor.raw(), + enable_d, + ) + } +} + +/// Registers returned by a pending TMEM load; decode only after an explicit wait. +#[repr(transparent)] +pub struct TmemRaw32( + #[doc = "Pending raw registers, consumed only after the caller’s explicit load wait."] + CuSimd, +); + +/// Issue one warp-collective TMEM load; no wait is inserted. +/// +/// # Safety +/// All lanes must issue uniformly with valid addresses and completed producers. +/// Keep the allocation live and call [`super::tmem::load_wait`] before +/// consuming the registers. +#[inline(always)] +pub unsafe fn load_32x32b_x32( + src: &impl TmemOperand>, +) -> TmemRaw32 { + TmemRaw32(unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(src.raw_address()) }) +} +impl TmemRaw32 { + /// Interpret completed load registers in their original order. + /// + /// # Safety + /// The issuing thread must execute [`super::tmem::load_wait`] after this load and before + /// decoding. The wrapper does not model the hardware scoreboard. + #[inline(always)] + pub unsafe fn decode_f32(self) -> CuSimd { + // Destructure by value. Borrowing the carrier for get::() forces + // an addressable 128-byte local before downstream inlining/SROA. + /// Bitcast each completed register by value, avoiding an addressable local carrier. + macro_rules! decode { ($($r:ident),*) => {{ + let [$($r),*] = self.0.to_array(); + CuSimd::new([$(f32::from_bits($r)),*]) + }}; } + decode!( + r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12, r13, r14, r15, r16, r17, r18, + r19, r20, r21, r22, r23, r24, r25, r26, r27, r28, r29, r30, r31 + ) + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mod.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mod.rs new file mode 100644 index 0000000000..ec4d56d62a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/mod.rs @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! SM100 tensor-memory allocation, FP16 MMA and register readback. +//! +//! Instructions require PTX 8.6 and a tcgen05-capable target profile; the current +//! CTM consumer validates `sm_100a`. TCGEN05 is unavailable on SM120a. +//! +//! Start with [`descriptors::DenseSharedF16`] for supported shapes and +//! [`tmem::TmemGuard`] for allocation ownership. Issue, commit, wait and release +//! remain explicit; [`tmem::load_wait`] and [`tmem::store_wait`] are the two +//! completion waits. The crate README gives an issue-site example and the +//! supported compiler profiles. + +pub mod descriptors; +pub mod mma; +pub mod tmem; + +/// Commit the issuing thread's group-one MMA work to a ready local barrier. +/// +/// # Safety +/// The elected issuer must own the pending group-one MMA operations. This live +/// barrier phase must expect their completion arrival, and operands must remain +/// live until completion. This function does not elect, wait or change phases. +/// All tcgen05 operations in the kernel must use group-one mode consistently. +/// +/// ```compile_fail +/// use ctm_rs::mbarrier::{GeneralBarrier, ManagedBarrier, Uninit}; +/// # fn rejected(barrier: &ManagedBarrier) { +/// unsafe { ctm_rs::tcgen05::commit(barrier) }; +/// # } +/// ``` +#[inline(always)] +pub unsafe fn commit( + barrier: &crate::mbarrier::ManagedBarrier, +) { + unsafe { cuda_device::tcgen05::tcgen05_commit(barrier.as_ptr().cast_mut().cast()) } +} + +/// Relinquish the group-one allocation permit after allocation completes. +/// +/// # Safety +/// The allocating warp must participate uniformly, and the kernel must use +/// group-one TMEM operations consistently. No deallocation is performed. +#[inline(always)] +pub unsafe fn relinquish_alloc_permit() { + cuda_device::tcgen05::tcgen05_relinquish_alloc_permit() +} + +/// Order later tcgen05 accesses after thread synchronization. +/// All tcgen05 operations in the kernel must use the same CTA-group mode. +#[inline(always)] +pub fn fence_after_thread_sync() { + cuda_device::tcgen05::tcgen05_fence_after_thread_sync() +} + +/// Order prior tcgen05 accesses before thread synchronization. +/// All tcgen05 operations in the kernel must use the same CTA-group mode. +#[inline(always)] +pub fn fence_before_thread_sync() { + cuda_device::tcgen05::tcgen05_fence_before_thread_sync() +} + +/// Release the group-2 tensor-memory allocation permit. +/// One full warp in each peer CTA must execute uniformly. All tcgen05 +/// operations in the kernel must use the same CTA-group mode. +#[inline(always)] +pub fn relinquish_alloc_permit_cg2() { + cuda_device::tcgen05::tcgen05_relinquish_alloc_permit_cg2() +} + +/// Commit group-2 MMA completion to the selected CTAs' local barrier offsets. +/// +/// # Safety +/// The MMA issuer must commit its own operations. Every selected CTA must be +/// active with a matching initialized barrier expecting the completion arrival. +/// Keep all associated resources live until the barrier signals completion. +#[inline(always)] +pub unsafe fn commit_multicast_cg2( + barrier: super::memory::SharedAddress<'_, cuda_device::barrier::Barrier>, + mask: super::cluster::CtaMask, +) { + unsafe { + cuda_device::tcgen05::tcgen05_commit_multicast_cg2(barrier.as_raw().cast(), mask.bits()) + } +} + +/// Commit group-1 MMA completion to a local barrier (SM100), without waiting. +/// +/// # Safety +/// The MMA issuer must commit its own operations exactly once to a live +/// initialized CTA-shared barrier expecting this completion arrival. Keep all +/// associated resources live until completion. All tcgen05 operations in the +/// kernel must use group-1 mode. +#[inline(always)] +pub unsafe fn commit_address(barrier: super::memory::SharedAddress<'_, crate::mbarrier::Barrier>) { + unsafe { cuda_device::tcgen05::tcgen05_commit(barrier.as_raw().cast()) } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/tmem.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/tmem.rs new file mode 100644 index 0000000000..c8b40b26dd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/tcgen05/tmem.rs @@ -0,0 +1,507 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Borrowed TMEM views and explicit collective allocation instructions. +//! +//! Raw register transfers require PTX 8.6 and a tcgen05-capable architecture +//! profile (`sm_100a`, `sm_101a`, `sm_103a` or `sm_110a` in the intrinsic catalog). +//! The current CTM consumer validates `sm_100a`; these are unavailable on SM120a. +use crate::groups::{CtaGroup, CtaGroup1}; +use core::marker::PhantomData; +pub use cuda_device::tcgen05::{TmemAddress as TmemAddr, TmemDeallocated, TmemReady, TmemUninit}; + +/// A typed handle to TMEM with explicit collective allocation and release. +/// +/// The handle contains only its shared address-slot pointer and zero-sized +/// state/group markers. It does not elect threads, insert synchronization, +/// relinquish allocation permits, or deallocate on drop. Callers coordinate +/// every GPU participant and explicitly release the allocation after all uses +/// and asynchronous accesses have completed. +/// +/// A ready handle can lend typed views, which prevent consuming that local +/// handle while a view remains in use. Other threads' handles and asynchronous +/// hardware operations still require caller-managed lifetime coordination. +#[repr(transparent)] +pub struct TmemGuard { + /// Live shared-memory slot containing the allocation’s packed TMEM address. + smem_ptr: *mut u32, + /// Records lifecycle state and CTA group without owning or freeing hardware storage. + _state: PhantomData<(State, Group)>, +} + +// Safety: Constructors and collective operations require callers to coordinate +// access to the shared slot and its allocation across all participating threads. +unsafe impl Send for TmemGuard {} + +/// Restricts allocation sizes, view formats and instruction operands. +mod sealed { + /// Admits only supported power-of-two TMEM allocation widths. + pub trait Columns {} + /// Marks supported accumulator and readback view shapes. + pub trait Format {} + /// Marks views that borrow a live allocation guard. + pub trait Operand {} +} +/// Sealed proof that the allocation width is one of 32, 64, 128, 256 or 512. +pub trait ValidColumns: sealed::Columns {} +/// Admit the instruction’s 32, 64, 128, 256 and 512 column allocation sizes. +macro_rules! columns { ($($n:literal),*) => { $(impl sealed::Columns<$n> for () {} impl ValidColumns<$n> for () {})* }; } +columns!(32, 64, 128, 256, 512); + +/// Offset in TMEM columns from the allocation's base, not a byte offset. +/// Range checks occur when creating a view. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ColumnOffset(#[doc = "Column displacement checked when forming a typed view."] u32); +impl ColumnOffset { + /// Wrap a TMEM column displacement; view construction checks the range. + pub const fn new(value: u32) -> Self { + Self(value) + } +} +/// Group-two FP32 accumulator: collective `M = 256`, `N = 192` or `160`. +/// Each CTA holds 128 rows and N columns. +pub struct TmemF32; +/// Local FP32 readback window; only `(ROWS, COLS) = (32, 32)` is supported. +pub struct TmemReadF32; +/// Sealed row/column shape used to validate a borrowed TMEM view. +pub trait TmemFormat: sealed::Format { + /// Number of consecutive TMEM columns used by the view. + const COLUMNS: u32; + /// Number of rows used in this CTA’s 128-row tensor-memory partition. + const LOCAL_ROWS: u32; +} +/// Marks view shapes accepted as MMA accumulators. +pub trait MmaTmemFormat: TmemFormat {} +/// Define an FP32 accumulator with 128 local rows and a collective M256 shape. +macro_rules! mma_format { + ($n:literal) => { + impl sealed::Format for TmemF32<256, $n> {} + impl TmemFormat for TmemF32<256, $n> { + const COLUMNS: u32 = $n; + const LOCAL_ROWS: u32 = 128; + } + impl MmaTmemFormat for TmemF32<256, $n> {} + }; +} +mma_format!(192); +mma_format!(160); +impl sealed::Format for TmemReadF32<32, 32> {} +impl TmemFormat for TmemReadF32<32, 32> { + const COLUMNS: u32 = 32; + const LOCAL_ROWS: u32 = 32; +} + +/// Why a selected row/column window cannot be packed into this allocation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TmemRangeError { + /// The requested column interval exceeds the allocation or overflows. + Columns, + /// The local row is not a 32-row boundary or the window exceeds 128 rows. + Rows, + /// Adding the selected window would overflow the packed instruction address. + AddressOverflow, +} + +#[inline(always)] +/// Check local row/column bounds and add their packed fields to a TMEM base address. +fn pack_address( + base: u32, + row: u32, + column: ColumnOffset, + width: u32, + rows: u32, +) -> Result { + let end = column.0.checked_add(width).ok_or(TmemRangeError::Columns)?; + if end > N { + return Err(TmemRangeError::Columns); + } + if !row.is_multiple_of(32) || row.checked_add(rows).ok_or(TmemRangeError::Rows)? > 128 { + return Err(TmemRangeError::Rows); + } + if (base & 0xffff) + .checked_add(end) + .ok_or(TmemRangeError::AddressOverflow)? + > 0x10000 + { + return Err(TmemRangeError::AddressOverflow); + } + base.checked_add((row << 16) | column.0) + .ok_or(TmemRangeError::AddressOverflow) +} + +/// A local borrow of a live allocation; other GPU participants still require coordination. +pub struct TmemView<'g, F: TmemFormat, const N: u32, G: CtaGroup> { + /// Packed TMEM address after checking the selected row and column window. + address: u32, + /// Prevents consuming the local allocation handle while this view is borrowed. + guard: PhantomData<&'g TmemGuard>, + /// Carries the view’s row and column extents without runtime storage. + format: PhantomData, +} +/// Sealed instruction operand that borrows live tensor memory. +pub trait TmemOperand: sealed::Operand { + /// CTA group that owns the underlying tensor-memory allocation. + type Group: CtaGroup; + /// Validated row/column shape addressed by the operand. + type Format: TmemFormat; + /// Escape to the packed instruction address. + /// + /// # Safety + /// The address may only be used while this view and its guard remain live; + /// callers must also enforce cross-thread lifetime and completion rules. + unsafe fn raw_address(&self) -> u32; +} +impl sealed::Operand for TmemView<'_, F, N, G> {} +impl TmemOperand for TmemView<'_, F, N, G> { + type Group = G; + type Format = F; + #[inline(always)] + unsafe fn raw_address(&self) -> u32 { + self.address + } +} + +impl TmemGuard +where + (): ValidColumns, +{ + /// Adopt the shared-memory address slot without allocating or synchronizing. + /// + /// # Safety + /// `slot` must be aligned live shared storage through allocation and release. + #[inline(always)] + pub unsafe fn from_raw_slot(slot: *mut u32) -> Self { + Self { + smem_ptr: slot, + _state: PhantomData, + } + } + /// Issue allocation without election, fences or barriers. + /// + /// # Safety + /// One complete warp per required CTA must issue uniformly. All tcgen05 + /// operations must use the same group mode. Synchronize before other users + /// adopt the resulting allocation or read its shared address slot. + #[inline(always)] + pub unsafe fn alloc_collective(self) -> TmemGuard { + unsafe { G::alloc(self.smem_ptr, N) }; + TmemGuard { + smem_ptr: self.smem_ptr, + _state: PhantomData, + } + } +} +impl TmemGuard +where + (): ValidColumns, +{ + /// Read the allocated TMEM address from the synchronized shared slot. + #[inline(always)] + pub fn address(&self) -> TmemAddr { + unsafe { TmemAddr::from_raw(self.raw_address()) } + } + + /// Read the raw TMEM address for low-level instruction interoperability. + #[inline(always)] + pub fn raw_address(&self) -> u32 { + unsafe { *self.smem_ptr } + } + + /// Return the number of columns in this allocation. + #[inline(always)] + pub const fn n_cols(&self) -> u32 { + N + } + + /// Adopt another participant's allocation after caller-provided synchronization. + /// + /// # Safety + /// The shared slot and its group-`G` allocation of `N` columns must remain + /// valid until all views and asynchronous accesses complete. This handle + /// does not establish global ownership or permission to deallocate. + #[inline(always)] + pub unsafe fn adopt_ready(slot: *mut u32) -> Self { + Self { + smem_ptr: slot, + _state: PhantomData, + } + } + /// Issue deallocation without election, fences or barriers. + /// + /// # Safety + /// All users and asynchronous operations must be finished. The required + /// complete warps must participate uniformly, exactly once per allocation. + #[inline(always)] + pub unsafe fn dealloc_collective(self) -> TmemGuard { + unsafe { G::dealloc(self.raw_address(), N) }; + TmemGuard { + smem_ptr: self.smem_ptr, + _state: PhantomData, + } + } + /// Borrow all 128 local rows and `F::COLUMNS` columns for MMA. + /// `column` counts TMEM columns. The entire view must fit in this allocation. + #[inline(always)] + pub fn view( + &self, + column: ColumnOffset, + ) -> Result, TmemRangeError> { + Ok(TmemView { + address: pack_address::(self.raw_address(), 0, column, F::COLUMNS, F::LOCAL_ROWS)?, + guard: PhantomData, + format: PhantomData, + }) + } + /// Borrow a local 32-row × 32-column window without loading it. + /// `local_row` must be 0, 32, 64 or 96; the column range must fit in the + /// allocation. Wait for the producer before issuing a load from this view. + /// + /// ```no_run + /// use ctm_rs::{groups::CtaGroup2, tcgen05::tmem::*}; + /// # fn inspect(guard: &TmemGuard) -> Result<(), TmemRangeError> { + /// let window = guard.read_32x32(32, ColumnOffset::new(192))?; + /// // Local rows 32..64, columns 192..224; no load has been issued. + /// # Ok(()) + /// # } + /// ``` + #[inline(always)] + pub fn read_32x32( + &self, + local_row: u32, + column: ColumnOffset, + ) -> Result, N, G>, TmemRangeError> { + Ok(TmemView { + address: pack_address::(self.raw_address(), local_row, column, 32, 32)?, + guard: PhantomData, + format: PhantomData, + }) + } +} + +/// Issue an SM100 `32x32b.x2` TMEM load into 2 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address, after producer completion and the required after-thread-sync fence. +/// All tcgen05 instructions must agree on CTA-group mode. True register +/// dependencies may consume the result without an immediate wait. Call +/// [`load_wait`] before conflicting memory operations, memory anti-dependencies +/// or a consumer-release handoff, with the required before-thread-sync fence. +/// Keep the allocation live until all accesses complete. See the +/// [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). +#[inline(always)] +pub unsafe fn load_32x32b_x2(address: u32) -> crate::memory::CuSimd { + unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x2_raw(address) } +} + +/// Issue an SM100 `32x32b.x2` TMEM store from 2 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address; coordinate conflicting accesses explicitly. All tcgen05 instructions +/// must agree on CTA-group mode. Call [`store_wait`] before conflicting memory +/// operations or completion handoffs, with the required before-thread-sync fence. +/// Keep the allocation live until every participant's accesses complete. +#[inline(always)] +pub unsafe fn store_32x32b_x2(address: u32, data: crate::memory::CuSimd) { + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x2_raw(address, data) } +} + +/// Issue an SM100 `32x32b.x4` TMEM load into 4 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address, after producer completion and the required after-thread-sync fence. +/// All tcgen05 instructions must agree on CTA-group mode. True register +/// dependencies may consume the result without an immediate wait. Call +/// [`load_wait`] before conflicting memory operations, memory anti-dependencies +/// or a consumer-release handoff, with the required before-thread-sync fence. +/// Keep the allocation live until all accesses complete. See the +/// [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). +#[inline(always)] +pub unsafe fn load_32x32b_x4(address: u32) -> crate::memory::CuSimd { + unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x4_raw(address) } +} + +/// Issue an SM100 `32x32b.x4` TMEM store from 4 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address; coordinate conflicting accesses explicitly. All tcgen05 instructions +/// must agree on CTA-group mode. Call [`store_wait`] before conflicting memory +/// operations or completion handoffs, with the required before-thread-sync fence. +/// Keep the allocation live until every participant's accesses complete. +#[inline(always)] +pub unsafe fn store_32x32b_x4(address: u32, data: crate::memory::CuSimd) { + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x4_raw(address, data) } +} + +/// Issue an SM100 `32x32b.x8` TMEM load into 8 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address, after producer completion and the required after-thread-sync fence. +/// All tcgen05 instructions must agree on CTA-group mode. True register +/// dependencies may consume the result without an immediate wait. Call +/// [`load_wait`] before conflicting memory operations, memory anti-dependencies +/// or a consumer-release handoff, with the required before-thread-sync fence. +/// Keep the allocation live until all accesses complete. See the +/// [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). +#[inline(always)] +pub unsafe fn load_32x32b_x8(address: u32) -> crate::memory::CuSimd { + unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x8_raw(address) } +} + +/// Issue an SM100 `32x32b.x8` TMEM store from 8 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address; coordinate conflicting accesses explicitly. All tcgen05 instructions +/// must agree on CTA-group mode. Call [`store_wait`] before conflicting memory +/// operations or completion handoffs, with the required before-thread-sync fence. +/// Keep the allocation live until every participant's accesses complete. +#[inline(always)] +pub unsafe fn store_32x32b_x8(address: u32, data: crate::memory::CuSimd) { + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x8_raw(address, data) } +} + +/// Issue an SM100 `32x32b.x16` TMEM load into 16 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address, after producer completion and the required after-thread-sync fence. +/// All tcgen05 instructions must agree on CTA-group mode. True register +/// dependencies may consume the result without an immediate wait. Call +/// [`load_wait`] before conflicting memory operations, memory anti-dependencies +/// or a consumer-release handoff, with the required before-thread-sync fence. +/// Keep the allocation live until all accesses complete. See the +/// [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). +#[inline(always)] +pub unsafe fn load_32x32b_x16(address: u32) -> crate::memory::CuSimd { + unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x16_raw(address) } +} + +/// Issue an SM100 `32x32b.x16` TMEM store from 16 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address; coordinate conflicting accesses explicitly. All tcgen05 instructions +/// must agree on CTA-group mode. Call [`store_wait`] before conflicting memory +/// operations or completion handoffs, with the required before-thread-sync fence. +/// Keep the allocation live until every participant's accesses complete. +#[inline(always)] +pub unsafe fn store_32x32b_x16(address: u32, data: crate::memory::CuSimd) { + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x16_raw(address, data) } +} + +/// Issue an SM100 `32x32b.x32` TMEM load into 32 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address, after producer completion and the required after-thread-sync fence. +/// All tcgen05 instructions must agree on CTA-group mode. True register +/// dependencies may consume the result without an immediate wait. Call +/// [`load_wait`] before conflicting memory operations, memory anti-dependencies +/// or a consumer-release handoff, with the required before-thread-sync fence. +/// Keep the allocation live until all accesses complete. See the +/// [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). +#[inline(always)] +pub unsafe fn load_32x32b_x32(address: u32) -> crate::memory::CuSimd { + unsafe { cuda_device::tcgen05::tcgen05_ld_32x32b_x32_raw(address) } +} + +/// Issue an SM100 `32x32b.x32` TMEM store from 32 raw registers per lane. +/// This uses an encoded instruction address; no bounds checks or waits are inserted. +/// +/// # Safety +/// The encoded address must cover a live allocation's selected tile and be +/// row-group aligned. All 32 warp lanes must execute convergently with the same +/// address; coordinate conflicting accesses explicitly. All tcgen05 instructions +/// must agree on CTA-group mode. Call [`store_wait`] before conflicting memory +/// operations or completion handoffs, with the required before-thread-sync fence. +/// Keep the allocation live until every participant's accesses complete. +#[inline(always)] +pub unsafe fn store_32x32b_x32(address: u32, data: crate::memory::CuSimd) { + unsafe { cuda_device::tcgen05::tcgen05_st_32x32b_x32_raw(address, data) } +} + +/// Complete the executing thread's previously issued SM100 TMEM loads. +/// No CTA barrier, producer wait or memory/proxy fence is inserted. +/// +/// # Safety +/// All 32 warp lanes must execute convergently. The loads must target live, +/// valid TMEM allocations with correctly ordered producers. All tcgen05 +/// operations must use one CTA-group mode. +/// +/// ```compile_fail,E0133 +/// ctm_rs::tcgen05::tmem::load_wait(); +/// ``` +#[inline(always)] +pub unsafe fn load_wait() { + cuda_device::tcgen05::tcgen05_load_wait() +} + +/// Complete the executing thread's previously issued SM100 TMEM stores. +/// No CTA barrier, consumer synchronization or memory/proxy fence is inserted. +/// +/// # Safety +/// All 32 warp lanes must execute convergently. The stores must target live, +/// valid TMEM allocations with no conflicting accesses. All tcgen05 operations +/// must use one CTA-group mode. +/// +/// ```compile_fail,E0133 +/// ctm_rs::tcgen05::tmem::store_wait(); +/// ``` +#[inline(always)] +pub unsafe fn store_wait() { + cuda_device::tcgen05::tcgen05_store_wait() +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + /// Valid TMEM windows must preserve allocation bits while rejecting row/column overflow. + #[test] + fn checked_tmem_ranges_preserve_allocation_bits() { + assert_eq!( + pack_address::<512>(0x0200_0000, 0, ColumnOffset::new(0), 192, 128), + Ok(0x0200_0000) + ); + assert_eq!( + pack_address::<512>(0x0200_0000, 0, ColumnOffset::new(192), 160, 128), + Ok(0x0200_00c0) + ); + for row in [0, 32, 64, 96] { + assert_eq!( + pack_address::<512>(0x0200_0000, row, ColumnOffset::new(192), 32, 32), + Ok(0x0200_00c0 | row << 16) + ); + } + assert!(pack_address::<512>(0, 128, ColumnOffset::new(0), 32, 32).is_err()); + assert!(pack_address::<512>(0, 1, ColumnOffset::new(0), 32, 32).is_err()); + assert!(pack_address::<256>(0, 0, ColumnOffset::new(192), 160, 128).is_err()); + assert!(pack_address::<512>(0, 0, ColumnOffset::new(u32::MAX), 32, 32).is_err()); + assert!(pack_address::<512>(0x0000_fff0, 0, ColumnOffset::new(32), 32, 32).is_err()); + assert!(pack_address::<512>(0xffff_0000, 32, ColumnOffset::new(0), 32, 32).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/tma.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/tma.rs new file mode 100644 index 0000000000..048a53469e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/tma.rs @@ -0,0 +1,223 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed tensor maps with canonical opaque descriptor storage. +use super::memory::PhysicalFormat; +use super::{ + cluster::CtaMask, + groups::CtaGroup2, + memory::{ClusterSharedAddress, SharedAddress, SharedRegion}, +}; +use cuda_device::TmaDescriptor; +use cuda_device::barrier::Barrier; + +/// Signed element coordinates; descriptor-defined out-of-bounds behavior is retained. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Coordinates2D { + /// Coordinate in the descriptor’s contiguous dimension, measured in storage elements. + inner: i32, + /// Coordinate in the descriptor’s outer dimension, normally the global row. + outer: i32, +} +impl Coordinates2D { + /// Create descriptor coordinates in inner-then-outer element order. + pub const fn new(inner: i32, outer: i32) -> Self { + Self { inner, outer } + } + /// Return the signed coordinate in the contiguous element dimension. + pub const fn inner(self) -> i32 { + self.inner + } + /// Return the signed coordinate in the outer, normally row, dimension. + pub const fn outer(self) -> i32 { + self.outer + } +} + +/// Issue one ordinary TMA load into this CTA's shared storage. +/// +/// This is the CuTeDSL primitives name for the unicast `cp.async.bulk.tensor` +/// load; the rank is fixed by [`TensorMap`]. No issuer election, fence, +/// barrier arrival or completion wait is inserted. +/// +/// # Safety +/// The caller must select the issuer explicitly and publish an initialized +/// completion barrier with the expected byte count. Keep the map, global +/// allocation, destination and barrier live until completion. Coordinate all +/// accesses to the destination and wait explicitly before reading its contents. +#[inline(always)] +pub unsafe fn cp_async_bulk_tensor_shared_cta_global>( + dst: &SharedRegion<'_, F>, + map: &TensorMap, + coords: Coordinates2D, + completion: SharedAddress<'_, Barrier>, +) { + unsafe { + cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta( + dst.address().as_raw().cast(), + map.as_descriptor(), + coords.inner, + coords.outer, + completion.as_raw(), + ) + } +} + +/// Issue one group-2 multicast TMA load, using the default cache hint (zero). +/// +/// This is the CuTeDSL primitives name for the cluster-multicast +/// `cp.async.bulk.tensor` load; the `CtaMask` argument fixes the +/// `.multicast::cluster` and `cta_group::2` qualifiers. +/// +/// # Safety +/// The elected issuer must satisfy group-2 TMA rules. Every mask bit must name +/// a participating CTA with matching destination storage. The remote completion +/// barrier must be initialized with the expected byte count. Keep the map, +/// allocation, destinations and barriers live until completion; wait explicitly. +#[inline(always)] +pub unsafe fn cp_async_bulk_tensor_shared_cluster_global>( + dst: &SharedRegion<'_, F>, + map: &TensorMap, + coords: Coordinates2D, + completion: ClusterSharedAddress<'_, Barrier>, + mask: CtaMask, +) { + unsafe { + cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast_cg2( + dst.address().as_raw().cast(), + map.as_descriptor(), + coords.inner, + coords.outer, + completion.as_raw(), + mask.bits(), + ) + } +} +/// Issue one shared-to-global tensor store without a fence, commit or wait. +/// +/// This is the CuTeDSL primitives name for the shared-to-global +/// `cp.async.bulk.tensor` store. +/// +/// # Safety +/// Publish shared writes to the async proxy before issuing. Source and global +/// storage must remain live through the appropriate explicit completion waits. +#[inline(always)] +pub unsafe fn cp_async_bulk_tensor_global_shared_cta>( + src: &SharedRegion<'_, F>, + map: &TensorMap, + coords: Coordinates2D, +) { + unsafe { + cuda_device::tma::cp_async_bulk_tensor_2d_s2g( + src.address().as_raw().cast(), + map.as_descriptor(), + coords.inner, + coords.outer, + ) + } +} +/// Prefetch TMA descriptor metadata into the tensor-map cache (PTX `prefetch.tensormap`). +/// +/// `descriptor` is either a typed [`TensorMap`] reference or a raw +/// `*const TmaDescriptor` encoded by another tensor abstraction. The raw form +/// does not interpret element or shared-layout types, which permits +/// interoperability with layout-typed descriptors such as CuTe tensor maps +/// without adopting a CTM physical format. +/// +/// # Safety +/// `descriptor` must name a valid, 64-byte-aligned, 128-byte CUDA tensor map in +/// device-accessible storage that remains live through the prefetch. +#[inline(always)] +pub unsafe fn prefetch_tensormap(descriptor: D) { + unsafe { + cuda_device::tma::prefetch_tma_descriptor(sealed::Descriptor::descriptor_ptr(descriptor)) + } +} +/// Sealed proof that a value names a 128-byte TMA descriptor [`prefetch_tensormap`] may read. +pub trait TmaDescriptorRef: sealed::Descriptor {} +impl> sealed::Descriptor for &TensorMap { + #[inline(always)] + fn descriptor_ptr(self) -> *const TmaDescriptor { + self.as_descriptor() + } +} +impl> TmaDescriptorRef for &TensorMap {} +impl sealed::Descriptor for *const TmaDescriptor { + #[inline(always)] + fn descriptor_ptr(self) -> *const TmaDescriptor { + self + } +} +impl TmaDescriptorRef for *const TmaDescriptor {} +/// Commit this thread’s pending bulk copies as one group; does not wait. +#[inline(always)] +pub fn cp_async_bulk_commit_group() { + cuda_device::tma::cp_async_bulk_commit_group() +} +/// Limits bulk-copy wait counts and prefetch operands to the forms the instructions accept. +mod sealed { + /// Select a literal maximum number of committed groups that may remain pending. + pub trait Count { + /// Wait until at most N groups still read their shared source buffers. + fn wait(); + } + /// Resolve a descriptor reference to the raw descriptor bytes the instruction reads. + pub trait Descriptor { + /// Return the address of the 128-byte descriptor. + fn descriptor_ptr(self) -> *const cuda_device::TmaDescriptor; + } +} +/// Sealed proof that a bulk-copy wait immediate is in 0..=7. +pub trait ValidWaitCount: sealed::Count {} +/// Generate the supported immediate bulk-copy read waits. +macro_rules! waits { ($($n:literal),*) => { $(impl sealed::Count<$n> for () { #[inline(always)] fn wait() { cuda_device::tma::cp_async_bulk_wait_group_read($n) } } impl ValidWaitCount<$n> for () {})* }; } +waits!(0, 1, 2, 3, 4, 5, 6, 7); +/// Wait until at most `N` committed bulk groups still read their shared sources. +/// This is not proof that their global writes are visible to other threads. +#[inline(always)] +pub fn cp_async_bulk_wait_group_read() +where + (): ValidWaitCount, +{ + <() as sealed::Count>::wait() +} +use core::marker::PhantomData; + +/// Host-encoded tensor map matching one reviewed physical shared format. +/// +/// A copy of this descriptor does not extend the global allocation's lifetime. +#[repr(transparent)] +pub struct TensorMap> { + /// Opaque CUDA descriptor bytes; copying them does not retain the tensor allocation. + raw: TmaDescriptor, + /// Records the matching element type and physical shared format. + marker: PhantomData (T, F)>, +} +impl> Copy for TensorMap {} +impl> Clone for TensorMap { + fn clone(&self) -> Self { + *self + } +} +impl> TensorMap { + /// Adopt a descriptor encoded outside the typed host encoder. + /// + /// # Safety + /// Bytes must describe element type `T`, rank two and the tile/swizzle + /// specified by `F`. Underlying global storage must remain live through + /// every asynchronous operation using this map. + #[inline(always)] + pub const unsafe fn from_encoded(raw: TmaDescriptor) -> Self { + Self { + raw, + marker: PhantomData, + } + } + /// Borrow canonical opaque descriptor storage for instruction interoperability. + #[inline(always)] + pub const fn as_descriptor(&self) -> &TmaDescriptor { + &self.raw + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/src/warp.rs b/cuda-oxide/cute-rs-preview/ctm-rs/src/warp.rs new file mode 100644 index 0000000000..91120bcc94 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/src/warp.rs @@ -0,0 +1,269 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Explicit lane identity, election and register exchange. +//! No other primitive elects implicitly. + +/// Exchange an FP32 register with the lane selected by XORing `lane_mask`. +/// +/// This forwards the full-warp shuffle with participation mask `u32::MAX`. +/// All 32 lanes must participate with the same lane mask; no reduction or +/// memory synchronization is inserted. +#[inline(always)] +pub fn shuffle_xor_f32(value: f32, lane_mask: u32) -> f32 { + cuda_device::warp::shuffle_xor_f32(value, lane_mask) +} + +/// Elect one lane from `mask`, returning its lane ID and whether this lane won. +/// Every lane named in the mask must be converged at the call (sm_90+). +#[inline(always)] +pub fn elect_sync(mask: u32) -> (u32, bool) { + cuda_device::warp::elect_sync(mask) +} + +/// Return whether this lane is elected from `mask` (sm_90+). +/// Every lane named in the mask must be converged at the call. +#[inline(always)] +pub fn is_elected_sync(mask: u32) -> bool { + cuda_device::warp::is_elected_sync(mask) +} + +/// Return `threadIdx.x / 32`, the warp ID for a one-dimensional thread block. +#[inline(always)] +pub fn warp_id() -> u32 { + cuda_device::warp::warp_id() +} + +/// Read a participating lane's integer register under an explicit warp mask. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask, including the +/// caller. Each requested source lane must be in `0..32` and participate. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shuffle_sync(u32::MAX, 1, 0); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync(mask: u32, value: u32, src_lane: u32) -> u32 { + cuda_device::warp::shuffle_sync(mask, value, src_lane) +} + +/// Sum FP32 values across a full warp in XOR-offset order 16, 8, 4, 2, 1. +/// Every lane receives the sum. No memory synchronization is inserted. +/// +/// # Safety +/// All 32 lanes must be active and execute convergently. Missing logical +/// elements must contribute zero while their lanes continue to participate. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::reduce_sum_f32(1.0); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn reduce_sum_f32(value: f32) -> f32 { + cuda_device::warp::reduce_sum_f32(value) +} + +/// Return whether every participating lane's predicate is true, using a full-warp mask. +/// +/// # Safety +/// All 32 lanes must be active and execute this vote convergently. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::all(true); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn all(predicate: bool) -> bool { + cuda_device::warp::all(predicate) +} + +/// Synchronize warp lanes selected by an explicit participation mask (PTX `bar.warp.sync`). +/// +/// # Safety +/// The caller's lane must be included in `mask`. Every lane named in `mask` +/// must reach this call with the same mask. All accesses ordered by the barrier +/// must satisfy their allocation contracts. +/// +/// ```compile_fail,E0133 +/// ctm_rs::warp::bar_warp_sync(u32::MAX); +/// ``` +#[inline(always)] +pub unsafe fn bar_warp_sync(mask: u32) { + cuda_device::warp::sync_mask(mask) +} + +/// Exchange a register with the lane selected by XORing `lane_mask`. +/// This inserts no memory synchronization. +/// +/// # Safety +/// All 32 lanes must be active and execute convergently with the same lane +/// mask. `lane_mask` must select a source lane within the warp. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shuffle_xor(0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_xor(value: u32, lane_mask: u32) -> u32 { + cuda_device::warp::shuffle_xor(value, lane_mask) +} + +/// Return this lane's index within its warp, from the `%laneid` register. +#[must_use] +#[inline(always)] +pub fn lane_id() -> u32 { + cuda_device::warp::lane_id() +} + +/// Return the mask of the warp's currently active lanes (`activemask.b32`). +/// Bit `k` is set when lane `k` has neither exited nor diverged away here. +#[must_use] +#[inline(always)] +pub fn activemask() -> u32 { + cuda_device::warp::active_mask() +} + +/// Ballot across the lanes in `mask` (`vote.sync.ballot.b32`): bit `k` is set +/// when lane `k` is in `mask` and its predicate is true. +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::vote_sync_ballot(u32::MAX, true); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn vote_sync_ballot(mask: u32, predicate: bool) -> u32 { + cuda_device::warp::ballot_sync(mask, predicate) +} + +/// Return whether `predicate` holds for at least one lane in `mask` (`vote.sync.any.pred`). +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::vote_sync_any(u32::MAX, true); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn vote_sync_any(mask: u32, predicate: bool) -> bool { + cuda_device::warp::any_sync(mask, predicate) +} + +/// Return whether `predicate` holds for every lane in `mask` (`vote.sync.all.pred`). +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::vote_sync_all(u32::MAX, true); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn vote_sync_all(mask: u32, predicate: bool) -> bool { + cuda_device::warp::all_sync(mask, predicate) +} + +/// Exchange a register with the lane selected by XORing `lane_mask`, under an +/// explicit participation mask (`shfl.sync.bfly.b32`). This inserts no memory +/// synchronization. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and lane mask, +/// including the caller. `lane_mask` must select a source lane within the warp. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_bfly(u32::MAX, 0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_bfly(mask: u32, value: u32, lane_mask: u32) -> u32 { + cuda_device::warp::shuffle_xor_sync(mask, value, lane_mask) +} + +/// FP32 form of [`shfl_sync_bfly`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and lane mask, +/// including the caller. `lane_mask` must select a source lane within the warp. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_bfly_f32(u32::MAX, 0.0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_bfly_f32(mask: u32, value: f32, lane_mask: u32) -> f32 { + cuda_device::warp::shuffle_xor_f32_sync(mask, value, lane_mask) +} + +/// Read the register of the lane `delta` below this lane, under an explicit +/// participation mask (`shfl.sync.up.b32`); lanes below `delta` keep their own +/// value. This inserts no memory synchronization. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `delta`, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_up(u32::MAX, 0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_up(mask: u32, value: u32, delta: u32) -> u32 { + cuda_device::warp::shuffle_up_sync(mask, value, delta) +} + +/// FP32 form of [`shfl_sync_up`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `delta`, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_up_f32(u32::MAX, 0.0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_up_f32(mask: u32, value: f32, delta: u32) -> f32 { + cuda_device::warp::shuffle_up_f32_sync(mask, value, delta) +} + +/// Read the register of the lane `delta` above this lane, under an explicit +/// participation mask (`shfl.sync.down.b32`); the top `delta` lanes keep their +/// own value. This inserts no memory synchronization. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `delta`, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_down(u32::MAX, 0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_down(mask: u32, value: u32, delta: u32) -> u32 { + cuda_device::warp::shuffle_down_sync(mask, value, delta) +} + +/// FP32 form of [`shfl_sync_down`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `delta`, +/// including the caller. +/// +/// ```compile_fail,E0133 +/// let _ = ctm_rs::warp::shfl_sync_down_f32(u32::MAX, 0.0, 1); +/// ``` +#[must_use] +#[inline(always)] +pub unsafe fn shfl_sync_down_f32(mask: u32, value: f32, delta: u32) -> f32 { + cuda_device::warp::shuffle_down_f32_sync(mask, value, delta) +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/accumulator_tile.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/accumulator_tile.rs new file mode 100644 index 0000000000..003f94cd4d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/accumulator_tile.rs @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![doc = "Checks const-indexed accumulator fragments and static instruction expansion order."] +use ctm_rs::mma::{self, Accumulator16x8, AccumulatorTile16x8}; + +/// Const-indexed accumulator access must preserve fragment bands, columns and register order. +#[test] +fn static_fragment_access_preserves_bands_columns_and_register_order() { + let mut tile = AccumulatorTile16x8::<2, 8>::zero(); + mma::static_for!(ROW in [0, 1] { + mma::static_for!(COL in [0, 1, 2, 3, 4, 5, 6, 7] { + assert_eq!(tile.fragment::().registers(), [0.0; 4]); + let tag = (ROW * 100 + COL * 10) as f32; + *tile.fragment_mut::() = + Accumulator16x8::from_registers([tag, tag + 1.0, tag + 2.0, tag + 3.0]); + }); + }); + // Asymmetric tags detect a row/column swap or an aliased fragment update. + assert_eq!(tile.fragment::<0, 0>().registers(), [0., 1., 2., 3.]); + assert_eq!(tile.fragment::<0, 7>().registers(), [70., 71., 72., 73.]); + assert_eq!( + tile.fragment::<1, 0>().registers(), + [100., 101., 102., 103.] + ); + assert_eq!( + tile.fragment::<1, 7>().registers(), + [170., 171., 172., 173.] + ); +} + +/// Nested static expansion must visit the requested order with usable const-generic indices. +#[test] +fn static_expansion_nests_in_source_order_with_const_generic_indices() { + /// Return the statically supplied row and column so expansion order can be observed. + fn coordinate() -> (usize, usize) { + (ROW, COL) + } + let mut visited = Vec::new(); + mma::static_for!(ROW in [0, 1] { + mma::static_for!(COL in [0, 1] { + visited.push(coordinate::()); + }); + }); + assert_eq!(visited, [(0, 0), (0, 1), (1, 0), (1, 1)]); +} + +/// Static expansion must preserve an explicitly supplied nonconsecutive index list. +#[test] +fn static_expansion_preserves_the_requested_index_list() { + let mut visited = Vec::new(); + mma::static_for!(INDEX in [3, 1, 0, 1] { + visited.push(INDEX); + }); + assert_eq!(visited, [3, 1, 0, 1]); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/attention_api.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/attention_api.rs new file mode 100644 index 0000000000..3507d3b0ea --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/attention_api.rs @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Compile-pass coverage for attention primitives and canonical storage adapters."] +/// Compile the attention primitive probe with typed instruction operands and canonical storage. +#[test] +fn attention_primitives_accept_instruction_operands_and_canonical_storage() { + let tests = trybuild::TestCases::new(); + tests.pass("tests/compile_pass/attention_primitives.rs"); + tests.pass("tests/compile_pass/sm100_attention_primitives.rs"); +} + +/// The importer identifies loop markers by their defining crate before device inlining. +#[test] +fn unroll_marker_retains_its_canonical_definition() { + assert_eq!( + core::any::type_name_of_val(&ctm_rs::thread::__unroll_config::<0>), + core::any::type_name_of_val(&cuda_device::thread::__unroll_config::<0>), + ); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs new file mode 100644 index 0000000000..fe30f6f5cf --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a fragment row outside the accumulator tile during const evaluation. +use ctm_rs::mma::AccumulatorTile16x8; + +/// Valid two-row, eight-column accumulator tile used by the bounds probe. + +const TILE: AccumulatorTile16x8<2, 8> = AccumulatorTile16x8::zero(); +/// Force compile-time rejection of a fragment beyond the tile’s last row. +const OUTSIDE: ctm_rs::mma::Accumulator16x8 = TILE.fragment::<2, 0>(); + +/// Reject a fragment row outside the accumulator tile during const evaluation. + +fn main() { + let _ = OUTSIDE; +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.stderr new file mode 100644 index 0000000000..cd3ac71da7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_bounds.stderr @@ -0,0 +1,24 @@ +error[E0080]: evaluation panicked: accumulator fragment out of bounds + --> $RUST/core/src/panic.rs + | + | $crate::panicking::panic_fmt($crate::const_format_args!($($t)+)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `ctm_rs::mma::AccumulatorTile16x8::<2, 8>::fragment::<2, 0>::{constant#0}` failed here + | + ::: src/mma.rs + | + | / assert!( + | | ROW < ROWS && COL < COLS, + | | "accumulator fragment out of bounds" + | | ) + | |_____________- in this macro invocation + +note: erroneous constant encountered + --> src/mma.rs + | + | / const { + | | assert!( + | | ROW < ROWS && COL < COLS, + | | "accumulator fragment out of bounds" + | | ) + | | } + | |_________^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs new file mode 100644 index 0000000000..0cd3f9110d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject swapped A/B fragment and scale roles in tiled MMA. +use ctm_rs::mma::{AccumulatorTile16x8, Mxf4A, Mxf4B, ScalePairA, ScalePairB}; + +/// Reject swapped A/B fragment and scale roles in tiled MMA. + +fn main() { + let mut tile = AccumulatorTile16x8::<2, 8>::zero(); + unsafe { + tile.mma_mxf4_pair::<0>( + [Mxf4B::from_registers([0; 2]); 2], + [Mxf4A::from_registers([0; 4]); 2], + [ScalePairB::from_bits(0); 2], + [ScalePairA::from_bits(0); 2], + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr new file mode 100644 index 0000000000..e568182aec --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_roles.stderr @@ -0,0 +1,23 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:12:14 + | +12 | [Mxf4B::from_registers([0; 2]); 2], + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `Mxf4A`, found `Mxf4B` + +error[E0308]: mismatched types + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:13:14 + | +13 | [Mxf4A::from_registers([0; 4]); 2], + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `Mxf4B`, found `Mxf4A` + +error[E0308]: mismatched types + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:14:14 + | +14 | [ScalePairB::from_bits(0); 2], + | ^^^^^^^^^^^^^^^^^^^^^^^^ expected `ScalePairA`, found `ScalePairB` + +error[E0308]: mismatched types + --> tests/compile_fail/primitives_accumulator_tile_roles.rs:15:14 + | +15 | [ScalePairA::from_bits(0); 2], + | ^^^^^^^^^^^^^^^^^^^^^^^^ expected `ScalePairB`, found `ScalePairA` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs new file mode 100644 index 0000000000..00781c7bf9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject the two-row MMA method on a one-row accumulator tile. +use ctm_rs::mma::{AccumulatorTile16x8, Mxf4A, Mxf4B, ScalePairA, ScalePairB}; + +/// Reject the two-row MMA method on a one-row accumulator tile. + +fn main() { + let mut tile = AccumulatorTile16x8::<1, 8>::zero(); + unsafe { + tile.mma_mxf4_pair::<0>( + [Mxf4A::from_registers([0; 4]); 2], + [Mxf4B::from_registers([0; 2]); 2], + [ScalePairA::from_bits(0); 2], + [ScalePairB::from_bits(0); 2], + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr new file mode 100644 index 0000000000..1b2cf9f02a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_accumulator_tile_rows.stderr @@ -0,0 +1,7 @@ +error[E0599]: no method named `mma_mxf4_pair` found for struct `AccumulatorTile16x8<1, 8>` in the current scope + --> tests/compile_fail/primitives_accumulator_tile_rows.rs:11:14 + | +11 | tile.mma_mxf4_pair::<0>( + | -----^^^^^^^^^^^^^ method not found in `AccumulatorTile16x8<1, 8>` + | + = note: the method was found for `AccumulatorTile16x8<2, COLS>` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.rs new file mode 100644 index 0000000000..2fc24b349f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: keep local shared addresses distinct from cluster shared addresses. +use ctm_rs::memory::{SharedAddress,ClusterSharedAddress}; +/// Accept only a cluster address, making a local-address substitution fail. +fn cluster(_: ClusterSharedAddress<'_,u32>) {} +/// Keep local shared addresses distinct from cluster shared addresses. +fn wrong(address: SharedAddress<'_,u32>) { cluster(address); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.stderr new file mode 100644 index 0000000000..e29a79ca9e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_address_domain.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_address_domain.rs:10:52 + | +10 | fn wrong(address: SharedAddress<'_,u32>) { cluster(address); } + | ------- ^^^^^^^ expected `ClusterSharedAddress<'_, u32>`, found `SharedAddress<'_, u32>` + | | + | arguments to this function are incorrect + | + = note: expected struct `ClusterSharedAddress<'_, u32>` + found struct `SharedAddress<'_, u32>` +note: function defined here + --> tests/compile_fail/primitives_address_domain.rs:8:4 + | + 8 | fn cluster(_: ClusterSharedAddress<'_,u32>) {} + | ^^^^^^^ ------------------------------- diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs new file mode 100644 index 0000000000..1944cfbf2e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: require the cp.async wait count to be a compile-time immediate. +use ctm_rs::async_copy::cp_async_wait_group; + +/// Require the cp.async wait count to be a compile-time immediate. + +fn wrong(pending: u32) { + unsafe { cp_async_wait_group::<{ pending }>() }; +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr new file mode 100644 index 0000000000..323273f924 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_async_wait_runtime.stderr @@ -0,0 +1,7 @@ +error[E0435]: attempt to use a non-constant value in a constant + --> tests/compile_fail/primitives_async_wait_runtime.rs:11:38 + | +10 | fn wrong(pending: u32) { + | ------- this would need to be a `const` +11 | unsafe { cp_async_wait_group::<{ pending }>() }; + | ^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.rs new file mode 100644 index 0000000000..f45a018831 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.rs @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Reject implicit reinterpretation between BF16 and FP16 matrix operands. +use ctm_rs::mma::{Accumulator16x8, Bf16A, Bf16B, F16A, F16B, mma_bf16, mma_f16}; + +/// Neither instruction accepts the other dtype's packed register carrier. +fn wrong(a: Bf16A, b: Bf16B, half_a: F16A, half_b: F16B) { + unsafe { mma_bf16(Accumulator16x8::zero(), half_a, half_b) }; + unsafe { mma_f16(Accumulator16x8::zero(), a, b) }; +} + +/// The invalid device calls are type checked without executing GPU operations. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.stderr new file mode 100644 index 0000000000..8cc8e99965 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_dtype.stderr @@ -0,0 +1,27 @@ +error[E0308]: arguments to this function are incorrect + --> tests/compile_fail/primitives_bf16_mma_dtype.rs:10:14 + | +10 | unsafe { mma_bf16(Accumulator16x8::zero(), half_a, half_b) }; + | ^^^^^^^^ ------ ------ expected `Bf16B`, found `F16B` + | | + | expected `Bf16A`, found `F16A` + | +note: function defined here + --> src/mma.rs + | + | pub unsafe fn mma_bf16(c: Accumulator16x8, a: Bf16A, b: Bf16B) -> Accumulator16x8 { + | ^^^^^^^^ + +error[E0308]: arguments to this function are incorrect + --> tests/compile_fail/primitives_bf16_mma_dtype.rs:11:14 + | +11 | unsafe { mma_f16(Accumulator16x8::zero(), a, b) }; + | ^^^^^^^ - - expected `F16B`, found `Bf16B` + | | + | expected `F16A`, found `Bf16A` + | +note: function defined here + --> src/mma.rs + | + | pub unsafe fn mma_f16(c: Accumulator16x8, a: F16A, b: F16B) -> Accumulator16x8 { + | ^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.rs new file mode 100644 index 0000000000..5d2fceb1be --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Reject swapped BF16 matrix operand roles. +use ctm_rs::mma::{Accumulator16x8, Bf16A, Bf16B, mma_bf16}; + +/// Roles remain distinct even though both operands contain packed BF16 values. +fn wrong(a: Bf16A, b: Bf16B) { + unsafe { mma_bf16(Accumulator16x8::zero(), b, a) }; +} + +/// The invalid device call is type checked without executing a GPU operation. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.stderr new file mode 100644 index 0000000000..090c994fd6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bf16_mma_roles.stderr @@ -0,0 +1,18 @@ +error[E0308]: arguments to this function are incorrect + --> tests/compile_fail/primitives_bf16_mma_roles.rs:10:14 + | +10 | unsafe { mma_bf16(Accumulator16x8::zero(), b, a) }; + | ^^^^^^^^ - - expected `Bf16B`, found `Bf16A` + | | + | expected `Bf16A`, found `Bf16B` + | +note: function defined here + --> src/mma.rs + | + | pub unsafe fn mma_bf16(c: Accumulator16x8, a: Bf16A, b: Bf16B) -> Accumulator16x8 { + | ^^^^^^^^ +help: swap these arguments + | +10 - unsafe { mma_bf16(Accumulator16x8::zero(), b, a) }; +10 + unsafe { mma_bf16(Accumulator16x8::zero(), a, b) }; + | diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs new file mode 100644 index 0000000000..616f2ee6df --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject a bulk-copy wait count above seven. +use ctm_rs::tma::cp_async_bulk_wait_group_read; +/// Reject a bulk-copy wait count above seven. +fn wrong() { cp_async_bulk_wait_group_read::<8>(); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr new file mode 100644 index 0000000000..54bef41b69 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_count.stderr @@ -0,0 +1,23 @@ +error[E0277]: the trait bound `(): ValidWaitCount<8>` is not satisfied + --> tests/compile_fail/primitives_bulk_wait_count.rs:8:46 + | +8 | fn wrong() { cp_async_bulk_wait_group_read::<8>(); } + | ^ the trait `ValidWaitCount<8>` is not implemented for `()` + | + = help: `()` implements trait `ValidWaitCount`: + ValidWaitCount<0> + ValidWaitCount<1> + ValidWaitCount<2> + ValidWaitCount<3> + ValidWaitCount<4> + ValidWaitCount<5> + ValidWaitCount<6> + ValidWaitCount<7> +note: required by a bound in `cp_async_bulk_wait_group_read` + --> src/tma.rs + | + | pub fn cp_async_bulk_wait_group_read() + | ----------------------------- required by a bound in this function + | where + | (): ValidWaitCount, + | ^^^^^^^^^^^^^^^^^ required by this bound in `cp_async_bulk_wait_group_read` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs new file mode 100644 index 0000000000..02d160af6c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: require bulk-copy wait counts as const generics rather than runtime arguments. +use ctm_rs::tma::cp_async_bulk_wait_group_read; +/// Require bulk-copy wait counts as const generics rather than runtime arguments. +fn wrong(n:u32) { cp_async_bulk_wait_group_read(n); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr new file mode 100644 index 0000000000..942d1dded7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_bulk_wait_runtime.stderr @@ -0,0 +1,16 @@ +error[E0061]: this function takes 0 arguments but 1 argument was supplied + --> tests/compile_fail/primitives_bulk_wait_runtime.rs:8:19 + | +8 | fn wrong(n:u32) { cp_async_bulk_wait_group_read(n); } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - unexpected argument of type `u32` + | +note: function defined here + --> src/tma.rs + | + | pub fn cp_async_bulk_wait_group_read() + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +help: remove the extra argument + | +8 - fn wrong(n:u32) { cp_async_bulk_wait_group_read(n); } +8 + fn wrong(n:u32) { cp_async_bulk_wait_group_read(); } + | diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs new file mode 100644 index 0000000000..8ae2553fc0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject swapped FP16 MMA A and B register operands. +use ctm_rs::mma::{Accumulator16x8, F16A, F16B, mma_f16}; + +/// Reject swapped FP16 MMA A and B register operands. + +fn wrong(a: F16A, b: F16B) { + unsafe { mma_f16(Accumulator16x8::zero(), b, a) }; +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr new file mode 100644 index 0000000000..1b3fabc7bd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_f16_mma_roles.stderr @@ -0,0 +1,18 @@ +error[E0308]: arguments to this function are incorrect + --> tests/compile_fail/primitives_f16_mma_roles.rs:11:14 + | +11 | unsafe { mma_f16(Accumulator16x8::zero(), b, a) }; + | ^^^^^^^ - - expected `F16B`, found `F16A` + | | + | expected `F16A`, found `F16B` + | +note: function defined here + --> src/mma.rs + | + | pub unsafe fn mma_f16(c: Accumulator16x8, a: F16A, b: F16B) -> Accumulator16x8 { + | ^^^^^^^ +help: swap these arguments + | +11 - unsafe { mma_f16(Accumulator16x8::zero(), b, a) }; +11 + unsafe { mma_f16(Accumulator16x8::zero(), a, b) }; + | diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs new file mode 100644 index 0000000000..0e3302dae9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: prevent access to a tensor-memory guard after deallocation consumes it. +use ctm_rs::{tcgen05::tmem::{TmemGuard,TmemReady},groups::CtaGroup2}; +/// Prevent access to a tensor-memory guard after deallocation consumes it. +fn wrong(g:TmemGuard) { let _ = unsafe { g.dealloc_collective() }; let _ = g.raw_address(); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr new file mode 100644 index 0000000000..95999c7b85 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_guard_consumed.stderr @@ -0,0 +1,14 @@ +error[E0382]: borrow of moved value: `g` + --> tests/compile_fail/primitives_guard_consumed.rs:8:101 + | +8 | fn wrong(g:TmemGuard) { let _ = unsafe { g.dealloc_collective() }; let _ = g.raw_address(); } + | - -------------------- ^ value borrowed here after move + | | | + | | `g` moved due to this method call + | move occurs because `g` has type `TmemGuard`, which does not implement the `Copy` trait + | +note: `ctm_rs::tcgen05::tmem::TmemGuard::::dealloc_collective` takes ownership of the receiver `self`, which moves `g` + --> src/tcgen05/tmem.rs + | + | pub unsafe fn dealloc_collective(self) -> TmemGuard { + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs new file mode 100644 index 0000000000..42f77f0bf2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject a 33-column tensor-memory allocation. +use ctm_rs::{tcgen05::tmem::{TmemGuard,TmemUninit},groups::CtaGroup2}; +/// Reject a 33-column tensor-memory allocation. +fn wrong(slot:*mut u32) { let _ = unsafe { TmemGuard::::from_raw_slot(slot) }; } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr new file mode 100644 index 0000000000..d78b042085 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_invalid_columns.stderr @@ -0,0 +1,8 @@ +error[E0599]: the associated function or constant `from_raw_slot` exists for struct `TmemGuard`, but its trait bounds were not satisfied + --> tests/compile_fail/primitives_invalid_columns.rs:8:82 + | +8 | fn wrong(slot:*mut u32) { let _ = unsafe { TmemGuard::::from_raw_slot(slot) }; } + | ^^^^^^^^^^^^^ associated function or constant cannot be called on `TmemGuard` due to unsatisfied trait bounds + | + = note: the following trait bounds were not satisfied: + `(): ValidColumns<33>` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs new file mode 100644 index 0000000000..fdbb520f0d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject an N160 B tile in an N192 instruction. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject an N160 B tile in an N192 instruction. +fn wrong(d:&impl TmemOperand>,a:&SmemOperand,b:&SmemOperand) { + unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr new file mode 100644 index 0000000000..d6b7774669 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_b_shape.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_mma_b_shape.rs:14:45 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; + | ----------------------------- ^ expected `96`, found `80` + | | + | arguments to this function are incorrect + | + = note: expected reference `&SmemOperand<'_, OperandB, Swizzled2D<96, 64, Swizzle128>>` + found reference `&SmemOperand<'_, OperandB, Swizzled2D<80, 64, Swizzle128>>` +note: function defined here + --> src/tcgen05/mma.rs + | + | pub unsafe fn mma( + | ^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs new file mode 100644 index 0000000000..26ae19312a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.rs @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject unsupported CTA-group, M-dimension and operand-major combinations. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Unsupported external major-mode marker. +struct MnMajor; +/// Reject unsupported CTA-group, M-dimension and operand-major combinations. +fn wrong() { + let _ = InstructionDescriptor::>::new(); + let _ = InstructionDescriptor::>::new(); + let _ = InstructionDescriptor::>::new(); +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr new file mode 100644 index 0000000000..a328f98b84 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_configuration.stderr @@ -0,0 +1,122 @@ +error[E0277]: the trait bound `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` is not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:16:10 + | +16 | let _ = InstructionDescriptor::>::new(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | + = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` +help: the following other types implement trait `MmaSpec` + --> src/tcgen05/descriptors.rs + | + | impl MmaSpec for DenseSharedF16 { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` +... + | spec!(192, F16TileB96); + | ---------------------- in this macro invocation + | spec!(160, F16TileB80); + | ---------------------- in this macro invocation +note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescriptor` + --> src/tcgen05/descriptors.rs + | + | pub struct InstructionDescriptor { + | ^^^^^^^ required by this bound in `InstructionDescriptor` + = note: this error originates in the macro `spec` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the associated function or constant `new` exists for struct `ctm_rs::tcgen05::descriptors::InstructionDescriptor>`, but its trait bounds were not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:16:84 + | +16 | let _ = InstructionDescriptor::>::new(); + | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds + | + ::: src/tcgen05/descriptors.rs + | + | pub struct DenseSharedF16( + | ------------------------------------------------------------------------ doesn't satisfy `_: MmaSpec` + | + = note: the following trait bounds were not satisfied: + `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` + +error[E0277]: the trait bound `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` is not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:17:10 + | +17 | let _ = InstructionDescriptor::>::new(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | + = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` +help: the following other types implement trait `MmaSpec` + --> src/tcgen05/descriptors.rs + | + | impl MmaSpec for DenseSharedF16 { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` +... + | spec!(192, F16TileB96); + | ---------------------- in this macro invocation + | spec!(160, F16TileB80); + | ---------------------- in this macro invocation +note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescriptor` + --> src/tcgen05/descriptors.rs + | + | pub struct InstructionDescriptor { + | ^^^^^^^ required by this bound in `InstructionDescriptor` + = note: this error originates in the macro `spec` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the associated function or constant `new` exists for struct `ctm_rs::tcgen05::descriptors::InstructionDescriptor>`, but its trait bounds were not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:17:84 + | +17 | let _ = InstructionDescriptor::>::new(); + | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds + | + ::: src/tcgen05/descriptors.rs + | + | pub struct DenseSharedF16( + | ------------------------------------------------------------------------ doesn't satisfy `_: MmaSpec` + | + = note: the following trait bounds were not satisfied: + `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` + +error[E0277]: the trait bound `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` is not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:18:10 + | +18 | let _ = InstructionDescriptor::>::new(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ unsatisfied trait bound + | + = help: the trait `MmaSpec` is not implemented for `ctm_rs::tcgen05::descriptors::DenseSharedF16` +help: the following other types implement trait `MmaSpec` + --> src/tcgen05/descriptors.rs + | + | impl MmaSpec for DenseSharedF16 { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` + | `ctm_rs::tcgen05::descriptors::DenseSharedF16` +... + | spec!(192, F16TileB96); + | ---------------------- in this macro invocation + | spec!(160, F16TileB80); + | ---------------------- in this macro invocation +note: required by a bound in `ctm_rs::tcgen05::descriptors::InstructionDescriptor` + --> src/tcgen05/descriptors.rs + | + | pub struct InstructionDescriptor { + | ^^^^^^^ required by this bound in `InstructionDescriptor` + = note: this error originates in the macro `spec` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the associated function or constant `new` exists for struct `ctm_rs::tcgen05::descriptors::InstructionDescriptor>`, but its trait bounds were not satisfied + --> tests/compile_fail/primitives_mma_configuration.rs:18:85 + | +18 | let _ = InstructionDescriptor::>::new(); + | ^^^ associated function or constant cannot be called due to unsatisfied trait bounds + | + ::: src/tcgen05/descriptors.rs + | + | pub struct DenseSharedF16( + | ------------------------------------------------------------------------ doesn't satisfy `_: MmaSpec` + | + = note: the following trait bounds were not satisfied: + `ctm_rs::tcgen05::descriptors::DenseSharedF16: MmaSpec` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs new file mode 100644 index 0000000000..f66ae38120 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject an N160 instruction descriptor paired with an N192 MMA specification. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject an N160 instruction descriptor paired with an N192 MMA specification. +fn wrong(d:&impl TmemOperand>,a:&SmemOperand,b:&SmemOperand) { + unsafe { mma::(d,a,b,&InstructionDescriptor::::new(),false) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr new file mode 100644 index 0000000000..e98583d409 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_descriptor.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_mma_descriptor.rs:14:47 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::::new(),false) }; + | ----------------------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `192`, found `160` + | | + | arguments to this function are incorrect + | + = note: expected reference `&InstructionDescriptor>` + found reference `&InstructionDescriptor>` +note: function defined here + --> src/tcgen05/mma.rs + | + | pub unsafe fn mma( + | ^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.rs new file mode 100644 index 0000000000..2ce590baa0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject an N160 tensor-memory destination for an N192 MMA instruction. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject an N160 tensor-memory destination for an N192 MMA instruction. +fn wrong(d:&impl TmemOperand>,a:&SmemOperand,b:&SmemOperand) { + unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr new file mode 100644 index 0000000000..0e624927fe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_destination.stderr @@ -0,0 +1,17 @@ +error[E0271]: type mismatch resolving `> as TmemOperand>::Format == TmemF32<256, 192>` + --> tests/compile_fail/primitives_mma_destination.rs:14:41 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; + | ----------------------------- ^ expected `192`, found `160` + | | + | required by a bound introduced by this call + | + = note: expected struct `TmemF32<256, 192>` + found struct `TmemF32<256, 160>` +note: required by a bound in `ctm_rs::tcgen05::mma::mma` + --> src/tcgen05/mma.rs + | + | pub unsafe fn mma( + | --- required by a bound in this function + | dst: &impl TmemOperand, + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `mma` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs new file mode 100644 index 0000000000..f2d42eac7c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.rs @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject K16 slice index four in a four-slice shared operand. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject K16 slice index four in a four-slice shared operand. +fn wrong(a:&SmemOperand) { let _ = a.k16::<4>(); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr new file mode 100644 index 0000000000..11de33be85 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_k_step.stderr @@ -0,0 +1,27 @@ +error[E0277]: the trait bound `(): ValidKStep<4>` is not satisfied + --> tests/compile_fail/primitives_mma_k_step.rs:13:65 + | +13 | fn wrong(a:&SmemOperand) { let _ = a.k16::<4>(); } + | ^ the trait `ValidKStep<4>` is not implemented for `()` + | +help: `()` implements trait `ValidKStep` + --> src/tcgen05/descriptors.rs + | + | macro_rules! step { ($($n:literal),*) => { $(impl sealed::Step<$n> for () {} impl ValidKStep<$n> for () {})* }; } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | + | `ValidKStep<0>` + | `ValidKStep<1>` + | `ValidKStep<2>` + | `ValidKStep<3>` + | step!(0, 1, 2, 3); + | ----------------- in this macro invocation +note: required by a bound in `ctm_rs::tcgen05::descriptors::SmemOperand::<'s, R, F>::k16` + --> src/tcgen05/descriptors.rs + | + | pub fn k16(&self) -> Self + | --- required by a bound in this associated function + | where + | (): ValidKStep, + | ^^^^^^^^^^^^^^^^ required by this bound in `SmemOperand::<'s, R, F>::k16` + = note: this error originates in the macro `step` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs new file mode 100644 index 0000000000..ea12f5bf92 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.rs @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject the output-tile format as a shared MMA input operand. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject the output-tile format as a shared MMA input operand. +fn wrong(a:&SharedRegion) { let _ = SmemOperand::::from_region(a); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr new file mode 100644 index 0000000000..134227ce3a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_output_format.stderr @@ -0,0 +1,37 @@ +error[E0277]: the trait bound `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` is not satisfied + --> tests/compile_fail/primitives_mma_output_format.rs:13:47 + | +13 | fn wrong(a:&SharedRegion) { let _ = SmemOperand::::from_region(a); } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `MmaInputFormat` is not implemented for `Swizzled2D<128, 32, Swizzle64>` + | +help: the following other types implement trait `MmaInputFormat` + --> src/tcgen05/descriptors.rs + | + | macro_rules! input { ($($f:ty),*) => { $(impl sealed::Input for $f {} impl MmaInputFormat for $f {})* }; } + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | + | `Swizzled2D<128, 64, Swizzle128>` + | `Swizzled2D<80, 64, Swizzle128>` + | `Swizzled2D<96, 64, Swizzle128>` + | input!(F16Tile128, F16TileB96, F16TileB80); + | ------------------------------------------ in this macro invocation +note: required by a bound in `ctm_rs::tcgen05::descriptors::SmemOperand` + --> src/tcgen05/descriptors.rs + | + | pub struct SmemOperand<'s, R: OperandRole, F: MmaInputFormat> { + | ^^^^^^^^^^^^^^ required by this bound in `SmemOperand` + = note: this error originates in the macro `input` (in Nightly builds, run with -Z macro-backtrace for more info) + +error[E0599]: the associated function or constant `from_region` exists for struct `ctm_rs::tcgen05::descriptors::SmemOperand<'_, ctm_rs::tcgen05::descriptors::OperandA, Swizzled2D<128, 32, Swizzle64>>`, but its trait bounds were not satisfied + --> tests/compile_fail/primitives_mma_output_format.rs:13:81 + | +13 | fn wrong(a:&SharedRegion) { let _ = SmemOperand::::from_region(a); } + | ^^^^^^^^^^^ associated function or constant cannot be called due to unsatisfied trait bounds + | + ::: src/memory.rs + | + | pub struct Swizzled2D( + | -------------------------------------------------------------------- doesn't satisfy `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` + | + = note: the following trait bounds were not satisfied: + `Swizzled2D<128, 32, Swizzle64>: MmaInputFormat` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.rs new file mode 100644 index 0000000000..8bbac278fe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject shared MMA operands with reversed A and B roles. +#![allow(unused_imports)] +use ctm_rs::{groups::{CtaGroup1,CtaGroup2},memory::*,tcgen05::{descriptors::*,mma::*,tmem::*}}; +/// Supported group-two M256/N192 specification used to distinguish incompatible MMA configurations. +type L=DenseSharedF16; +/// Supported group-two M256/N160 specification used to distinguish incompatible MMA configurations. +type R=DenseSharedF16; +/// Reject shared MMA operands with reversed A and B roles. +fn wrong(d:&impl TmemOperand>,a:&SmemOperand,b:&SmemOperand) { + unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr new file mode 100644 index 0000000000..63f0bfe4ad --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_mma_roles.stderr @@ -0,0 +1,25 @@ +error[E0308]: arguments to this function are incorrect + --> tests/compile_fail/primitives_mma_roles.rs:14:11 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | +note: expected `&SmemOperand<'_, OperandA, _>`, found `&SmemOperand<'_, OperandB, _>` + --> tests/compile_fail/primitives_mma_roles.rs:14:43 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; + | ^ + = note: expected reference `&SmemOperand<'_, ctm_rs::tcgen05::descriptors::OperandA, Swizzled2D<128, 64, Swizzle128>>` + found reference `&SmemOperand<'_, ctm_rs::tcgen05::descriptors::OperandB, Swizzled2D<128, 64, Swizzle128>>` +note: expected `&SmemOperand<'_, OperandB, _>`, found `&SmemOperand<'_, OperandA, _>` + --> tests/compile_fail/primitives_mma_roles.rs:14:45 + | +14 | unsafe { mma::(d,a,b,&InstructionDescriptor::new(),false) }; + | ^ + = note: expected reference `&SmemOperand<'_, ctm_rs::tcgen05::descriptors::OperandB, Swizzled2D<96, 64, Swizzle128>>` + found reference `&SmemOperand<'_, ctm_rs::tcgen05::descriptors::OperandA, Swizzled2D<96, 64, Swizzle128>>` +note: function defined here + --> src/tcgen05/mma.rs + | + | pub unsafe fn mma( + | ^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs new file mode 100644 index 0000000000..34f4d8d0ca --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: require explicit collective allocation rather than an implicit guard allocation method. +use ctm_rs::{tcgen05::tmem::{TmemGuard,TmemUninit},groups::{CtaGroup1,CtaGroup2}}; +/// Reject implicit allocation for a group-one TMEM guard. +fn group1(g:TmemGuard) { let _ = unsafe { g.alloc() }; } +/// Reject implicit allocation for a group-two TMEM guard. +fn group2(g:TmemGuard) { let _ = unsafe { g.alloc() }; } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr new file mode 100644 index 0000000000..dc48b33275 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_no_implicit_allocation.stderr @@ -0,0 +1,11 @@ +error[E0599]: no method named `alloc` found for struct `TmemGuard` in the current scope + --> tests/compile_fail/primitives_no_implicit_allocation.rs:8:71 + | +8 | fn group1(g:TmemGuard) { let _ = unsafe { g.alloc() }; } + | ^^^^^ method not found in `TmemGuard` + +error[E0599]: no method named `alloc` found for struct `TmemGuard` in the current scope + --> tests/compile_fail/primitives_no_implicit_allocation.rs:10:71 + | +10 | fn group2(g:TmemGuard) { let _ = unsafe { g.alloc() }; } + | ^^^^^ method not found in `TmemGuard` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs new file mode 100644 index 0000000000..149351a384 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: register counts must be multiples of eight. +use ctm_rs::thread::setmaxnreg_dec; +/// Trigger validation of a count that is in range but not aligned. +fn main() { + unsafe { setmaxnreg_dec::<25>() }; +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr new file mode 100644 index 0000000000..0f732206a4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_alignment.stderr @@ -0,0 +1,45 @@ +error[E0277]: the trait bound `(): ValidRegisterCount<25>` is not satisfied + --> tests/compile_fail/primitives_register_count_alignment.rs:9:31 + | +9 | unsafe { setmaxnreg_dec::<25>() }; + | ^^ the trait `ValidRegisterCount<25>` is not implemented for `()` + | + = help: `()` implements trait `ValidRegisterCount`: + ValidRegisterCount<104> + ValidRegisterCount<112> + ValidRegisterCount<120> + ValidRegisterCount<128> + ValidRegisterCount<136> + ValidRegisterCount<144> + ValidRegisterCount<152> + ValidRegisterCount<160> + ValidRegisterCount<168> + ValidRegisterCount<176> + ValidRegisterCount<184> + ValidRegisterCount<192> + ValidRegisterCount<200> + ValidRegisterCount<208> + ValidRegisterCount<216> + ValidRegisterCount<224> + ValidRegisterCount<232> + ValidRegisterCount<240> + ValidRegisterCount<248> + ValidRegisterCount<24> + ValidRegisterCount<256> + ValidRegisterCount<32> + ValidRegisterCount<40> + ValidRegisterCount<48> + ValidRegisterCount<56> + ValidRegisterCount<64> + ValidRegisterCount<72> + ValidRegisterCount<80> + ValidRegisterCount<88> + ValidRegisterCount<96> +note: required by a bound in `setmaxnreg_dec` + --> src/lib.rs + | + | pub unsafe fn setmaxnreg_dec() + | -------------- required by a bound in this function + | where + | (): ValidRegisterCount, + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `setmaxnreg_dec` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.rs new file mode 100644 index 0000000000..b70838154c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: register counts may not exceed 256. +use ctm_rs::thread::setmaxnreg_inc; +/// Trigger validation of an aligned count beyond the instruction range. +fn main() { + unsafe { setmaxnreg_inc::<264>() }; +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr new file mode 100644 index 0000000000..022ee950f2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_range.stderr @@ -0,0 +1,45 @@ +error[E0277]: the trait bound `(): ValidRegisterCount<264>` is not satisfied + --> tests/compile_fail/primitives_register_count_range.rs:9:31 + | +9 | unsafe { setmaxnreg_inc::<264>() }; + | ^^^ the trait `ValidRegisterCount<264>` is not implemented for `()` + | + = help: `()` implements trait `ValidRegisterCount`: + ValidRegisterCount<104> + ValidRegisterCount<112> + ValidRegisterCount<120> + ValidRegisterCount<128> + ValidRegisterCount<136> + ValidRegisterCount<144> + ValidRegisterCount<152> + ValidRegisterCount<160> + ValidRegisterCount<168> + ValidRegisterCount<176> + ValidRegisterCount<184> + ValidRegisterCount<192> + ValidRegisterCount<200> + ValidRegisterCount<208> + ValidRegisterCount<216> + ValidRegisterCount<224> + ValidRegisterCount<232> + ValidRegisterCount<240> + ValidRegisterCount<248> + ValidRegisterCount<24> + ValidRegisterCount<256> + ValidRegisterCount<32> + ValidRegisterCount<40> + ValidRegisterCount<48> + ValidRegisterCount<56> + ValidRegisterCount<64> + ValidRegisterCount<72> + ValidRegisterCount<80> + ValidRegisterCount<88> + ValidRegisterCount<96> +note: required by a bound in `setmaxnreg_inc` + --> src/lib.rs + | + | pub unsafe fn setmaxnreg_inc() + | -------------- required by a bound in this function + | where + | (): ValidRegisterCount, + | ^^^^^^^^^^^^^^^^^^^^^ required by this bound in `setmaxnreg_inc` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs new file mode 100644 index 0000000000..541dee1bb8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: a register count is an instruction immediate. +use ctm_rs::thread::setmaxnreg_inc; +/// Reject a runtime register count. +fn wrong(count: u32) { + unsafe { setmaxnreg_inc::<{ count }>() }; +} +/// Unused host entry point. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr new file mode 100644 index 0000000000..af7e13a5a0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_register_count_runtime.stderr @@ -0,0 +1,7 @@ +error[E0435]: attempt to use a non-constant value in a constant + --> tests/compile_fail/primitives_register_count_runtime.rs:9:33 + | +8 | fn wrong(count: u32) { + | ----- this would need to be a `const` +9 | unsafe { setmaxnreg_inc::<{ count }>() }; + | ^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.rs new file mode 100644 index 0000000000..126b54f139 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject a local shared address passed to a remote barrier arrival. +use cuda_device::barrier::Barrier; +use ctm_rs::{memory::SharedAddress,cluster::arrive_remote}; +/// Reject a local shared address passed to a remote barrier arrival. +fn wrong(local:SharedAddress) { unsafe { arrive_remote(local) }; } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr new file mode 100644 index 0000000000..defc28ec88 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_domain.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_remote_domain.rs:9:65 + | +9 | fn wrong(local:SharedAddress) { unsafe { arrive_remote(local) }; } + | ------------- ^^^^^ expected `ClusterSharedAddress<'_, Barrier>`, found `SharedAddress<'_, Barrier>` + | | + | arguments to this function are incorrect + | + = note: expected struct `ClusterSharedAddress<'_, cuda_device::Barrier>` + found struct `SharedAddress<'_, cuda_device::Barrier>` +note: function defined here + --> src/cluster.rs + | + | pub unsafe fn arrive_remote(address: ClusterSharedAddress<'_, Barrier>) { + | ^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs new file mode 100644 index 0000000000..6329b3d9fb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.rs @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject a local shared address passed to a relaxed remote arrival. +use ctm_rs::{cluster::arrive_remote_relaxed, memory::SharedAddress}; +use cuda_device::barrier::Barrier; +/// Reject a local shared address passed to a relaxed remote arrival. +fn wrong(local: SharedAddress) { + unsafe { arrive_remote_relaxed(local) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr new file mode 100644 index 0000000000..abfc2f8a92 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_remote_relaxed_domain.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_remote_relaxed_domain.rs:10:36 + | +10 | unsafe { arrive_remote_relaxed(local) }; + | --------------------- ^^^^^ expected `ClusterSharedAddress<'_, Barrier>`, found `SharedAddress<'_, Barrier>` + | | + | arguments to this function are incorrect + | + = note: expected struct `ClusterSharedAddress<'_, cuda_device::Barrier>` + found struct `SharedAddress<'_, cuda_device::Barrier>` +note: function defined here + --> src/cluster.rs + | + | pub unsafe fn arrive_remote_relaxed(address: ClusterSharedAddress<'_, Barrier>) { + | ^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs new file mode 100644 index 0000000000..2af3e7c34f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject an FP32 descriptor element type paired with an FP16 physical tile. +#![feature(f16)] +use ctm_rs::{memory::F16Tile128,tma::TensorMap}; +/// Reject an FP32 descriptor element type paired with an FP16 physical tile. +fn wrong(_: &TensorMap) {} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr new file mode 100644 index 0000000000..50f6c8b8c4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tensor_map_element.stderr @@ -0,0 +1,11 @@ +error[E0271]: type mismatch resolving ` as PhysicalFormat>::Element == f32` + --> tests/compile_fail/primitives_tensor_map_element.rs:9:13 + | +9 | fn wrong(_: &TensorMap) {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `f32`, found `f16` + | +note: required by a bound in `TensorMap` + --> src/tma.rs + | + | pub struct TensorMap> { + | ^^^^^^^^^^^ required by this bound in `TensorMap` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.rs new file mode 100644 index 0000000000..d02b4b221d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: reject a shared output tile paired with a different TMA descriptor format. +#![feature(f16)] +use ctm_rs::{memory::*,tma::*}; +/// Reject a shared output tile paired with a different TMA descriptor format. +fn wrong(src:&SharedRegion,map:&TensorMap) { unsafe { cp_async_bulk_tensor_global_shared_cta(src,map,Coordinates2D::new(0,0)) }; } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.stderr new file mode 100644 index 0000000000..7d35a1b8e7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tma_format.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_tma_format.rs:9:124 + | +9 | ...F16Tile128>) { unsafe { cp_async_bulk_tensor_global_shared_cta(src,map,Coordinates2D::new(0,0)) }; } + | -------------------------------------- ^^^ expected `32`, found `64` + | | + | arguments to this function are incorrect + | + = note: expected reference `&TensorMap>` + found reference `&TensorMap>` +note: function defined here + --> src/tma.rs + | + | pub unsafe fn cp_async_bulk_tensor_global_shared_cta>( + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs new file mode 100644 index 0000000000..822af8e9bb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: TMEM transfers retain their exact canonical register count. +use ctm_rs::{memory::CuSimd, tcgen05::tmem::store_32x32b_x4}; +/// Reject a store with too few register operands. +unsafe fn wrong(address: u32, data: CuSimd) { + unsafe { store_32x32b_x4(address, data) }; +} +/// Unused host entry point. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr new file mode 100644 index 0000000000..d9c28e9d4b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_tmem_carrier_width.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/primitives_tmem_carrier_width.rs:9:39 + | +9 | unsafe { store_32x32b_x4(address, data) }; + | --------------- ^^^^ expected `4`, found `2` + | | + | arguments to this function are incorrect + | + = note: expected struct `CuSimd` + found struct `CuSimd` +note: function defined here + --> src/tcgen05/tmem.rs + | + | pub unsafe fn store_32x32b_x4(address: u32, data: crate::memory::CuSimd) { + | ^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs new file mode 100644 index 0000000000..1d84b47b26 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: prevent extracting an instruction address before the barrier is initialized. +use cuda_device::barrier::{ManagedBarrier,Uninit,TmaBarrier}; +use ctm_rs::mbarrier::BarrierAddressExt; +/// Prevent extracting an instruction address before the barrier is initialized. +fn wrong(b:ManagedBarrier) { let _ = b.shared_address(); } +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr new file mode 100644 index 0000000000..e6077af33f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_unready_barrier.stderr @@ -0,0 +1,5 @@ +error[E0599]: no method named `shared_address` found for struct `ManagedBarrier` in the current scope + --> tests/compile_fail/primitives_unready_barrier.rs:9:59 + | +9 | fn wrong(b:ManagedBarrier) { let _ = b.shared_address(); } + | ^^^^^^^^^^^^^^ method not found in `ManagedBarrier` diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs new file mode 100644 index 0000000000..171b511e97 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-fail probe: keep a tensor-memory allocation borrowed until its last view use. +use ctm_rs::{groups::CtaGroup2,tcgen05::tmem::{TmemGuard,TmemReady,TmemF32,ColumnOffset,TmemOperand}}; +/// Keep a TMEM view live after the attempted deallocation. +fn consume_view(_:impl TmemOperand) {} +/// Keep a tensor-memory allocation borrowed until its last view use. +fn wrong(g:TmemGuard) { + let view = g.view::>(ColumnOffset::new(0)).unwrap(); + let _dead = unsafe { g.dealloc_collective() }; + consume_view(view); +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr new file mode 100644 index 0000000000..1be02f4a62 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_fail/primitives_view_after_deallocation.stderr @@ -0,0 +1,11 @@ +error[E0505]: cannot move out of `g` because it is borrowed + --> tests/compile_fail/primitives_view_after_deallocation.rs:12:26 + | +10 | fn wrong(g:TmemGuard) { + | - binding `g` declared here +11 | let view = g.view::>(ColumnOffset::new(0)).unwrap(); + | - borrow of `g` occurs here +12 | let _dead = unsafe { g.dealloc_collective() }; + | ^ move out of `g` occurs here +13 | consume_view(view); + | ---- borrow later used here diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/attention_primitives.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/attention_primitives.rs new file mode 100644 index 0000000000..599f608bcf --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/attention_primitives.rs @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-pass probe: type-check attention instruction primitives and their canonical storage adapters. +#![allow(dead_code, unused_imports)] + +use ctm_rs::memory::{DisjointSlice, DynamicSharedArray}; +use ctm_rs::mma::{Accumulator16x8, F16A, F16B, mma_f16}; +use ctm_rs::{ + async_copy, convert, cuda_module, float, kernel, launch_bounds, launch_contract, warp, +}; + +// A valid instruction sequence type-checks without executing device operations +// on the host. A and B have different register counts and operand roles. +/// Type-check copy waits, FP16 MMA, warp reduction, softmax arithmetic and FP16 packing; never executed. +unsafe fn attention_step(a: [u32; 4], b: [u32; 2], c: [f32; 4]) -> u32 { + unsafe { + async_copy::cp_async_commit_group(); + async_copy::cp_async_wait_group::<0>(); + async_copy::cp_async_wait_group::<7>(); + let c = mma_f16( + Accumulator16x8::from_registers(c), + F16A::from_registers(a), + F16B::from_registers(b), + ) + .registers(); + let maximum = warp::shuffle_xor_f32(c[0], 2); + let prefill = float::ex2_approx_ftz_f32(c[0] - maximum); + let decode = float::ex2_approx_f32(c[1] - maximum); + let reciprocal = float::rcp_approx_ftz_f32(prefill + decode); + convert::cvt_f16x2_f32(prefill * reciprocal, decode * reciprocal) + } +} + +// Storage aliases must retain the canonical types used by importer and macros. +/// Verify the ctm-rs alias is exactly the canonical device disjoint-slice type. +fn canonical_slice<'s>(slice: DisjointSlice<'s, f32>) -> cuda_device::DisjointSlice<'s, f32> { + slice +} +/// Verify the ctm-rs alias preserves the canonical dynamic-shared-storage type. +fn canonical_dynamic_storage( + storage: DynamicSharedArray, +) -> cuda_device::DynamicSharedArray { + storage +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs new file mode 100644 index 0000000000..30137d3a96 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/compile_pass/sm100_attention_primitives.rs @@ -0,0 +1,80 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Compile-pass probe for SM100 raw instruction operands and canonical carriers. +#![allow(dead_code)] +use ctm_rs::{async_copy, f32x2, fence, float, mbarrier, memory, tcgen05, thread, warp}; + +/// Canonical carrier identity is retained across both public crates. +fn canonical_simd(value: memory::CuSimd) -> cuda_device::CuSimd { + value +} + +/// Canonical packed global-memory element identity is retained. +fn canonical_u32x4(value: memory::U32x4) -> cuda_device::vector::U32x4 { + value +} + +/// Type-check explicit copy, barrier, group-1 commit and TMEM instructions; never run on host. +unsafe fn sm100_step(bar: *mut mbarrier::Barrier, shared: *mut u32, global: *const u32, addr: u32) { + unsafe { + mbarrier::mbarrier_init(bar, 1); + let _ = mbarrier::mbarrier_arrive(bar); + let _ = mbarrier::mbarrier_arrive_expect_tx(bar, 1, 4); + let _ = mbarrier::mbarrier_arrive_expect_tx_cluster(bar, 1, 4); + let _ = mbarrier::mbarrier_try_wait_parity(bar, 0); + let _ = mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, 1); + fence::fence_proxy_async_acquire_sync_restrict(); + fence::fence_proxy_async_release_sync_restrict(); + fence::fence_sync_restrict_acquire(); + async_copy::cp_async_ca_4_raw(shared, global); + async_copy::cp_async_mbarrier_arrive_noinc_raw(bar); + tcgen05::commit_address(memory::SharedAddress::from_raw(bar)); + thread::setmaxnreg_inc::<256>(); + thread::setmaxnreg_dec::<24>(); + let data: cuda_device::CuSimd = tcgen05::tmem::load_32x32b_x2(addr); + tcgen05::tmem::load_wait(); + tcgen05::tmem::store_32x32b_x2(addr, data); + let data: cuda_device::CuSimd = tcgen05::tmem::load_32x32b_x4(addr); + tcgen05::tmem::load_wait(); + tcgen05::tmem::store_32x32b_x4(addr, data); + let data: cuda_device::CuSimd = tcgen05::tmem::load_32x32b_x8(addr); + tcgen05::tmem::load_wait(); + tcgen05::tmem::store_32x32b_x8(addr, data); + let data: cuda_device::CuSimd = tcgen05::tmem::load_32x32b_x16(addr); + tcgen05::tmem::load_wait(); + tcgen05::tmem::store_32x32b_x16(addr, data); + let data: cuda_device::CuSimd = tcgen05::tmem::load_32x32b_x32(addr); + tcgen05::tmem::load_wait(); + tcgen05::tmem::store_32x32b_x32(addr, data); + tcgen05::tmem::store_wait(); + } + unsafe { tcgen05::relinquish_alloc_permit() }; + let x = f32x2::pack_f32x2(1.0, 2.0); + let y = f32x2::add_f32x2(x, x); + let z = f32x2::mul_f32x2(x, y); + let _ = f32x2::fma_f32x2(x, y, z); + let _ = float::lg2_approx_f32(f32x2::unpack_f32x2_hi(z)); + let _ = unsafe { warp::all(true) }; + let _ = unsafe { warp::shuffle_xor(0, 1) }; + let _ = unsafe { warp::shuffle_sync(u32::MAX, 0, 0) }; + unsafe { warp::bar_warp_sync(u32::MAX) }; + let mut i = 0; + while i < 2 { + thread::__unroll_config::<0>(); + i += 1; + } + let mut j = 0; + while j < 2048 { + thread::__unroll_config::<1024>(); + j += 1; + } +} + +/// Host pack/unpack operations preserve both lanes' exact bits. +fn main() { + let packed = f32x2::pack_f32x2(f32::from_bits(0x8000_0000), f32::from_bits(0x7fc1_2345)); + assert_eq!(f32x2::unpack_f32x2_lo(packed).to_bits(), 0x8000_0000); + assert_eq!(f32x2::unpack_f32x2_hi(packed).to_bits(), 0x7fc1_2345); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/core_intrinsics.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/core_intrinsics.rs new file mode 100644 index 0000000000..9cda3ecf2f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/core_intrinsics.rs @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Core facade contracts: bit conversions, fragment order and collective types. + +use ctm_rs::{convert, debug, float, mma, thread, warp}; + +/// Macro rewriting must preserve the CTM prefix and the canonical scope witness. +#[ctm_rs::kernel] +pub fn scoped_index_probe(mut output: ctm_rs::memory::DisjointSlice) { + if let Some(value) = output.get_mut(thread::index_1d()) { + *value = 1; + } +} + +/// Conversions retain ties-to-even, signed zeros, infinities and NaNs. +#[test] +fn bf16_rounding_keeps_ties_even_and_special_values_distinct() { + for (input, expected) in [ + (0x0000_0000, 0x0000), + (0x8000_0000, 0x8000), + (0x3f80_7fff, 0x3f80), + (0x3f80_8000, 0x3f80), + (0x3f80_8001, 0x3f81), + (0x3f81_8000, 0x3f82), + (0xbf80_8000, 0xbf80), + (0xbf81_8000, 0xbf82), + (0x7f80_0000, 0x7f80), + (0xff80_0000, 0xff80), + (0x7f80_0001, 0x7fc0), + (0xff80_0001, 0xffc0), + ] { + assert_eq!(convert::f32_to_bf16_rne(f32::from_bits(input)), expected); + } +} + +/// Every sixteen-bit payload widens exactly, including noncanonical NaNs. +#[test] +fn bf16_widening_preserves_every_storage_pattern() { + for bits in 0..=u16::MAX { + assert_eq!(convert::bf16_to_f32(bits).to_bits(), u32::from(bits) << 16); + } +} + +/// Fragment pairs stay adjacent within each row and cover the whole MMA tile once. +#[test] +fn accumulator_coordinates_preserve_pair_order_and_bijection() { + const LAST: (usize, usize) = mma::acc_coords(31, 3); + assert_eq!(LAST, (15, 7)); + assert_eq!(mma::acc_coords(0, 0), (0, 0)); + assert_eq!(mma::acc_coords(0, 1), (0, 1)); + assert_eq!(mma::acc_coords(0, 2), (8, 0)); + assert_eq!(mma::acc_coords(0, 3), (8, 1)); + assert_eq!(mma::acc_coords(31, 0), (7, 6)); + assert_eq!(mma::acc_coords(32, 4), (0, 0)); + let mut visits = [[0u8; 8]; 16]; + for lane in 0..32 { + for register in 0..4 { + let (row, col) = mma::acc_coords(lane, register); + visits[row][col] += 1; + } + } + assert!(visits.into_iter().flatten().all(|count| count == 1)); +} + +/// Collective participation stays unsafe without executing GPU instructions on the host. +#[test] +fn collective_and_instruction_signatures_preserve_their_contracts() { + let _: unsafe fn(u32, u32, u32) -> u32 = warp::shuffle_sync; + let _: unsafe fn(f32) -> f32 = warp::reduce_sum_f32; + let _: unsafe fn(bool) -> bool = warp::all; + let _: unsafe fn(u32) = warp::bar_warp_sync; + let _: unsafe fn() = thread::setmaxnreg_inc::<24>; + let _: unsafe fn() = thread::setmaxnreg_inc::<256>; + let _: unsafe fn() = thread::setmaxnreg_dec::<24>; + let _: unsafe fn() = thread::setmaxnreg_dec::<256>; + let _: fn(f32) -> f32 = float::rsqrt_approx_f32; + let _: fn(f32) -> f32 = float::lg2_approx_f32; + let _: fn(f32) -> f32 = float::sin_approx_f32; + let _: fn() -> ! = debug::trap; + let _: fn() -> thread::ThreadIndex<'static> = thread::index_1d; +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/packed_memory.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/packed_memory.rs new file mode 100644 index 0000000000..512da3c95b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/packed_memory.rs @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Packed global words preserve payloads without numeric conversion. + +#[test] +fn aligned_word_load_preserves_each_half_and_word_offset() { + let words = [0x3f80_8000u32, 0x7fc1_0001, 0xffff_7f80]; + let halves = words.as_ptr().cast::(); + for (index, expected) in words.iter().enumerate() { + // SAFETY: a u32 allocation guarantees alignment, initialization and + // four readable bytes at each even u16 offset for this whole borrow. + let actual = unsafe { ctm_rs::memory::load_aligned_u32(halves.add(index * 2).cast()) }; + assert_eq!(actual, *expected); + } +} + +#[test] +fn raw_word_load_requires_an_explicit_unsafe_boundary() { + let _: unsafe fn(*const u32) -> u32 = ctm_rs::memory::load_aligned_u32; +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_mxf4_storage.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_mxf4_storage.rs new file mode 100644 index 0000000000..625afdb10c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_mxf4_storage.rs @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![doc = "Checks MXFP4 physical storage validation and typed local TMA signatures."] +use ctm_rs::{ + memory::{ + F16Tile128, MemoryError, Mxf4ScaleTile, Mxf4Tile128, PhysicalFormat, SharedAddress, + SharedRegion, + }, + tma::{Coordinates2D, TensorMap, cp_async_bulk_tensor_shared_cta_global}, +}; +use cuda_device::barrier::Barrier; + +/// MXFP4 value and scale regions must report exact instruction footprints and element types. +#[test] +fn blockscaled_stage_storage_matches_instruction_footprints() { + assert_eq!(Mxf4Tile128::FOOTPRINT_BYTES, 8192); + assert_eq!(Mxf4Tile128::BASE_ALIGNMENT, 512); + assert_eq!(Mxf4Tile128::SWIZZLE_BYTES, 64); + assert_eq!(Mxf4ScaleTile::FOOTPRINT_BYTES, 512); + assert_eq!(Mxf4ScaleTile::BASE_ALIGNMENT, 128); + assert_eq!(Mxf4ScaleTile::SWIZZLE_BYTES, 0); + assert_eq!(F16Tile128::FOOTPRINT_BYTES, 16384); + assert_eq!(size_of::>(), 128); + assert_eq!(size_of::>(), 128); + assert_eq!(align_of::>(), 64); + assert_eq!(align_of::>(), 64); +} + +/// Region construction must reject insufficient capacity and misaligned shared bases. +#[test] +fn blockscaled_storage_rejects_short_or_misaligned_stages() { + unsafe { + assert!( + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1080), + 512 + ) + .is_ok() + ); + assert!(matches!( + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1000), + 8191 + ), + Err(MemoryError::InsufficientCapacity) + )); + assert!(matches!( + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1010), + 8192 + ), + Err(MemoryError::Misaligned) + )); + assert!(matches!( + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1080), + 511 + ), + Err(MemoryError::InsufficientCapacity) + )); + assert!(matches!( + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1010), + 512 + ), + Err(MemoryError::Misaligned) + )); + } +} + +/// Both value and scale load signatures must use a local shared completion barrier. +#[test] +// The explicit instruction signatures are the contract checked by this test. +#[allow(clippy::type_complexity)] +fn ordinary_tma_load_uses_cta_local_completion_for_both_stage_types() { + let _: unsafe fn( + &SharedRegion<'_, Mxf4Tile128>, + &TensorMap, + Coordinates2D, + SharedAddress<'_, Barrier>, + ) = cp_async_bulk_tensor_shared_cta_global::; + let _: unsafe fn( + &SharedRegion<'_, Mxf4ScaleTile>, + &TensorMap, + Coordinates2D, + SharedAddress<'_, Barrier>, + ) = cp_async_bulk_tensor_shared_cta_global::; +} + +/// Successful trapping constructors must retain the same validated addresses as checked constructors. +#[test] +fn trapping_regions_preserve_checked_constructor_validation_and_address() { + /// Aligned host buffer used to test numeric region validation without GPU accesses. + #[repr(align(1024))] + struct Storage( + #[doc = "Backing bytes with enough room for every tested physical format."] [u8; 32768], + ); + + /// Compare checked and trapping region constructors for one supported physical format. + fn check() + where + F::Element: std::panic::RefUnwindSafe, + { + let mut storage = Storage([0; 32768]); + let base = storage.0.as_mut_ptr().cast::(); + // Host-side representation check only: neither constructor dereferences + // the pointer. Device execution must supply actual shared storage. + for offset in [0, F::BASE_ALIGNMENT] { + let pointer = base.cast::().wrapping_add(offset).cast(); + for capacity in [F::FOOTPRINT_BYTES, F::FOOTPRINT_BYTES + F::BASE_ALIGNMENT] { + let region = + unsafe { SharedRegion::::from_raw_parts_or_trap(pointer, capacity) }; + assert_eq!(region.address().as_raw(), pointer); + } + } + + for (pointer, bytes, error) in [ + ( + core::ptr::null_mut(), + F::FOOTPRINT_BYTES, + MemoryError::NullAddress, + ), + ( + base.cast::().wrapping_add(1).cast(), + F::FOOTPRINT_BYTES, + MemoryError::Misaligned, + ), + ( + base, + F::FOOTPRINT_BYTES - 1, + MemoryError::InsufficientCapacity, + ), + ] { + assert!(matches!( + unsafe { SharedRegion::::from_raw_parts(pointer, bytes) }, + Err(actual) if actual == error + )); + let failure = std::panic::catch_unwind(|| unsafe { + SharedRegion::::from_raw_parts_or_trap(pointer, bytes); + }) + .expect_err("invalid storage must reach the device trap"); + let message = failure + .downcast_ref::<&str>() + .copied() + .or_else(|| failure.downcast_ref::().map(String::as_str)); + assert!( + message + .unwrap_or_default() + .contains("trap called outside CUDA kernel context") + ); + } + } + check::(); + check::(); + check::(); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_safety.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_safety.rs new file mode 100644 index 0000000000..8a461409a1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_safety.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![doc = "Compile-fail coverage for instruction operands, address domains and hardware resource lifetimes."] +/// Compile-fail probes must reject invalid shapes, address domains, roles and resource lifetimes. +#[test] +fn primitives_reject_invalid_type_contracts() { + trybuild::TestCases::new().compile_fail("tests/compile_fail/primitives_*.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_sm120.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_sm120.rs new file mode 100644 index 0000000000..ae76090995 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_sm120.rs @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#![doc = "Checks SM120 register carriers and packed scale-byte order."] +use ctm_rs::mma::{Accumulator16x8, Mxf4A, Mxf4B, ScalePairA, ScalePairB}; +/// Register and scale carriers must preserve exact instruction ordering and packed bytes. +#[test] +fn register_carriers_preserve_instruction_order_and_scale_bytes() { + assert_eq!( + Mxf4A::from_registers([1, 2, 3, 4]).registers(), + [1, 2, 3, 4] + ); + assert_eq!(Mxf4B::from_registers([5, 6]).registers(), [5, 6]); + assert_eq!(Accumulator16x8::zero().registers(), [0.0; 4]); + assert_eq!( + Accumulator16x8::from_registers([1., 2., 3., 4.]).registers(), + [1., 2., 3., 4.] + ); + assert_eq!(ScalePairA::from_bits(0x807e).bits(), 0x807e); + assert_eq!(ScalePairB::from_bits(0x7e80).bits(), 0x7e80); +} diff --git a/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_types.rs b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_types.rs new file mode 100644 index 0000000000..76396ec518 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/ctm-rs/tests/primitives_types.rs @@ -0,0 +1,87 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![doc = "Checks instruction-facing storage formats, opaque descriptor ABI and typed address domains."] +#![feature(f16)] +use ctm_rs::{memory::*, tma::TensorMap}; + +/// Physical formats and opaque tensor maps must preserve expected footprints, alignment and Copy behavior. +#[test] +fn physical_formats_and_tensor_map_preserve_kernel_abi() { + assert_eq!(size_of::>(), 128); + assert_eq!(align_of::>(), 64); + assert_eq!(F16Tile128::FOOTPRINT_BYTES, 16384); + assert_eq!(F16TileB96::FOOTPRINT_BYTES, 12288); + assert_eq!(F16TileB80::FOOTPRINT_BYTES, 10240); + assert_eq!(F16TileC::FOOTPRINT_BYTES, 8192); + assert_eq!(F16Tile128::BASE_ALIGNMENT, 1024); + assert_eq!(F16TileC::BASE_ALIGNMENT, 512); + /// Require the tensor-map type to implement Copy without adding marker bounds. + fn copy() {} + copy::>(); +} + +/// Shared-region checks must reject null, short or misaligned storage before exposing an address. +#[test] +fn shared_region_rejects_invalid_storage_before_use() { + // These invalid layouts return Err before any address can be used. + assert!( + unsafe { SharedRegion::::from_raw_parts(core::ptr::null_mut(), 16384) } + .is_err() + ); + assert!( + unsafe { + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1002), + 16384, + ) + } + .is_err() + ); + assert!( + unsafe { + SharedRegion::::from_raw_parts( + core::ptr::without_provenance_mut(0x1000), + 16383, + ) + } + .is_err() + ); +} + +/// Masks must remain nonempty and signed tensor coordinates must retain negative values. +#[test] +fn multicast_masks_and_signed_coordinates_preserve_instruction_domains() { + use ctm_rs::{cluster::CtaMask, groups::CtaGroup2, tma::Coordinates2D}; + assert!(CtaMask::::try_from_bits(0).is_err()); + assert_eq!( + CtaMask::::try_from_bits(0b1111).unwrap().bits(), + 0b1111 + ); + assert_eq!(Coordinates2D::new(-32, 4096).inner(), -32); + assert_eq!(Coordinates2D::new(-32, 4096).outer(), 4096); +} + +/// The relaxed remote-arrival signature must retain the cluster-shared address domain. +#[test] +fn relaxed_remote_arrival_requires_a_cluster_barrier_address() { + use ctm_rs::cluster::arrive_remote_relaxed; + use cuda_device::barrier::Barrier; + let _: for<'a> unsafe fn(ClusterSharedAddress<'a, Barrier>) = arrive_remote_relaxed; +} + +/// BF16 operands retain packed bits and distinct roles without changing the FP32 result carrier. +#[test] +fn bf16_mma_carriers_preserve_bits_roles_and_fp32_accumulation() { + use ctm_rs::mma::{Accumulator16x8, Bf16A, Bf16B, mma_bf16}; + let a = [0x8000_0000, 0x7fc1_3f80, 0xffff_7f80, 0x0001_ff80]; + let b = [0xbf80_3f80, 0x7f81_0080]; + assert_eq!(Bf16A::from_registers(a).registers(), a); + assert_eq!(Bf16B::from_registers(b).registers(), b); + assert_eq!(size_of::(), size_of::<[u32; 4]>()); + assert_eq!(align_of::(), align_of::<[u32; 4]>()); + assert_eq!(size_of::(), size_of::<[u32; 2]>()); + assert_eq!(align_of::(), align_of::<[u32; 2]>()); + let _: unsafe fn(Accumulator16x8, Bf16A, Bf16B) -> Accumulator16x8 = mma_bf16; +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/Cargo.toml b/cuda-oxide/cute-rs-preview/cute-layout/Cargo.toml new file mode 100644 index 0000000000..6fd31e22c6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "cute-layout" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Coordinate-to-memory mapping and validation for CuTe-style GPU tiles" + +[dependencies] diff --git a/cuda-oxide/cute-rs-preview/cute-layout/README.md b/cuda-oxide/cute-rs-preview/cute-layout/README.md new file mode 100644 index 0000000000..d5795d4b10 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/README.md @@ -0,0 +1,51 @@ +# Layout algebra for cuda-oxide + +`cute-layout` describes and evaluates coordinate-to-memory mappings. It is a +`no_std` crate using `alloc`, shared by the Rust tensor API and the compiler. + +```rust +use cute_layout::{IntTuple, Layout}; + +let layout: Layout = "(2,3):(3,1)".parse().unwrap(); +assert_eq!(layout.checked_call(&IntTuple::Leaf(1)), Some(3)); +``` + +## Module navigation + +| Module | Responsibility | +| --- | --- | +| [`tuple`](src/tuple.rs) | Nested `IntTuple` shapes, strides and coordinates | +| [`layout`](src/layout.rs) | `Layout` construction and coordinate evaluation | +| [`parse`](src/parse.rs) | Text formatting, parsing and `ParseLayoutError` | +| [`algebra::mapping`](src/algebra/mapping.rs) | Flattening, coalescing, composition, complements, inverse maps and reshaping | +| [`algebra::tiling`](src/algebra/tiling.rs) | Logical and zipped division, blocked products and raked products | +| [`algebra::thread_value`](src/algebra/thread_value.rs) | Thread/value ownership maps | +| [`algebra::alignment`](src/algebra/alignment.rs) | Static vector-access alignment | +| [`composed`](src/composed.rs) and [`swizzle`](src/swizzle.rs) | Composed addresses, offset units and shared-memory swizzles | +| [`predication`](src/predication.rs) | Safe source selection for partial copy atoms | +| [`validation`](src/validation/mod.rs) | Bounded validation for thread/value layouts, cooperative copies, TMA and ldmatrix | + +The root imports `IntTuple`, `Layout`, `ParseLayoutError` and `crd2idx` remain +available. Every existing `algebra::*` function remains available through an +explicit reexport. Grouped imports such as `cute_layout::layout::Layout` and +`cute_layout::algebra::mapping::composition` refer to the same definitions. + +Use checked evaluation and the validation modules for layouts decoded from +external compiler input. The unchecked constructors and algebra retain their +documented arithmetic and shape preconditions. + +## Validation + +Run from `cuda-oxide/` with its pinned toolchain: + +```sh +cargo test -p cute-layout +cargo clippy -p cute-layout --all-targets -- -D warnings +RUSTDOCFLAGS="-D warnings" cargo doc -p cute-layout --no-deps +cargo fmt -p cute-layout --check +``` + +Tests cover the defining algebra equations, checked arithmetic, representation, +parsing, static validation, and interoperability of the grouped and existing +public paths. Algebra contract tests live beside their shared fixtures under +[`src/algebra/tests`](src/algebra/tests). diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/alignment.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/alignment.rs new file mode 100644 index 0000000000..35c38b64c1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/alignment.rs @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Static vector-access alignment derived from layout mappings and division. + +use crate::Layout; + +use super::mapping::{coalesce, right_inverse}; +use super::tiling::logical_divide; +use super::{AlgebraError, Result}; + +/// Returns the alignment factor guaranteed by the layout's offset pattern. +/// +/// This examines only the static strides; it does not inspect a memory +/// pointer. For `(8,16):(1,4)`, the first coordinate gives an eight-offset +/// unit-stride run, while changing the second coordinate shifts that run by +/// four: +/// +/// ```text +/// layout coordinate offset walk +/// (0 through 7, 0) 0 1 2 3 4 5 6 7 run length 8 +/// (0 through 7, 1) 4 5 6 7 8 9 10 11 shifted by 4 +/// +/// greatest common divisor of 8 and 4 = 4 offset units +/// ``` +/// +/// The result is `4` because width 4 divides both the eight-unit run and its +/// four-unit shift. Four-unit vector chunks therefore have compatible starts. +/// This does **not** prove that the base pointer is aligned; the compiler must +/// check the pointer separately before selecting a vector instruction. +/// +/// The unit is whatever the layout's strides count: elements for an +/// element-stride layout or bytes for a byte-stride layout. A negative stride +/// returns `1` because this analysis cannot prove a larger factor for a +/// backwards walk. +pub fn max_alignment(layout: &Layout) -> Result { + let shapes = layout.shape.flatten(); + let strides = layout.stride.flatten(); + if shapes.iter().any(|&shape| shape <= 0) { + return Err(AlgebraError( + "max_alignment: every shape extent must be positive".into(), + )); + } + if strides.iter().any(|&stride| stride < 0) { + return Ok(1); + } + + let flat = coalesce(layout); + let inverse = right_inverse(&flat)?; + let permuted = logical_divide(&flat, &inverse)?; + let modes = permuted.modes(); + if modes.len() != 2 { + return Err(AlgebraError( + "max_alignment: expected vector and remainder modes".into(), + )); + } + + let vector_size = modes[0].size(); + let rest_stride_gcd = modes[1] + .stride + .flatten() + .into_iter() + .fold(0, gcd_nonnegative); + Ok(gcd_nonnegative(vector_size, rest_stride_gcd).max(1)) +} + +/// Compute the greatest common divisor of nonnegative layout alignment factors. +fn gcd_nonnegative(a: i64, b: i64) -> i64 { + let mut a = a.unsigned_abs(); + let mut b = b.unsigned_abs(); + while b != 0 { + (a, b) = (b, a % b); + } + // Both inputs came from i64 layout values, so their greatest common + // divisor also fits in i64. + a as i64 +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mapping.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mapping.rs new file mode 100644 index 0000000000..0e78c51a64 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mapping.rs @@ -0,0 +1,494 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Simplification, composition, inverse maps, and reshaping. + +use alloc::format; +use alloc::vec; +use alloc::vec::Vec; + +use crate::{IntTuple, Layout}; + +use super::{AlgebraError, Result, flat_modes, layout_from_flat}; + +/// Removes the parentheses that group [modes](crate#modes), as CuTe's +/// `flatten`. +/// +/// ```text +/// ((2,3),4):((12,4),1) -> (2,3,4):(12,4,1) +/// ``` +/// +/// Every flat coordinate reaches the same offset as before; only the nesting +/// of the tuple coordinate changes, so the result has one mode per integer +/// of the original shape. A layout with a single integer mode is returned +/// without parentheses. [`IntTuple::flatten`] does the same for one tuple. +pub fn flatten(l: &Layout) -> Layout { + layout_from_flat(flat_modes(l)) +} + +/// Rewrites a layout as fewer [modes](crate#modes) while preserving its flat +/// coordinate-to-offset sequence. +/// +/// A **one-cell mode** means an `extent:stride` pair whose extent is `1`. Its +/// coordinate can only be zero, so it never changes the offset: +/// +/// ```text +/// before = (1,4):(99,1) +/// +/// mode pair 1:99 4:1 +/// allowed coordinate only 0 0, 1, 2, 3 +/// offset contribution 0 * 99 = 0 0, 1, 2, 3 +/// +/// tuple coordinate (0,0) (0,1) (0,2) (0,3) +/// before offset 0 1 2 3 +/// after = 4:1 0 1 2 3 +/// ``` +/// +/// **Adjacent modes** are consecutive shape/stride pairs inside the same +/// layout, not two separate shapes. In `(2,4):(1,2)`, they are `2:1` and +/// `4:2`: +/// +/// ```text +/// tuple coordinate (0,0) (1,0) | (0,1) (1,1) | (0,2) (1,2) | (0,3) (1,3) +/// before offset 0 1 | 2 3 | 4 5 | 6 7 +/// flat coordinate 0 1 2 3 4 5 6 7 +/// after = 8:1 0 1 2 3 4 5 6 7 +/// ``` +/// +/// The pairs merge because the second stride is `2 = 2 * 1`: its first step +/// starts immediately after visiting both positions of the first pair. In +/// general, `s0:d0` and `s1:d1` merge when `d1 == s0 * d0`. +/// +/// Coalescing means the same flat sequence can be written as one +/// constant-stride walk. It does **not** require adjacent memory: +/// `(2,4):(3,6)` maps to `0,3,6,...,21` and coalesces to `8:3`. Successive +/// offsets are physically adjacent only when the resulting stride has +/// magnitude `1`. +/// +/// Tuple coordinate structure may change during this rewrite. For example, +/// `(2,4):(1,2)` accepts `(first, second)`, while `8:1` accepts one flat +/// coordinate; their flat walks are the same. +pub fn coalesce(l: &Layout) -> Layout { + let mut out: Vec<(i64, i64)> = Vec::new(); + for (s, d) in flat_modes(l) { + if s == 1 { + continue; + } + if let Some(last) = out.last_mut() + && d == last.0 * last.1 + { + last.0 *= s; + continue; + } + out.push((s, d)); + } + layout_from_flat(out) +} + +/// Removes broadcast modes before simplifying a layout. +/// +/// A stride-zero mode repeats the same offset for every coordinate: +/// `3:0` maps `0,1,2 -> 0,0,0`. It therefore occupies no additional offsets +/// when [`complement`] asks which locations remain unused. +fn filter(l: &Layout) -> Layout { + let kept: Vec<(i64, i64)> = flat_modes(l) + .into_iter() + .map(|(s, d)| if d == 0 { (1, 0) } else { (s, d) }) + .collect(); + coalesce(&layout_from_flat(kept)) +} + +/// Replaces two consecutive mappings with one direct mapping. +/// +/// `b` first turns the input coordinate into a coordinate for `a`; `a` then +/// turns that coordinate into the final offset. For example: +/// +/// ```text +/// a = 8:2 coordinate 0 1 2 3 4 5 6 7 +/// offset 0 2 4 6 8 10 12 14 +/// +/// b = 4:2 input c 0 1 2 3 +/// b(c) 0 2 4 6 coordinates selected in a +/// a(b(c)) 0 4 8 12 final offsets +/// +/// result = 4:4 +/// input c 0 1 2 3 +/// offset 0 4 8 12 +/// ``` +/// +/// In symbols, `result(c) = a(b(c))`, also written `a ∘ b`. The compiler +/// uses this to replace chains such as "thread and value -> tile cell -> +/// memory offset" with one direct offset calculation. +/// +/// If `b` has several [modes](crate#modes), the result keeps their nesting. +/// Some input walks cannot be represented by the supported constant-stride +/// mode structure; those cases return [`AlgebraError`]. An imperfect final +/// tile may be rounded up, so later predication must disable coordinates past +/// the real tensor boundary. +pub fn composition(a: &Layout, b: &Layout) -> Result { + if let (IntTuple::Tuple(bs), IntTuple::Tuple(bd)) = (&b.shape, &b.stride) { + let modes = bs + .iter() + .zip(bd) + .map(|(s, d)| { + composition( + a, + &Layout { + shape: s.clone(), + stride: d.clone(), + }, + ) + }) + .collect::>>()?; + // Each output mode corresponds to the same input mode in `b`; only + // that mode's mapping has passed through `a`. + return Ok(Layout { + shape: IntTuple::Tuple(modes.iter().map(|m| m.shape.clone()).collect()), + stride: IntTuple::Tuple(modes.into_iter().map(|m| m.stride).collect()), + }); + } + + let (IntTuple::Leaf(b_shape), IntTuple::Leaf(b_stride)) = (&b.shape, &b.stride) else { + return Err(AlgebraError( + "composition: rhs shape/stride not congruent".into(), + )); + }; + compose_int(a, *b_shape, *b_stride) +} + +/// Composes `a` with one input mode described by `b_shape:b_stride`. +fn compose_int(a: &Layout, b_shape: i64, b_stride: i64) -> Result { + // A stride-zero input always selects coordinate zero in `a`, so the + // result is the same broadcast regardless of `a`'s strides. + if b_stride == 0 { + return Ok(layout_from_flat(vec![(b_shape, 0)])); + } + let ac = coalesce(a); + let modes = flat_modes(&ac); + // With one mode in `a`, each step of `b` advances `b_stride` coordinates, + // and each of those advances `a` by its stride. Multiplication combines + // the two steps. + if modes.len() == 1 { + return Ok(layout_from_flat(vec![(b_shape, b_stride * modes[0].1)])); + } + + let mut result: Vec<(i64, i64)> = Vec::new(); + let mut rest_shape = b_shape; + let mut rest_stride = b_stride; + for &(curr_shape, curr_stride) in &modes[..modes.len() - 1] { + // A step through `b` must align with a digit boundary in `a`, or be + // small enough to live entirely inside the current mode. Otherwise a + // simple affine mode cannot represent the composed walk. + if rest_stride % curr_shape != 0 && rest_stride.abs() >= curr_shape { + return Err(AlgebraError(format!( + "composition: stride {rest_stride} does not divide lhs mode of shape {curr_shape} \ + (stride divisibility condition)" + ))); + } + let abs_rest_stride = rest_stride.abs(); + let sign = if rest_stride < 0 { -1 } else { 1 }; + let next_shape = ceil_div(curr_shape, abs_rest_stride); + let next_stride = ceil_div(abs_rest_stride, curr_shape) * sign; + + if next_shape == 1 || rest_shape == 1 { + rest_stride = next_stride; + continue; + } + let new_shape = next_shape.min(rest_shape); + // Every non-final piece must split the remaining coordinates into + // equal groups. The final piece is allowed to describe a partial + // overhang, which higher-level predication will mask. + if rest_shape % new_shape != 0 { + return Err(AlgebraError(format!( + "composition: shape {rest_shape} is not divisible by {new_shape} \ + (shape divisibility condition)" + ))); + } + result.push((new_shape, rest_stride * curr_stride)); + rest_shape /= new_shape; + rest_stride = next_stride; + } + + let last_stride = modes[modes.len() - 1].1; + if result.is_empty() { + return Ok(layout_from_flat(vec![( + rest_shape, + rest_stride * last_stride, + )])); + } + if rest_shape != 1 { + result.push((rest_shape, rest_stride * last_stride)); + } + Ok(coalesce(&layout_from_flat(result))) +} + +/// Round a nonnegative extent upward to a whole positive-divisor tile count. +fn ceil_div(a: i64, b: i64) -> i64 { + (a + b - 1) / b +} + +/// Builds the extra coordinate needed to fill the offsets that `l` skips. +/// +/// Take `l = 4:2`, which reaches only the even offsets below `8`: +/// +/// ```text +/// l coordinate 0 1 2 3 +/// l offset 0 2 4 6 +/// ``` +/// +/// `complement(l, 8)` returns `2:1`. Its coordinate chooses an extra offset of +/// zero or one. Joining the two layouts means adding their contributions: +/// +/// ```text +/// combined coordinate (0,0) (1,0) (2,0) (3,0) | (0,1) (1,1) (2,1) (3,1) +/// l contribution 0 2 4 6 | 0 2 4 6 +/// extra contribution 0 0 0 0 | 1 1 1 1 +/// combined offset 0 2 4 6 | 1 3 5 7 +/// ``` +/// +/// Together they reach every offset from `0` through `7` exactly once. The +/// complement is therefore an *extra choice to add*, not a list of all unused +/// offsets by itself. +/// +/// `cosize_hi` is one past the last requested offset: a value of `8` requests +/// offsets `0` through `7`. It need not equal [`Layout::cosize`] of `l`. +/// [`logical_divide`](super::tiling::logical_divide) uses the returned [modes](crate#modes) for its “which +/// tile?” coordinate. +/// +/// When the requested endpoint does not fit a whole repetition, the result may +/// cover a rounded-up superset. For example, `complement(4:1, 6)` is `2:4`; +/// joining them as `(4,2):(1,4)` reaches offsets `0` through `7`, although only +/// `0` through `5` were requested. Later predication must disable `6` and `7`. +pub fn complement(l: &Layout, cosize_hi: i64) -> Result { + let f = filter(l); + let modes = flat_modes(&f); + // A broadcast or one-cell layout contributes no varying offset. The new + // coordinate must therefore span the entire requested interval. + if modes.len() == 1 && (modes[0].1 == 0 || modes[0].0 == 1) { + return Ok(coalesce(&Layout::col_major(IntTuple::Leaf(cosize_hi)))); + } + + let mut sd = modes; + sd.sort_by_key(|&(_, d)| d); + + let mut out_shapes: Vec = Vec::new(); + let mut out_strides: Vec = vec![1]; + for &(curr_shape, curr_stride) in &sd[..sd.len() - 1] { + let prev_stride = *out_strides.last().unwrap(); + if curr_stride % prev_stride != 0 { + return Err(AlgebraError(format!( + "complement: stride {curr_stride} not divisible by {prev_stride}" + ))); + } + out_shapes.push(curr_stride / prev_stride); + out_strides.push(curr_shape * curr_stride); + } + let (last_shape, last_stride) = *sd.last().unwrap(); + let prev_stride = *out_strides.last().unwrap(); + if last_stride % prev_stride != 0 { + return Err(AlgebraError(format!( + "complement: stride {last_stride} not divisible by {prev_stride}" + ))); + } + out_shapes.push(last_stride / prev_stride); + let new_stride = last_stride * last_shape; + let rest_shape = ceil_div(cosize_hi, new_stride); + out_shapes.push(rest_shape); + // The final coordinate advances beyond everything covered so far. Its + // stride is paired with the final shape appended just above. + out_strides.push(new_stride); + + let flat: Vec<(i64, i64)> = out_shapes.into_iter().zip(out_strides).collect(); + Ok(coalesce(&layout_from_flat(flat))) +} + +/// Builds the reverse lookup for a layout: given an output value, finds a flat +/// input coordinate that produces it. +/// +/// The main use is assigning work inside a GPU tile. While planning a tile, it +/// is natural to describe who owns each cell: +/// +/// ```text +/// tile cell -> (thread, value number) +/// ``` +/// +/// A running GPU thread needs the opposite answer. It knows its thread number +/// and which of its values it is processing, and must find the corresponding +/// tile cell: +/// +/// ```text +/// (thread, value number) -> tile cell +/// ``` +/// +/// Consider two threads, each responsible for three values. `T0V2` means +/// "thread 0, value number 2": +/// +/// ```text +/// tile cell c 0 1 2 | 3 4 5 +/// owner T0V0 T0V1 T0V2 | T1V0 T1V1 T1V2 +/// ``` +/// +/// A [`Layout`] returns one integer, so the pair `(thread, value number)` is +/// temporarily encoded as an **owner ID**. Thread number changes fastest: +/// +/// ```text +/// owner ID = thread + 2 * value_number +/// +/// owner T0V0 T1V0 | T0V1 T1V1 | T0V2 T1V2 +/// owner ID 0 1 | 2 3 | 4 5 +/// ``` +/// +/// In layout notation, `thr = 2:1` describes the two threads and `val = 3:1` +/// describes their three value numbers. [`raked_product`](super::tiling::raked_product) combines them into +/// the forward ownership layout: +/// +/// ```text +/// a = raked_product(thr, val) = (3,2):(2,1) +/// +/// tile cell c 0 1 2 | 3 4 5 +/// owner T0V0 T0V1 T0V2 | T1V0 T1V1 T1V2 +/// a(c), owner ID 0 2 4 | 1 3 5 +/// ``` +/// +/// Here the integer produced by `a` is an owner ID, not a memory address. The +/// right inverse `r` reverses the lookup: +/// +/// ```text +/// r = right_inverse(a) = (2,3):(3,1) +/// +/// owner ID x 0 1 2 3 4 5 +/// r(x), tile cell 0 3 1 4 2 5 +/// a(r(x)), owner ID 0 1 2 3 4 5 +/// ``` +/// +/// For example, thread 1 processing its value number 0 has owner ID +/// `1 + 2 * 0 = 1`, and `r(1) = 3` tells it to use tile cell 3. +/// [`make_layout_tv`](super::thread_value::make_layout_tv) performs this construction in three steps: +/// +/// ```text +/// raked_product tile cell -> owner ID +/// right_inverse owner ID -> tile cell +/// with_shape (thread, value number) -> tile cell +/// ``` +/// +/// In general, the returned layout guarantees `a(r(x)) = x` for every `x` that +/// `r` covers. It is called a **right** inverse because, in `a(r(x))`, `r` is +/// written to the right of `a` and runs first. +/// +/// This is not a universal lookup for every layout. If `a` skips an output, +/// there is no coordinate to return for it. If several coordinates produce the +/// same output, `r` may choose one but cannot recover the others. The compact +/// shape-and-stride format can also represent only certain reverse patterns, +/// so the returned range may be shorter than all reachable outputs. The +/// thread/value use requires every owner ID and tile cell exactly once; callers +/// validate that condition with +/// [validate_tv_layout](crate::validation::validate_tv_layout). +/// +/// Compare [`left_inverse`], which instead starts with an input, sends it +/// through `a`, and must recover that same input. +pub fn right_inverse(a: &Layout) -> Result { + let ac = coalesce(a); + let modes = flat_modes(&ac); + if modes.len() == 1 && modes[0].0 == 1 { + return Ok(layout_from_flat(vec![(1, 0)])); + } + // Before sorting modes by the offsets they reach, remember how each mode + // contributes to the original flat coordinate. These prefix products + // become strides in the inverse map. + let mut rstrides = Vec::with_capacity(modes.len()); + let mut running = 1i64; + for &(s, _) in &modes { + rstrides.push(running); + running *= s; + } + let mut dsa: Vec<(i64, i64, i64)> = modes + .iter() + .zip(&rstrides) + .map(|(&(s, d), &r)| (d, s, r)) + .collect(); + dsa.sort(); + + let mut current = 1i64; + let mut out: Vec<(i64, i64)> = Vec::new(); + for (stride, shape, rstride) in dsa { + if stride == 0 { + continue; + } + if current != stride { + break; + } + out.push((shape, rstride)); + current = shape * stride; + } + if out.is_empty() { + out.push((1, 0)); + } + Ok(coalesce(&layout_from_flat(out))) +} + +/// Recovers the original input after that input has passed through `a`. +/// +/// This round trip starts with an input, unlike [`right_inverse`]: +/// +/// ```text +/// left-inverse trip: input --a--> offset --r--> the same input +/// ``` +/// +/// For example, `a = 3:2` stores three logical positions in every other +/// offset. The gaps do not matter because the trip only visits offsets that +/// `a` actually produces: +/// +/// ```text +/// original input i 0 1 2 +/// a(i), the offset 0 2 4 +/// r(a(i)) 0 1 2 +/// +/// left_inverse(a) = (2,3):(3,1) +/// relevant lookups: r(0) = 0, r(2) = 1, r(4) = 2 +/// ``` +/// +/// The guarantee is `r(a(i)) = i` for every input `i`. It is called a +/// **left** inverse because `r` appears on the left in `r ∘ a` and runs +/// second. Values of `r` at unused offsets such as `1` and `3` have no meaning +/// for this guarantee. +/// +/// Every input must have its own offset. If two inputs share one, the offset +/// cannot tell them apart on the way back. In other words, `a` must be +/// **one-to-one** (also called *injective*). The duplicate-offset example from +/// [`right_inverse`], `(2,2):(1,0)`, therefore has no valid left inverse. +/// This function does not exhaustively check that requirement; callers must +/// validate externally supplied layouts before relying on the result. +/// +/// The implementation fills `a`'s gaps with [`complement`], then uses +/// [`right_inverse`]. +pub fn left_inverse(a: &Layout) -> Result { + let comp = complement(a, a.cosize())?; + let padded = Layout::from_modes(vec![a.clone(), comp]); + right_inverse(&padded) +} + +/// Replaces a flat input coordinate with a tuple-shaped coordinate, preserving +/// the same flat offset walk. +/// +/// When both have six cells, as in `l = 6:1` and `shape = (2,3)`, the left +/// coordinate changes fastest: +/// +/// ```text +/// old flat coordinate 0 1 2 3 4 5 +/// new tuple coordinate (0,0) (1,0) (0,1) (1,1) (0,2) (1,2) +/// offset 0 1 2 3 4 5 +/// +/// result = (2,3):(1,2) +/// ``` +/// +/// Only the way callers name a coordinate changes; the flattened +/// coordinate-to-offset sequence does not. This is useful when a flat inverse +/// or product must expose named `(thread, value)` [modes](crate#modes). +/// Preserving the complete walk requires `shape.size() == l.size()`. This +/// function does not enforce equality: a smaller shape selects a prefix, while +/// a larger shape may extend beyond `l`'s original coordinate range. +pub fn with_shape(l: &Layout, shape: IntTuple) -> Result { + composition(l, &Layout::col_major(shape)) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mod.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mod.rs new file mode 100644 index 0000000000..71fae927c9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/mod.rs @@ -0,0 +1,147 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builds new layouts from existing ones without generating runtime index math. +//! +//! Start with the basic idea from the crate root: a [`Layout`] maps a logical +//! coordinate to a linear offset. This module changes *descriptions* of those +//! maps at compile time. It never moves memory itself. +//! +//! The main operations answer practical tiling questions: +//! +//! ```text +//! coalesce Can adjacent modes be written as one simpler mode? +//! flatten What are the modes once their grouping parentheses are removed? +//! composition What is the final map after passing through two maps? +//! complement Which offsets does a layout leave unused? +//! inverse Which coordinate reaches a requested offset? +//! divide Where am I inside a tile, and which tile am I in? +//! product How do I repeat this small layout? +//! TV layout Which thread owns each value in a tile? +//! ``` +//! +//! A [**mode**](crate#modes) is one number in a coordinate that can change +//! while the other numbers stay fixed. In the examples below, `4:2` means +//! "four coordinate values, with an offset step of two": +//! +//! ```text +//! coordinate 0 1 2 3 +//! offset 0 2 4 6 +//! ``` +//! +//! `(2,3):(3,1)` has two modes: the pairs `2:3` and `3:1`. A **flat walk** +//! numbers all tuple coordinates with the left-hand mode changing fastest: +//! +//! ```text +//! flat coordinate 0 1 2 3 4 5 +//! tuple coordinate (0,0) (1,0) (0,1) (1,1) (0,2) (1,2) +//! offset 0 3 1 4 2 5 +//! ``` +//! +//! A **tile** is a fixed-size part of a larger tensor. A **TV layout** maps +//! `(thread index, value index)` to a cell inside that tile. +//! +//! Why the compiler needs algebra rather than recovered scalar arithmetic: +//! +//! ```text +//! Rust const-generic layouts +//! | +//! v +//! these operations +//! | +//! v +//! exact per-thread addresses + proof of contiguous vector accesses +//! ``` +//! +//! Some shape-and-stride combinations cannot be represented by the supported +//! layout form; those operations return [`AlgebraError`]. A tile size that does +//! not divide evenly is different: the operation may describe a rounded-up +//! final tile, and [`crate::predication`] handles its extra edge positions. +//! Untrusted compiler input is additionally checked and bounded by +//! [`crate::validation`] before it reaches expensive enumeration. +//! +//! The implementation follows cutegen, the C++ reference implementation. +//! Exhaustive small-layout tests check the defining equations, and a pinned +//! cutegen oracle compares full address maps so that formatting or +//! implementation details cannot hide a semantic difference. + +/// Static vector-access alignment derived from layout strides. +pub mod alignment; +/// Simplification, composition, inverse maps, and reshaping. +pub mod mapping; +/// Thread/value ownership for static tiles. +pub mod thread_value; +/// Tile division and repeated-layout products. +pub mod tiling; + +pub use alignment::max_alignment; +pub use mapping::{ + coalesce, complement, composition, flatten, left_inverse, right_inverse, with_shape, +}; +pub use thread_value::make_layout_tv; +pub use tiling::{blocked_product, logical_divide, logical_product, raked_product, zipped_divide}; + +use alloc::string::String; +use alloc::vec::Vec; +use core::fmt; + +use crate::{IntTuple, Layout}; + +/// Explains why a requested layout transformation is not mathematically valid. +/// +/// For example, splitting six cells into tiles of four leaves an incomplete +/// tile. Operations that require exact division report an error instead of +/// silently constructing a wrong address map. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AlgebraError( + #[doc = "Explanation of the unsupported or inconsistent layout transformation."] pub String, +); + +impl fmt::Display for AlgebraError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "layout algebra error: {}", self.0) + } +} + +impl core::error::Error for AlgebraError {} + +/// Layout-algebra result with a descriptive transformation error. +type Result = core::result::Result; + +/// Flattens nested modes into `(shape extent, stride)` pairs from left to right. +/// +/// `((2,3),4):((12,4),1)` becomes `[(2,12), (3,4), (4,1)]`. +fn flat_modes(l: &Layout) -> Vec<(i64, i64)> { + l.shape + .flatten() + .into_iter() + .zip(l.stride.flatten()) + .collect() +} + +/// Builds the simplest tuple structure that can hold a flat mode list. +/// +/// No modes produce the neutral one-cell map `1:0`; one mode does not need +/// parentheses; two or more modes become matching shape and stride tuples. +fn layout_from_flat(modes: Vec<(i64, i64)>) -> Layout { + match modes.len() { + 0 => Layout { + shape: IntTuple::Leaf(1), + stride: IntTuple::Leaf(0), + }, + 1 => Layout { + shape: IntTuple::Leaf(modes[0].0), + stride: IntTuple::Leaf(modes[0].1), + }, + _ => Layout { + shape: IntTuple::Tuple(modes.iter().map(|(s, _)| IntTuple::Leaf(*s)).collect()), + stride: IntTuple::Tuple(modes.iter().map(|(_, d)| IntTuple::Leaf(*d)).collect()), + }, + } +} + +/// Definitional equations and known layout maps, grouped by algebra family. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/alignment.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/alignment.rs new file mode 100644 index 0000000000..4223306362 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/alignment.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Contract tests for static alignment analysis. + +use super::*; + +/// Static alignment factors must agree with the pinned cutegen examples. +#[test] +fn max_alignment_matches_cutegen_static_examples() { + let cases = [ + ("(1,8):(1,1)", 8), + ("(4,2,(8,2)):(512,1,(8,8192))", 2), + ("((1,(2))):((1,(0)))", 1), + ("(8,16):(1,4)", 4), + ("(8,16):(1,6)", 2), + ("(8,16):(1,8)", 128), + ]; + for (text, expected) in cases { + assert_eq!(max_alignment(&l(text)).unwrap(), expected, "{text}"); + } +} + +/// Negative strides must not overstate the guaranteed access alignment. +#[test] +fn max_alignment_is_conservative_for_negative_strides() { + assert_eq!(max_alignment(&l("8:-1")).unwrap(), 1); +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mapping.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mapping.rs new file mode 100644 index 0000000000..d03993b94e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mapping.rs @@ -0,0 +1,199 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Contract tests for the mapping algebra family. + +use super::*; + +/// Coalescing must preserve the flat offset sequence across singleton, nested and strided modes. +#[test] +fn coalesce_examples() { + assert_eq!(coalesce(&l("(2,4):(1,2)")).to_string(), "8:1"); + assert_eq!(coalesce(&l("(2,1,4):(1,7,2)")).to_string(), "8:1"); + assert_eq!(coalesce(&l("(2,4):(4,1)")).to_string(), "(2,4):(4,1)"); + assert_eq!(coalesce(&l("1:0")).to_string(), "1:0"); +} + +/// Flattening must keep every flat coordinate's offset and expose one mode per +/// integer of the original shape. +#[test] +fn flatten_is_definitional() { + let nested = [l("((2,3),4):((12,4),1)"), l("(2,(3,4)):(1,(2,6))")]; + for a in corpus().into_iter().chain(nested) { + let flat = flatten(&a); + assert_eq!(flat.rank(), a.shape.flatten().len(), "{a}"); + assert_eq!(flat.shape.flatten(), a.shape.flatten(), "{a}"); + assert_eq!(flat.stride.flatten(), a.stride.flatten(), "{a}"); + for coord in 0..a.size() { + let coord = IntTuple::Leaf(coord); + assert_eq!(flat.call(&coord), a.call(&coord), "{a} at {coord}"); + } + } + assert_eq!( + flatten(&l("((2,3),4):((12,4),1)")).to_string(), + "(2,3,4):(12,4,1)" + ); + assert_eq!(flatten(&l("(4):(1)")).to_string(), "4:1"); +} + +/// Checks the definition of composition at every coordinate: +/// `result(i) == a(b(i))`. +/// +/// This test uses a one-mode `b`. Its stride must line up with a boundary +/// in flattened `a`. For shape modes `[s0, s1, ...]`, those boundaries are +/// `[1, s0, s0*s1, ...]`: the number of coordinates covered before the +/// next mode changes. The end of `b` must also divide one of those +/// boundaries exactly. +/// +/// Other tests cover multi-mode composition through divide, product, and +/// TV construction. Imperfect final tiles are deliberately excluded here +/// because composition may describe their overhang; predication, rather +/// than the pointwise identity, makes those tiles safe. +#[test] +fn composition_is_definitional() { + let mut checked = 0; + for a in corpus() { + let prefix: Vec = { + let mut p = vec![1i64]; + for (s, _) in super::flat_modes(&super::coalesce(&a)) { + p.push(p.last().unwrap() * s); + } + p + }; + let total = a.size(); + for n in [1i64, 2, 3, 4, 6, 8] { + for d in [0i64, 1, 2, 3, 4, 6, 8] { + let b = Layout { + shape: IntTuple::Leaf(n), + stride: IntTuple::Leaf(d), + }; + if b.cosize() > total { + continue; + } + let admissible = d == 0 + || (prefix.contains(&d) + && prefix.iter().any(|&p| p >= n * d && p % (n * d) == 0)); + if !admissible { + continue; + } + let Ok(r) = composition(&a, &b) else { + continue; // divisibility precondition not met: fine + }; + assert_eq!(r.size(), b.size(), "size mismatch for {a} o {b}"); + for i in 0..b.size() { + let expect = a.call(&IntTuple::Leaf(b.call(&IntTuple::Leaf(i)))); + let got = r.call(&IntTuple::Leaf(i)); + assert_eq!(got, expect, "({a}) o ({b}) at {i}"); + } + checked += 1; + } + } + } + assert!( + checked > 60, + "corpus too weak: only {checked} pairs checked" + ); +} + +/// Composed layouts must match known offset maps and reject unsupported mode walks. +#[test] +fn composition_known_examples() { + // A documented nested example: select a column-like walk through a + // row-major matrix, then preserve the resulting submode structure. + assert_eq!( + composition(&l("(6,2):(8,2)"), &l("(4,3):(3,1)")) + .unwrap() + .to_string(), + "((2,2),3):((24,2),8)" + ); + // Every top-level mode on the right remains a top-level result mode. + let r = composition(&l("(8,8):(8,1)"), &l("(4,2):(2,1)")).unwrap(); + for i in 0..8 { + assert_eq!( + r.call(&IntTuple::Leaf(i)), + l("(8,8):(8,1)").call(&IntTuple::Leaf(l("(4,2):(2,1)").call(&IntTuple::Leaf(i)))) + ); + } +} + +/// Joining a layout and its complement must reach each offset in `[0, M)` +/// exactly once. +#[test] +fn complement_is_definitional() { + let cases = [ + ("4:1", 24), + ("4:2", 24), + ("6:4", 24), + ("(2,3):(3,1)", 12), + ("(2,2):(1,8)", 16), + ("(4,4):(4,1)", 32), + ]; + for (text, m) in cases { + let a = l(text); + let c = complement(&a, m).unwrap(); + let combined = Layout::from_modes(vec![a.clone(), c.clone()]); + let total = combined.size(); + assert_eq!(total, m, "complement of {text} in {m} has wrong total size"); + let mut seen: Vec = (0..total) + .map(|i| combined.call(&IntTuple::Leaf(i))) + .collect(); + seen.sort_unstable(); + let expect: Vec = (0..m).collect(); + assert_eq!( + seen, expect, + "complement of {text} in {m} is not a bijection" + ); + } + // Even offsets plus a parity coordinate cover all 24 offsets. + assert_eq!( + complement(&l("4:2"), 24).unwrap().to_string(), + "(2,3):(1,8)" + ); +} + +/// A right inverse must turn every reachable consecutive offset back into +/// a coordinate that produces that offset. +#[test] +fn right_inverse_is_definitional() { + for a in corpus() { + let r = right_inverse(&a).unwrap(); + for i in 0..r.size() { + assert_eq!( + a.call(&IntTuple::Leaf(r.call(&IntTuple::Leaf(i)))), + i, + "right_inverse of {a} fails at {i}" + ); + } + // A compact one-to-one layout has no gap, so its inverse covers + // the full logical cell count. + if a.cosize() == a.size() { + assert_eq!(r.size(), a.size(), "right_inverse of compact {a} not full"); + } + } +} + +/// A left inverse must recover every original coordinate for a layout with +/// no duplicate offsets. +#[test] +fn left_inverse_is_definitional() { + for a in corpus() { + // Keep only layouts where every coordinate reaches a distinct + // offset; duplicate offsets cannot be undone. + let mut image: Vec = (0..a.size()).map(|i| a.call(&IntTuple::Leaf(i))).collect(); + image.sort_unstable(); + image.dedup(); + if image.len() as i64 != a.size() { + continue; + } + let r = left_inverse(&a).unwrap(); + for i in 0..a.size() { + assert_eq!( + r.call(&IntTuple::Leaf(a.call(&IntTuple::Leaf(i)))), + i, + "left_inverse of {a} fails at {i}" + ); + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mod.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mod.rs new file mode 100644 index 0000000000..11c7f7472b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/mod.rs @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared fixtures for the algebra family contract tests. + +use super::*; +use alloc::string::ToString; +use alloc::vec; + +/// Parse a test layout literal and fail immediately if the fixture is malformed. +fn l(text: &str) -> Layout { + text.parse().unwrap() +} + +/// Small layouts whose complete coordinate maps are cheap to enumerate. +fn corpus() -> Vec { + [ + "4:1", + "8:1", + "4:2", + "6:4", + "(2,3):(3,1)", + "(2,3):(1,2)", + "(4,4):(4,1)", + "(4,4):(1,4)", + "(2,2,2):(4,2,1)", + "(2,2,2):(1,2,4)", + "(8,4):(4,1)", + "(2,4):(1,8)", + ] + .iter() + .map(|t| l(t)) + .collect() +} + +mod alignment; +mod mapping; +mod thread_value; +mod tiling; diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/thread_value.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/thread_value.rs new file mode 100644 index 0000000000..22f369c1e9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/thread_value.rs @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Contract tests for the thread/value algebra family. + +use super::*; + +/// Checks the documented six-thread, four-values-per-thread example and +/// then verifies that its 24 assignments cover the `4 x 6` tile exactly. +#[test] +fn make_layout_tv_matches_dsl_docstring() { + let thr = l("(2,3):(3,1)"); + let val = l("(2,2):(2,1)"); + let (tiler, tv) = make_layout_tv(&thr, &val).unwrap(); + assert_eq!(tiler, vec![4, 6]); + assert_eq!(tv.to_string(), "((3,2),(2,2)):((8,2),(4,1))"); + // No two `(thread, value)` pairs may own the same tile cell, and no + // cell may be left unowned. + let mut seen: Vec = (0..tv.size()) + .map(|i| tv.call(&IntTuple::Leaf(i))) + .collect(); + seen.sort_unstable(); + let expect: Vec = (0..24).collect(); + assert_eq!(seen, expect); +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/tiling.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/tiling.rs new file mode 100644 index 0000000000..19b0cbb34c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tests/tiling.rs @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Contract tests for the tiling algebra family. + +use super::*; + +/// Zipped division changes the coordinate structure, not the address. The +/// inside-tile and tile-number coordinates are recombined mode by mode and +/// compared with the original layout. +#[test] +fn zipped_divide_is_definitional() { + let a = l("(8,6):(6,1)"); + let r = zipped_divide(&a, &[4, 3]).unwrap(); + for t0 in 0..4 { + for t1 in 0..3 { + for r0 in 0..2 { + for r1 in 0..2 { + let coord = IntTuple::Tuple(vec![ + IntTuple::Tuple(vec![IntTuple::Leaf(t0), IntTuple::Leaf(t1)]), + IntTuple::Tuple(vec![IntTuple::Leaf(r0), IntTuple::Leaf(r1)]), + ]); + let orig = IntTuple::Tuple(vec![ + IntTuple::Leaf(t0 + 4 * r0), + IntTuple::Leaf(t1 + 3 * r1), + ]); + assert_eq!( + r.call(&coord), + a.call(&orig), + "at t=({t0},{t1}) r=({r0},{r1})" + ); + } + } + } + } +} + +/// Logical product must repeat the inner block at the tiler’s expected displacements. +#[test] +fn logical_product_repeats_the_block() { + // Six copies of a four-cell block must reach all 24 offsets exactly + // once. + let r = logical_product(&l("4:1"), &l("6:1")).unwrap(); + let mut seen: Vec = (0..r.size()).map(|i| r.call(&IntTuple::Leaf(i))).collect(); + seen.sort_unstable(); + assert_eq!(seen, (0..24).collect::>()); +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/thread_value.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/thread_value.rs new file mode 100644 index 0000000000..888d374af0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/thread_value.rs @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Thread/value ownership maps derived from thread and per-thread value layouts. + +use alloc::vec; +use alloc::vec::Vec; + +use crate::{IntTuple, Layout}; + +use super::Result; +use super::mapping::{right_inverse, with_shape}; +use super::tiling::raked_product; + +/// Builds a direct map from `(thread, value_number)` to a cell in one tile. +/// +/// `thr` says how thread numbers are arranged; `val` says how one thread's +/// value numbers are arranged. For this example their input mappings are: +/// +/// ```text +/// thr = (2,3):(3,1) val = (2,2):(2,1) +/// +/// thread grid one thread's value pattern +/// 0 1 2 0 1 +/// 3 4 5 2 3 +/// ``` +/// +/// Combining the two patterns makes a `4 x 6` tile. `TnVm` means value `m` +/// owned by thread `n`: +/// +/// ```text +/// tile column +/// 0 1 2 3 4 5 +/// +------+------+------+------+------+------+ +/// tile row 0 | T0V0 | T0V1 | T1V0 | T1V1 | T2V0 | T2V1 | +/// +------+------+------+------+------+------+ +/// 1 | T0V2 | T0V3 | T1V2 | T1V3 | T2V2 | T2V3 | +/// +------+------+------+------+------+------+ +/// 2 | T3V0 | T3V1 | T4V0 | T4V1 | T5V0 | T5V1 | +/// +------+------+------+------+------+------+ +/// 3 | T3V2 | T3V3 | T4V2 | T4V3 | T5V2 | T5V3 | +/// +------+------+------+------+------+------+ +/// ``` +/// +/// The result is `(tiler, tv)`: +/// +/// ```text +/// tiler = [4, 6] tile shape +/// tv = ((3,2),(2,2)):((8,2),(4,1)) direct mapping +/// +/// input (thread 4, value 3) -> tv offset 15 -> tile cell (row 3, column 3) +/// ``` +/// +/// Here the tile is flattened with four rows changing fastest, so cell +/// `(3,3)` has flat offset `3 + 3 * 4 = 15`. The construction first makes a +/// "tile cell -> (thread, value)" map with [`raked_product`], then reverses it +/// with [`right_inverse`] and [`with_shape`]. +/// +/// A [copy atom](crate#atoms) is not part of this function. Later validation +/// groups the returned values into copy transactions and checks their memory +/// addresses. +/// +/// This algebra assumes the thread and value patterns combine into one +/// complete, one-owner-per-cell tile; it does not prove that precondition. +/// Pass the result to +/// [validate_tv_layout](crate::validation::validate_tv_layout) before using a +/// layout decoded from external compiler input. +pub fn make_layout_tv(thr: &Layout, val: &Layout) -> Result<(Vec, Layout)> { + let layout_mn = raked_product(thr, val)?; + let tiler_mn: Vec = layout_mn.modes().iter().map(|m| m.size()).collect(); + let inv = right_inverse(&layout_mn)?; + let tv = with_shape( + &inv, + IntTuple::Tuple(vec![IntTuple::Leaf(thr.size()), IntTuple::Leaf(val.size())]), + )?; + Ok((tiler_mn, tv)) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tiling.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tiling.rs new file mode 100644 index 0000000000..33ad43a75f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/algebra/tiling.rs @@ -0,0 +1,271 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Logical division and products that partition or repeat layouts. + +use alloc::format; +use alloc::vec; +use alloc::vec::Vec; + +use crate::{IntTuple, Layout}; + +use super::mapping::{coalesce, complement, composition}; +use super::{AlgebraError, Result}; + +/// Groups a coordinate space into tiles without moving any data. +/// +/// Think of eight numbered items packed into boxes that hold four items each. +/// One number from `0` through `7` can then be replaced by two smaller numbers: +/// the item's slot inside its box and the box number. +/// +/// ```text +/// original coordinate i 0 1 2 3 | 4 5 6 7 +/// slot inside box 0 1 2 3 | 0 1 2 3 +/// box number 0 0 0 0 | 1 1 1 1 +/// new coordinate (0,0) ... (3,0) | (0,1) ... (3,1) +/// result offset 0 1 2 3 | 4 5 6 7 +/// ``` +/// +/// In layout notation, the whole strip is `a = 8:1`, one box is `b = 4:1`, +/// and `logical_divide(a, b)` returns `(4,2):(1,4)`. For example, original +/// coordinate `6` becomes `(slot 2, box 1)`, and the result still reaches +/// offset `2 + 4 * 1 = 6`. +/// +/// The first top-level [mode](crate#modes) describes a position inside the +/// `b` pattern. The second chooses a translated copy of that pattern. `b` is a +/// layout, not just a tile size, so its strides also control the walk within a +/// tile. The final offsets still come from `a`. +/// +/// If the items do not fill the last box, the result still describes a whole +/// box. For example, `logical_divide(6:1, 4:1)` also returns +/// `(4,2):(1,4)`. Its last box includes offsets `6` and `7`, so later +/// predication must disable those two extra positions. +pub fn logical_divide(a: &Layout, b: &Layout) -> Result { + let rest = complement(b, a.size())?; + composition(a, &Layout::from_modes(vec![b.clone(), rest])) +} + +/// Divides every top-level coordinate into tiles, then groups the answers into +/// "position inside the tile" and "which tile". +/// +/// For example, divide a four-row, six-column grid into `2 x 3` tiles: +/// +/// ```text +/// a = (4,6):(6,1) tile_shape = [2, 3] +/// +/// column +/// 0 1 2 | 3 4 5 +/// +---------+---------+ +/// row 0 | A A A | B B B | +/// 1 | A A A | B B B | +/// +---------+---------+ +/// 2 | C C C | D D D | +/// 3 | C C C | D D D | +/// +---------+---------+ +/// +/// A = tile (0,0) B = tile (0,1) C = tile (1,0) D = tile (1,1) +/// ``` +/// +/// Row `3` is row `1` inside tile row `1`. Column `5` is column `2` inside +/// tile column `1`. The original coordinate `(3,5)` is therefore renamed as +/// `((1,2), (1,1))`: +/// +/// ```text +/// ((inside_row, inside_column), (tile_row, tile_column)) +/// ((1, 2), (1, 1)) +/// +/// original offset = 3 * 6 + 5 = 23 +/// result offset = 23 +/// result = ((2,3),(2,2)):((6,1),(12,3)) +/// ``` +/// +/// The name **zipped** describes how the coordinate pieces are regrouped, like +/// matching the two sides of a zipper: +/// +/// ```text +/// split each axis: ((inside_row, tile_row), (inside_column, tile_column)) +/// zip the answers: ((inside_row, inside_column), (tile_row, tile_column)) +/// ``` +/// +/// This changes how a coordinate is named, not the offset it reaches. Each +/// entry in `tile_shape` matches one top-level [mode](crate#modes) of `a`, so +/// the two ranks must be equal. If a tile extent does not divide its matching +/// input extent, the final tile is rounded up and later predication must +/// disable its extra edge positions. +/// +/// This is CuTe's `zipped_divide` with the tiler given as a shape, so every +/// tile mode is contiguous (stride 1). A tiler layout with its own strides +/// is not accepted. +pub fn zipped_divide(a: &Layout, tile_shape: &[i64]) -> Result { + let modes = a.modes(); + if modes.len() != tile_shape.len() { + return Err(AlgebraError(format!( + "zipped_divide: tiler rank {} does not match layout rank {}", + tile_shape.len(), + modes.len() + ))); + } + let mut tiles = Vec::with_capacity(modes.len()); + let mut rests = Vec::with_capacity(modes.len()); + for (mode, &t) in modes.iter().zip(tile_shape) { + let divided = logical_divide(mode, &Layout::contiguous(t))?; + let mut parts = divided.modes(); + if parts.len() != 2 { + return Err(AlgebraError( + "zipped_divide: mode divide was not rank-2".into(), + )); + } + rests.push(parts.pop().unwrap()); + tiles.push(parts.pop().unwrap()); + } + Ok(Layout::from_modes(vec![ + Layout::from_modes(tiles), + Layout::from_modes(rests), + ])) +} + +/// Builds a larger layout by placing translated copies of `a` according to +/// `b`. It describes the copies; it does not copy any data. +/// +/// Think of `a = 2:1` as a tray with two positions. With `b = 3:1`, the product +/// has three trays: +/// +/// ```text +/// copy A B C +/// copy number 0 1 2 +/// positions in each copy [0 1] [0 1] [0 1] +/// result offsets [0 1] [2 3] [4 5] +/// +/// product coordinate (0,0) (1,0) | (0,1) (1,1) | (0,2) (1,2) +/// result offset 0 1 | 2 3 | 4 5 +/// result = (2,3):(1,2) +/// ``` +/// +/// The result keeps two logical groups: `(position in a, copy selected by b)`. +/// If both inputs describe rows and columns, that grouping looks like: +/// +/// ```text +/// ((position_row, position_column), (copy_row, copy_column)) +/// ``` +/// +/// This is the grouping that [`blocked_product`] and [`raked_product`] later +/// pair up axis by axis. `b` is a layout, not merely a number of copies, so it +/// can also add spacing or change their order. For example, +/// `logical_product(2:1, 3:2)` is `(2,3):(1,4)` and places the copies at +/// offsets `[0,1]`, `[4,5]`, and `[8,9]`. +pub fn logical_product(a: &Layout, b: &Layout) -> Result { + let comp = complement(a, a.size() * b.cosize())?; + let repeated = composition(&comp, b)?; + Ok(Layout::from_modes(vec![a.clone(), repeated])) +} + +/// Adds non-varying one-cell modes until two product inputs have equal rank. +/// +/// A `1:0` mode contributes only offset zero, so padding changes structure but +/// not the original mapping. +fn pad_to_rank(l: &Layout, n: usize) -> Layout { + let mut modes = l.modes(); + while modes.len() < n { + modes.push(Layout { + shape: IntTuple::Leaf(1), + stride: IntTuple::Leaf(0), + }); + } + Layout::from_modes(modes) +} + +/// Pairs matching axes of a logical product, like closing a zipper. +/// +/// [`logical_product`] first keeps all `a` coordinates separate from all copy +/// coordinates. This helper pairs the matching coordinates instead: +/// +/// ```text +/// logical product: ((a_row, a_column), (copy_row, copy_column)) +/// +/// block first: ((a_row, copy_row), (a_column, copy_column)) +/// copy first: ((copy_row, a_row), (copy_column, a_column)) +/// ``` +/// +/// `block_first` selects the second or third line. Each pair is simplified +/// afterward, so a pair may print as one mode when it is one constant-stride +/// walk. If the inputs have different ranks, missing right-hand axes act like +/// one-choice `1:0` modes. +fn zipped_product(a: &Layout, b: &Layout, block_first: bool) -> Result { + let n = a.modes().len().max(b.modes().len()); + let ap = pad_to_rank(a, n); + let bp = pad_to_rank(b, n); + let lp = logical_product(&ap, &bp)?; + let parts = lp.modes(); + let (block, reps) = (&parts[0], &parts[1]); + let block_modes = block.modes(); + let rep_modes = reps.modes(); + if block_modes.len() != n || rep_modes.len() != n { + return Err(AlgebraError("product: unexpected mode structure".into())); + } + let zipped = block_modes + .into_iter() + .zip(rep_modes) + .map(|(bm, rm)| { + let pair = if block_first { + Layout::from_modes(vec![bm, rm]) + } else { + Layout::from_modes(vec![rm, bm]) + }; + coalesce(&pair) + }) + .collect(); + Ok(Layout::from_modes(zipped)) +} + +/// Repeats `a` while keeping each copy together as a block. +/// +/// For one-axis inputs, `a = 2:1` gives two positions per copy and `b = 3:1` +/// makes three copies. Letters name the copies; digits name positions inside a +/// copy: +/// +/// ```text +/// copies A = [0 1] B = [2 3] C = [4 5] +/// blocked walk A0 A1 | B0 B1 | C0 C1 +/// result offset 0 1 | 2 3 | 4 5 +/// result = 6:1 +/// ``` +/// +/// In plain English: finish copy A before starting B, then finish B before +/// starting C. [`raked_product`] visits the same six positions in a different +/// order. +/// +/// More precisely, this pairs each [mode](crate#modes) of `a` with the matching +/// repetition mode derived from `b`, putting the `a` coordinate first. For +/// layouts with several modes, "finish one whole copy" need not describe the +/// global flat walk; the blocked order applies inside each matching mode pair. +/// Strides still decide whether successive offsets are adjacent in memory. +pub fn blocked_product(a: &Layout, b: &Layout) -> Result { + zipped_product(a, b, true) +} + +/// Repeats `a` while visiting the same position across all copies first. +/// +/// Using the same one-axis example as [`blocked_product`], imagine taking the +/// first item from copies A, B, and C before coming back for their second item: +/// +/// ```text +/// copies A = [0 1] B = [2 3] C = [4 5] +/// raked walk A0 B0 C0 | A1 B1 C1 +/// result offset 0 2 4 | 1 3 5 +/// result = (3,2):(2,1) +/// ``` +/// +/// Both products cover the same six positions. Only their coordinate grouping +/// and flat-coordinate walk differ. More precisely, this pairs each +/// [mode](crate#modes) of `a` with the matching repetition mode derived from +/// `b`, putting the repetition coordinate first. With several modes, the raked +/// order applies inside each matching pair rather than across one global list +/// of copies. +/// +/// [`make_layout_tv`](super::thread_value::make_layout_tv) uses this interleaving to combine thread and per-thread +/// value patterns without assigning the same tile cell twice. +pub fn raked_product(a: &Layout, b: &Layout) -> Result { + zipped_product(a, b, false) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/composed.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/composed.rs new file mode 100644 index 0000000000..390411598d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/composed.rs @@ -0,0 +1,594 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Keep a layout and its swizzle together as one offset calculation. +//! +//! A [`Layout`] answers a basic question: “where does this logical coordinate +//! live relative to the start of a tile?” For example, a row-major 2-by-4 +//! layout maps `(row, column)` like this: +//! +//! ```text +//! logical coordinates plain element offsets +//! (0,0) (0,1) (0,2) (0,3) 0 1 2 3 +//! (1,0) (1,1) (1,2) (1,3) 4 5 6 7 +//! ``` +//! +//! A [`Swizzle`] then permutes offset bits, usually to spread shared-memory +//! requests across hardware banks. [`ComposedLayout`] stores both operations +//! so a caller cannot accidentally use the layout but forget its swizzle: +//! +//! ```text +//! coordinate +//! │ +//! ▼ +//! outer layout ──> plain offset ──> + static offset ──> inner swizzle +//! │ +//! ▼ +//! physical tile offset +//! +//! call(c) = inner(offset + outer(c)) +//! ``` +//! +//! The names “inner” and “outer” are CuTe's. They name each map's position +//! in `inner(offset + outer(c))`, not the order of evaluation: the outer +//! layout is applied first and the inner swizzle last, as CuTeDSL's +//! `ComposedLayout` documents. The stored `offset` is a fixed displacement +//! between them. +//! +//! ## Concrete example +//! +//! Start with row and column [modes](crate#modes), then swizzle the plain +//! row-major offsets: +//! +//! ```text +//! outer layout: (2,4):(4,1) +//! row mode: extent 2, stride 4 +//! column mode: extent 4, stride 1 +//! fixed offset: 0 +//! inner swizzle: S<1,0,2> (offset bit 2 toggles bit 0) +//! +//! coordinates plain offsets final offsets +//! (0,0) (0,1) (0,2) (0,3) 0 1 2 3 -> 0 1 2 3 +//! (1,0) (1,1) (1,2) (1,3) 4 5 6 7 -> 5 4 7 6 +//! ``` +//! +//! Every result is still a **tile-relative offset**, not an absolute pointer. +//! The allocation's base pointer is added only after this calculation: +//! +//! ```text +//! absolute address = allocation base pointer + physical tile offset +//! ``` +//! +//! Keeping the pointer outside the XOR makes the same tile mapping work no +//! matter where the allocation happens to begin. +//! +//! ## Elements and bytes are different units +//! +//! An *element offset* counts typed values: for `f32`, element offset 3 means +//! the fourth `f32`. A *byte offset* counts individual bytes: that same `f32` +//! begins at byte offset `3 × 4 = 12`. +//! +//! This distinction matters because a swizzle operates on the binary bits of +//! the offset. Bit 0 means one element in an element-unit layout, but one byte +//! in a byte-unit layout. [`OffsetUnit`] records which interpretation applies. +//! Use [`ComposedLayout::to_byte_offsets`] before byte-addressed validation or +//! code generation. + +use core::fmt; + +use crate::algebra::{AlgebraError, max_alignment}; +use crate::{IntTuple, Layout, Swizzle}; + +/// What one step in every stored and returned offset represents. +/// +/// The outer layout's strides, the fixed composed offset, and the number fed +/// to the swizzle must all use this same unit. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum OffsetUnit { + /// One step means one typed value, such as one `f32`. + /// + /// The element type and its byte size are supplied later by the operation + /// that uses the layout. + Elements, + /// One step means one byte, so the result can be added to a byte pointer. + Bytes, +} + +impl fmt::Display for OffsetUnit { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + OffsetUnit::Elements => write!(f, "elements"), + OffsetUnit::Bytes => write!(f, "bytes"), + } + } +} + +/// Why a composed layout could not be built or converted to byte offsets. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ComposedLayoutError { + /// The swizzle's source and target fields overlap or fall outside the + /// supported low 32-bit swizzle window. + InvalidSwizzle, + /// Element-to-byte conversion needs a positive power-of-two byte size. + /// The payload is the rejected size. + InvalidElementBytes(#[doc = "Rejected byte size for one element."] i64), + /// Scaling the offset or a stride exceeded the `i64` offset model. + ArithmeticOverflow, +} + +impl fmt::Display for ComposedLayoutError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ComposedLayoutError::InvalidSwizzle => { + write!(f, "composed layout has an invalid CuTe swizzle") + } + ComposedLayoutError::InvalidElementBytes(bytes) => write!( + f, + "element size must be a positive power of two, got {bytes} bytes" + ), + ComposedLayoutError::ArithmeticOverflow => { + write!(f, "composed-layout unit conversion overflowed") + } + } + } +} + +impl core::error::Error for ComposedLayoutError {} + +/// A complete, position-independent mapping from coordinates to swizzled +/// tile-relative offsets. +/// +/// It stores all three parts of `inner(offset + outer(coordinate))`, plus the +/// unit in which those offsets are measured. It deliberately does not store +/// the allocation's base pointer. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ComposedLayout { + /// Offset-bit permutation applied last, after the outer map and the fixed + /// displacement. + inner: Swizzle, + /// Fixed displacement in the recorded unit, added before the swizzle. + offset: i64, + /// Coordinate-to-offset map applied first, before the inner swizzle. + outer: Layout, + /// Whether offsets and strides count elements or bytes. + unit: OffsetUnit, +} + +impl ComposedLayout { + /// Build the mapping `inner(offset + outer(coordinate))`. + /// + /// The arguments follow CuTeDSL's `make_composed_layout(inner, offset, + /// outer)`. `offset`, every stride in `outer`, and the input expected by + /// `inner` must already use `unit`. Construction rejects a swizzle whose + /// bit fields overlap or fall outside the supported low 32-bit swizzle + /// window. Higher offset bits pass through unchanged. See the module's + /// [concrete example](crate::composed#concrete-example) for all three + /// steps with numbers. + pub fn new( + inner: Swizzle, + offset: i64, + outer: Layout, + unit: OffsetUnit, + ) -> Result { + if !inner.is_valid() { + return Err(ComposedLayoutError::InvalidSwizzle); + } + Ok(ComposedLayout { + inner, + offset, + outer, + unit, + }) + } + + /// Turn a plain layout into a composition that produces the same offsets. + /// + /// The identity swizzle changes no bits and the fixed offset is zero: + /// + /// ```text + /// outer layout (2,3):(3,1), coordinate (1,2) + /// outer offset 5 -> identity swizzle -> final offset 5 + /// ``` + pub fn from_layout(outer: Layout, unit: OffsetUnit) -> Self { + ComposedLayout { + inner: Swizzle::IDENTITY, + offset: 0, + outer, + unit, + } + } + + /// Return the inner swizzle, the map applied last. + /// + /// “Inner” is CuTe's name for the map on the outside of + /// `inner(offset + outer(c))`; it does not mean an inner tile or the map + /// that runs first. + pub const fn inner(&self) -> Swizzle { + self.inner + } + + /// Return the fixed displacement added after the outer layout and before + /// the inner swizzle, measured in [`Self::unit`]. + pub const fn offset(&self) -> i64 { + self.offset + } + + /// Return the outer layout, the plain coordinate-to-offset map applied + /// first. + pub const fn outer(&self) -> &Layout { + &self.outer + } + + /// Return the swizzle portion, the inner map applied last. + /// + /// This is CuTeDSL's `get_swizzle_portion(layout)` for a `ComposedLayout`: + /// the same [`Swizzle`] as [`Self::inner`], which the compiler reads. The + /// DSL form also accepts a plain layout and answers with the identity + /// swizzle `(0, 4, 3)`; a plain layout wrapped by [`Self::from_layout`] + /// carries that same [`Swizzle::IDENTITY`] here. + pub const fn get_swizzle_portion(&self) -> Swizzle { + self.inner + } + + /// Return the non-swizzle portion, the outer layout applied first. + /// + /// This is CuTeDSL's `get_nonswizzle_portion(layout)` for a + /// `ComposedLayout`: the same [`Layout`] as [`Self::outer`], which the + /// compiler reads. The DSL form returns a plain layout unchanged; here + /// the plain layout is the outer portion of [`Self::from_layout`]. + pub const fn get_nonswizzle_portion(&self) -> &Layout { + &self.outer + } + + /// Return whether this composition's offsets count elements or bytes. + pub const fn unit(&self) -> OffsetUnit { + self.unit + } + + /// Return the number of logical coordinates accepted by the outer layout. + /// + /// Swizzling only permutes locations, so it does not change this count. + pub fn size(&self) -> i64 { + self.outer.size() + } + + /// Map one logical coordinate to a physical tile-relative offset. + /// + /// The returned number uses [`Self::unit`]. This unchecked form assumes + /// that `coord` has the right structure and lies inside the outer layout; + /// compiler input should use [`Self::checked_call`] instead. The module's + /// [concrete example](crate::composed#concrete-example) shows the complete + /// outer-layout, fixed-offset, and swizzle calculation. + pub fn call(&self, coord: &IntTuple) -> i64 { + self.call_with_outer_delta(coord, 0) + } + + /// Map an in-range coordinate without panicking or overflowing. + /// + /// Returns `None` when the coordinate has the wrong structure, lies + /// outside the layout, or either addition exceeds `i64`. Compiler passes + /// use this for layouts decoded from the compiler's intermediate + /// representation (IR), which must be treated as untrusted input. + pub fn checked_call(&self, coord: &IntTuple) -> Option { + self.checked_call_with_outer_delta(coord, 0) + } + + /// Map a position shortly after the item selected by `coord`. + /// + /// `delta` is added to the outer layout's result before swizzling and uses + /// the composition's declared unit. For a byte-unit layout, `delta = 3` + /// means “three bytes after this item's first byte.” + /// + /// For a one-thread [copy atom](crate#atoms) such as `CopyG2S<16>`, + /// validation calls this method for deltas 0 through 15 to check that the + /// swizzle keeps those 16 bytes together and in order. + pub fn call_with_outer_delta(&self, coord: &IntTuple, delta: i64) -> i64 { + self.inner + .apply(self.offset + self.outer.call(coord) + delta) + } + + /// Checked form of [`Self::call_with_outer_delta`]. + /// + /// Returns `None` for an invalid coordinate or arithmetic overflow. + pub fn checked_call_with_outer_delta(&self, coord: &IntTuple, delta: i64) -> Option { + let input = self + .offset + .checked_add(self.outer.checked_call(coord)?)? + .checked_add(delta)?; + Some(self.inner.apply(input)) + } + + /// Convert an element-counting composition into an equivalent + /// byte-counting composition. + /// + /// For a power-of-two element size `2^k`, converting an element offset to + /// bytes shifts its binary value left by `k` places. Both swizzle fields + /// must move left by the same amount, so `M` increases by `k` while `B` + /// and `S` remain unchanged: + /// + /// ```text + /// four-byte elements: k = 2 + /// + /// element S<2,2,3>: [source 6..5][gap 4][target 3..2][unchanged 1..0] + /// byte S<2,4,3>: [source 8..7][gap 6][target 5..4][unchanged 3..0] + /// + /// S<2,2,3> in elements ──> S<2,4,3> in bytes + /// ``` + /// + /// The layout strides change units too: + /// + /// ```text + /// f32 layout: (2,3):(3,1) elements -> (2,3):(12,4) bytes + /// coordinate (1,2): 5 elements -> 20 bytes + /// ``` + /// + /// The fixed offset and every outer stride are multiplied by + /// `element_bytes` as well. The resulting physical byte offset is exactly + /// the old physical element offset multiplied by `element_bytes`. + /// + /// A byte-unit composition is returned unchanged. Zero, negative, and + /// non-power-of-two element sizes are rejected because a simple bit-field + /// shift cannot represent them. Overflow and a shifted swizzle whose + /// fields leave the supported low 32-bit swizzle window are also reported. + pub fn to_byte_offsets(&self, element_bytes: i64) -> Result { + if self.unit == OffsetUnit::Bytes { + return Ok(self.clone()); + } + if element_bytes <= 0 || !(element_bytes as u64).is_power_of_two() { + return Err(ComposedLayoutError::InvalidElementBytes(element_bytes)); + } + + let added_bits = element_bytes.trailing_zeros(); + let new_base = self + .inner + .num_base + .checked_add(added_bits) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + let inner = Swizzle::try_new(self.inner.num_bits, new_base, self.inner.num_shift) + .ok_or(ComposedLayoutError::InvalidSwizzle)?; + let offset = self + .offset + .checked_mul(element_bytes) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + let stride = checked_scale_tuple(&self.outer.stride, element_bytes) + .ok_or(ComposedLayoutError::ArithmeticOverflow)?; + + Ok(ComposedLayout { + inner, + offset, + outer: Layout::new(self.outer.shape.clone(), stride), + unit: OffsetUnit::Bytes, + }) + } + + /// Return a conservative access-alignment factor supported by every part + /// of this composition, measured in [`Self::unit`]. + /// + /// Alignment means an access may safely begin only at a multiple of some + /// value. A composition cannot promise more alignment than its weakest + /// part, so the result is the greatest common divisor of: + /// + /// ```text + /// outer layout's alignment + /// fixed offset + /// swizzle's chunk-alignment factor + /// │ + /// ▼ + /// gcd(outer alignment, fixed offset, swizzle alignment) + /// ``` + /// + /// Example: an outer layout aligned to 16 bytes plus an 8-byte fixed + /// offset can promise only 8-byte alignment, even if the swizzle preserves + /// 16-byte chunks. + /// + /// A swizzle with zero XOR bits changes no offset at all, so it adds no + /// restriction regardless of its placeholder `M` field; only a real + /// swizzle caps the result at its `2^M` preserved-chunk width. This is a + /// static layout property; the real allocation pointer must still be + /// suitably aligned. + pub fn max_alignment(&self) -> Result { + let outer_alignment = max_alignment(&self.outer)?; + // gcd(x, 0) = x, so a zero offset drops out on its own. + let with_offset = gcd(outer_alignment, unsigned_to_i64(self.offset.unsigned_abs())); + if self.inner.num_bits == 0 { + return Ok(with_offset.max(1)); + } + Ok(gcd(with_offset, self.inner.max_alignment()).max(1)) + } +} + +impl fmt::Display for ComposedLayout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{} o {} o {} [{}]", + self.inner, self.offset, self.outer, self.unit + ) + } +} + +/// Multiply each nested stride or displacement by the element byte size without overflow. +fn checked_scale_tuple(tuple: &IntTuple, scale: i64) -> Option { + match tuple { + IntTuple::Leaf(value) => value.checked_mul(scale).map(IntTuple::Leaf), + IntTuple::Tuple(items) => items + .iter() + .map(|item| checked_scale_tuple(item, scale)) + .collect::>() + .map(IntTuple::Tuple), + } +} + +/// Compute the common alignment factor of two unsigned magnitudes. +fn gcd(a: i64, b: i64) -> i64 { + let mut a = a.unsigned_abs(); + let mut b = b.unsigned_abs(); + while b != 0 { + (a, b) = (b, a % b); + } + unsigned_to_i64(a) +} + +/// Convert an alignment magnitude, saturating values too large for i64. +fn unsigned_to_i64(value: u64) -> i64 { + // Valid layout values normally fit after taking their absolute value. + // `i64::MIN` is the sole exception; returning 1 makes no unsafe alignment + // promise when that malformed edge case reaches this helper. + i64::try_from(value).unwrap_or(1) +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + use alloc::string::ToString; + + /// Parse a test layout literal and fail immediately if the fixture is malformed. + fn l(text: &str) -> Layout { + text.parse().unwrap() + } + + /// Composed calls must evaluate the outer layout, fixed displacement and inner swizzle in that order. + #[test] + fn call_uses_cutegen_composition_order() { + let layout = + ComposedLayout::new(Swizzle::new(2, 0, 3), 8, l("4:1"), OffsetUnit::Bytes).unwrap(); + + for coord in 0..4 { + let coord = IntTuple::Leaf(coord); + assert_eq!( + layout.call(&coord), + layout.inner.apply(8 + layout.outer.call(&coord)) + ); + } + assert_eq!(layout.to_string(), "S<2,0,3> o 8 o 4:1 [bytes]"); + } + + /// Converting element offsets to bytes must preserve the physical coordinate map. + #[test] + fn byte_conversion_preserves_the_physical_mapping() { + let elements = ComposedLayout::new( + Swizzle::new(2, 2, 3), + 8, + l("(4,4):(1,4)"), + OffsetUnit::Elements, + ) + .unwrap(); + let bytes = elements.to_byte_offsets(4).unwrap(); + + assert_eq!(bytes.unit(), OffsetUnit::Bytes); + assert_eq!(bytes.inner(), Swizzle::new(2, 4, 3)); + assert_eq!(bytes.offset(), 32); + assert_eq!(bytes.outer().to_string(), "(4,4):(4,16)"); + for coord in 0..elements.size() { + let coord = IntTuple::Leaf(coord); + assert_eq!(bytes.call(&coord), elements.call(&coord) * 4); + } + } + + /// The CuTeDSL portion accessors must name the inner swizzle and the outer layout. + #[test] + fn portion_accessors_split_the_swizzle_from_the_outer_layout() { + let swizzle = Swizzle::new(2, 2, 3); + let layout = ComposedLayout::new(swizzle, 8, l("(4,4):(1,4)"), OffsetUnit::Bytes).unwrap(); + + assert_eq!(layout.get_swizzle_portion(), swizzle); + assert_eq!(layout.get_swizzle_portion(), layout.inner()); + assert_eq!(layout.get_nonswizzle_portion(), &l("(4,4):(1,4)")); + assert_eq!(layout.get_nonswizzle_portion(), layout.outer()); + + let plain = ComposedLayout::from_layout(l("4:1"), OffsetUnit::Elements); + assert_eq!(plain.get_swizzle_portion(), Swizzle::new(0, 4, 3)); + assert_eq!(plain.get_swizzle_portion(), Swizzle::IDENTITY); + assert_eq!(plain.get_nonswizzle_portion(), &l("4:1")); + } + + /// Byte conversion must reject zero, non-power-of-two and overflowing element sizes. + #[test] + fn byte_conversion_rejects_unrepresentable_sizes() { + let layout = ComposedLayout::from_layout(l("4:1"), OffsetUnit::Elements); + assert_eq!( + layout.to_byte_offsets(3), + Err(ComposedLayoutError::InvalidElementBytes(3)) + ); + + let high_bits = + ComposedLayout::new(Swizzle::new(1, 30, 1), 0, l("4:1"), OffsetUnit::Elements).unwrap(); + assert_eq!( + high_bits.to_byte_offsets(2), + Err(ComposedLayoutError::InvalidSwizzle) + ); + } + + /// Composed construction must reject invalid public swizzle fields rather than trusting prior validation. + #[test] + fn construction_revalidates_deserialized_swizzles() { + // Compiler code can decode a Swizzle directly from its intermediate + // representation, bypassing Swizzle::new. ComposedLayout is therefore + // a second input boundary and rejects overlapping bit fields. + let invalid = Swizzle { + num_bits: 4, + num_base: 3, + num_shift: 3, + }; + assert_eq!( + ComposedLayout::new(invalid, 0, l("4:1"), OffsetUnit::Bytes), + Err(ComposedLayoutError::InvalidSwizzle) + ); + } + + /// Checked composed evaluation must return failure when the fixed displacement overflows. + #[test] + fn checked_call_rejects_offset_overflow() { + let layout = + ComposedLayout::new(Swizzle::IDENTITY, i64::MAX, l("2:1"), OffsetUnit::Bytes).unwrap(); + assert_eq!(layout.checked_call(&IntTuple::Leaf(1)), None); + } + + /// Alignment must account for outer strides, fixed displacement and swizzle-protected bits. + #[test] + fn max_alignment_combines_every_part() { + let layout = ComposedLayout::new( + Swizzle::new(3, 4, 4), + 0, + l("(32,1):(1,32)"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!(layout.max_alignment().unwrap(), 16); + + let offset_limited = ComposedLayout::new( + Swizzle::new(3, 4, 4), + 8, + l("(32,1):(1,32)"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!(offset_limited.max_alignment().unwrap(), 8); + + let outer_limited = ComposedLayout::new( + Swizzle::new(3, 4, 4), + 0, + l("(8,16):(1,4)"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!(outer_limited.max_alignment().unwrap(), 4); + } + + /// An identity swizzle must not reduce alignment already guaranteed by the outer map. + #[test] + fn zero_bit_swizzle_adds_no_alignment_cap() { + // Swizzle::IDENTITY is S<0,4,3>: its M = 4 is a placeholder, and with + // zero XOR bits it changes no offset. The offset must still count. + let offset_only = + ComposedLayout::new(Swizzle::IDENTITY, 32, l("128:1"), OffsetUnit::Bytes).unwrap(); + assert_eq!(offset_only.max_alignment().unwrap(), 32); + + let unrestricted = + ComposedLayout::new(Swizzle::IDENTITY, 0, l("128:1"), OffsetUnit::Bytes).unwrap(); + assert_eq!(unrestricted.max_alignment().unwrap(), 128); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/layout.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/layout.rs new file mode 100644 index 0000000000..cd82d79fc9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/layout.rs @@ -0,0 +1,464 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Layout construction and checked or unchecked coordinate evaluation. + +use alloc::vec; +use alloc::vec::Vec; + +use crate::IntTuple; + +/// Converts a logical coordinate into a linear offset using `shape:stride`. +/// +/// With a tuple coordinate, multiply each coordinate number by the matching +/// stride and add the results: +/// +/// ```text +/// layout = (2,3):(3,1) +/// coord = (1,2) row 1, column 2 +/// +/// offset = 1*3 + 2*1 = 5 +/// ``` +/// +/// A single integer is also accepted for a tuple shape. It follows the +/// [flat-coordinate walk](crate#flat-coordinate-walk): unpack it by changing +/// the leftmost [mode](crate#modes) fastest, then apply the strides. +/// +/// ```text +/// layout (2,3):(3,1) +/// flat coordinate 0 1 2 3 4 5 +/// unpacked coordinate (0,0) (1,0) (0,1) (1,1) (0,2) (1,2) +/// returned offset 0 3 1 4 2 5 +/// ``` +/// +/// Notice that flat coordinate `1` returns offset `3`, not offset `1`: the +/// shape controls how the flat number is unpacked, while the strides control +/// where that coordinate lives. +/// +/// # Panics and unchecked input +/// +/// This low-level function assumes that shape and stride have the same +/// structure, that a tuple coordinate matches them, and that coordinates and +/// arithmetic are valid. It may panic when those assumptions are broken. +/// Compiler attributes from outside this crate must go through the checked +/// validators instead. [`Layout::checked_call`] is the checked evaluator. +pub fn crd2idx(coord: &IntTuple, shape: &IntTuple, stride: &IntTuple) -> i64 { + match (coord, shape, stride) { + (IntTuple::Leaf(c), IntTuple::Leaf(_s), IntTuple::Leaf(d)) => { + debug_assert!(*c < *_s, "coordinate out of bounds"); + c * d + } + (IntTuple::Tuple(cs), IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => { + debug_assert!(cs.len() == ss.len() && ss.len() == ds.len()); + cs.iter() + .zip(ss.iter().zip(ds)) + .map(|(c, (s, d))| crd2idx(c, s, d)) + .sum() + } + (IntTuple::Leaf(c), IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => { + debug_assert!(ss.len() == ds.len()); + let mut rem = *c; + let mut idx = 0; + for (i, (s, d)) in ss.iter().zip(ds).enumerate() { + let sz = s.size(); + // The final mode receives the whole remaining quotient. For + // an out-of-range flat coordinate this extends that mode + // instead of silently wrapping it back into the shape. + let c_i = if i + 1 == ss.len() { rem } else { rem % sz }; + idx += crd2idx(&IntTuple::Leaf(c_i), s, d); + rem /= sz; + } + idx + } + _ => panic!("crd2idx: coord/shape/stride shapes are incompatible"), + } +} + +/// Safe evaluator used after the compiler decodes an untrusted layout. +/// +/// Unlike [`crd2idx`], this returns `None` for an out-of-range coordinate, an +/// incompatible tuple tree, an invalid shape, or overflowing arithmetic. +fn checked_crd2idx(coord: &IntTuple, shape: &IntTuple, stride: &IntTuple) -> Option { + match (coord, shape, stride) { + (IntTuple::Leaf(c), IntTuple::Leaf(s), IntTuple::Leaf(d)) => { + (*c >= 0 && *c < *s).then_some(())?; + c.checked_mul(*d) + } + (IntTuple::Tuple(cs), IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => { + (cs.len() == ss.len() && ss.len() == ds.len()).then_some(())?; + cs.iter() + .zip(ss.iter().zip(ds)) + .try_fold(0i64, |sum, (c, (s, d))| { + sum.checked_add(checked_crd2idx(c, s, d)?) + }) + } + (IntTuple::Leaf(c), IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => { + (ss.len() == ds.len() && *c >= 0).then_some(())?; + let total = checked_tuple_size(shape)?; + (*c < total).then_some(())?; + let mut rem = *c; + let mut idx = 0i64; + for (i, (s, d)) in ss.iter().zip(ds).enumerate() { + let size = checked_tuple_size(s)?; + (size > 0).then_some(())?; + let mode_coord = if i + 1 == ss.len() { rem } else { rem % size }; + idx = idx.checked_add(checked_crd2idx(&IntTuple::Leaf(mode_coord), s, d)?)?; + rem /= size; + } + Some(idx) + } + _ => None, + } +} + +/// Multiply positive shape leaves, rejecting empty, nonpositive or overflowing shapes. +fn checked_tuple_size(tuple: &IntTuple) -> Option { + match tuple { + IntTuple::Leaf(value) => (*value > 0).then_some(*value), + IntTuple::Tuple(items) => items.iter().try_fold(1i64, |size, item| { + size.checked_mul(checked_tuple_size(item)?) + }), + } +} + +/// Maps logical coordinates to linear offsets. +/// +/// Write a layout as `shape:stride`. In a flat layout, each shape number and +/// the stride in the same position describe one [mode](crate#modes): +/// +/// ```text +/// Layout (2,3):(3,1) +/// +/// mode shape says stride says +/// row row can be 0 or 1 changing row adds 3 +/// column column can be 0, 1, or 2 changing column adds 1 +/// +/// offset(row, column) = row*3 + column*1 +/// ``` +/// +/// Shape and stride must have the same parenthesis structure. In +/// `((2,4),3):((1,2),8)`, the nested pair `(2,4):(1,2)` is one grouped +/// top-level mode. The offset is usually an element index, but it can be a byte +/// index when the caller explicitly treats the strides as bytes. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Layout { + /// How many coordinate values each [mode](crate#modes) accepts. + /// + /// In `(2,3):(3,1)`, shape `(2,3)` means two rows and three columns. + pub shape: IntTuple, + /// How much the offset changes when each matching coordinate increases by + /// one. + /// + /// In `(2,3):(3,1)`, stride `(3,1)` means “add 3 for the next row” and + /// “add 1 for the next column.” Strides describe modes; they are not + /// additional [modes](crate#modes). + pub stride: IntTuple, +} + +impl Layout { + /// Creates `shape:stride` after checking that their parentheses match. + /// + /// ```text + /// shape (2,3) + stride (3,1) -> (2,3):(3,1) valid + /// shape (2,3) + stride 1 -> (2,3):1 invalid + /// ``` + /// + /// The invalid case has two shape entries but no two-entry stride tuple. + /// This constructor checks structure only; it does not reject zero or + /// negative extents. Use [`Layout::checked_size`] or the crate validators + /// when accepting untrusted layouts. + /// + /// # Panics + /// + /// Panics when `shape` and `stride` do not have identical nesting. + pub fn new(shape: IntTuple, stride: IntTuple) -> Self { + assert!( + shape.is_congruent(&stride), + "shape and stride must be congruent" + ); + Layout { shape, stride } + } + + /// Returns the logical cell count: the product of all shape numbers. + /// + /// `(2,3):(3,1)` has `2 * 3 = 6` cells. Strides do not affect this count: + /// `4:2` has four logical cells even though they occupy offsets + /// `0, 2, 4, 6`. This unchecked form may return a non-positive value or + /// overflow for an invalid shape; [`Layout::checked_size`] rejects those + /// cases. + pub fn size(&self) -> i64 { + self.shape.size() + } + + /// Returns the positive logical cell count, or `None` for an invalid size. + /// + /// ```text + /// shape (2,3) -> Some(6) + /// shape (2,0) -> None zero extent + /// shape (i64::MAX,2) -> None product does not fit in i64 + /// ``` + /// + /// Strides are irrelevant to the logical cell count. + pub fn checked_size(&self) -> Option { + checked_tuple_size(&self.shape) + } + + /// Maps one tuple or flat coordinate to its linear offset. + /// + /// ```text + /// layout (2,3):(3,1) + /// + /// tuple coordinate (1,2) -> 1*3 + 2*1 -> offset 5 + /// flat coordinate 4 -> unpacked as (0,2) -> offset 2 + /// ``` + /// + /// Flat coordinates use the [leftmost-mode-fastest walk](crate#flat-coordinate-walk). + /// This unchecked convenience method has the same input requirements as + /// [`crd2idx`]. + pub fn call(&self, coord: &IntTuple) -> i64 { + crd2idx(coord, &self.shape, &self.stride) + } + + /// Safely maps one tuple or flat coordinate to its linear offset. + /// + /// ```text + /// layout (2,3):(3,1) + /// coordinate (1,2) -> Some(5) + /// coordinate (2,0) -> None row 2 is outside a two-row shape + /// flat coordinate 6 -> None the six cells are numbered 0 through 5 + /// ``` + /// + /// `None` also means a tuple has the wrong parenthesis structure, the + /// shape is invalid, or the computed offset does not fit in `i64`. + pub fn checked_call(&self, coord: &IntTuple) -> Option { + checked_crd2idx(coord, &self.shape, &self.stride) + } + + /// Creates a one-dimensional layout of `n` adjacent cells: `n:1`. + /// + /// ```text + /// contiguous(4) = 4:1 + /// coordinate 0 1 2 3 + /// offset 0 1 2 3 + /// ``` + /// + /// This constructor does not validate `n`; [`Layout::checked_size`] + /// rejects zero or negative extents. + pub fn contiguous(n: i64) -> Self { + Layout { + shape: IntTuple::Leaf(n), + stride: IntTuple::Leaf(1), + } + } + + /// Reports whether every flat coordinate maps to the same-numbered offset. + /// + /// ```text + /// 4:1 4:2 + /// coord 0 1 2 3 coord 0 1 2 3 + /// offset 0 1 2 3 offset 0 2 4 6 + /// identity not identity + /// ``` + /// + /// A multi-mode layout can also be an identity map. Its flat-coordinate + /// walk must visit adjacent offsets in order: + /// + /// ```text + /// layout (2,3):(1,2) + /// flat coordinate 0 1 2 3 4 5 + /// tuple coordinate (0,0) (1,0) (0,1) (1,1) (0,2) (1,2) + /// offset 0 1 2 3 4 5 + /// ``` + /// + /// This checks the flat-coordinate mapping; it does not say that the + /// original tuple coordinate interface is one-dimensional. An identity + /// mapping describes one adjacent memory interval, so the compiler may be + /// able to replace scalar accesses with one vector access. The method + /// checks every coordinate and is intended for small compile-time tiles. + pub fn is_identity_map(&self) -> bool { + let Some(size) = self.checked_size() else { + return false; + }; + (0..size) + .all(|coordinate| self.checked_call(&IntTuple::Leaf(coordinate)) == Some(coordinate)) + } + + /// Returns the length of the offset interval needed to contain the layout. + /// + /// This is one more than the greatest reached offset, so skipped offsets + /// still count: + /// + /// ```text + /// layout 4:2 + /// coordinate 0 1 2 3 + /// reached offset 0 2 4 6 + /// needed positions 0 1 2 3 4 5 6 -> cosize 7 + /// + /// size = 4 logical cells, but cosize = 7 offset positions + /// ``` + /// + /// You may see the reached offsets called the *image* and their containing + /// interval called the *codomain* in layout-algebra material. The plain + /// rule is the one above: `cosize` is greatest reached offset plus one. + /// This calculation assumes non-negative strides; do not use it to bound + /// a negative-stride layout. + pub fn cosize(&self) -> i64 { + let n = self.size(); + if n == 0 { + return 0; + } + self.call(&IntTuple::Leaf(n - 1)) + 1 + } + + /// Builds a compact layout with the smallest stride on the leftmost + /// coordinate number. + /// + /// For a `(row,column)` shape, this is column-major: rows within one column + /// are adjacent, then the next column begins. + /// + /// ```text + /// shape (2,3), init 1 -> (2,3):(1,2) + /// + /// column + /// 0 1 2 + /// +---+---+---+ + /// row 0 | 0 | 2 | 4 | numbers are offsets + /// +---+---+---+ + /// row 1 | 1 | 3 | 5 | + /// +---+---+---+ + /// + /// flat-coordinate walk (0,0), (1,0), (0,1), (1,1), (0,2), (1,2) + /// returned offsets 0, 1, 2, 3, 4, 5 + /// ``` + /// + /// `init` is the distance between adjacent cells in that walk. With + /// `init = 4`, the result is `(2,3):(4,8)` and the flat coordinates map to + /// offsets `0, 4, 8, 12, 16, 20`. Shape extents and stride arithmetic are + /// unchecked; validate externally supplied layouts before using them. + pub fn col_major_with(shape: IntTuple, init: i64) -> Layout { + let leaves = shape.flatten(); + let mut strides = Vec::with_capacity(leaves.len()); + let mut running = init; + for s in &leaves { + strides.push(running); + running *= s; + } + let stride = shape.restructure(&mut strides.iter()); + Layout { shape, stride } + } + + /// Builds a compact column-major layout with unit spacing. + /// + /// This is [`Layout::col_major_with`] with `init = 1`. For example, shape + /// `(2,3)` becomes `(2,3):(1,2)`. + pub fn col_major(shape: IntTuple) -> Layout { + Layout::col_major_with(shape, 1) + } + + /// Splits one layout into one sublayout per top-level [mode](crate#modes). + /// + /// ```text + /// original layout (2,3):(3,1) + /// + /// top-level position 0 1 + /// shape entry 2 3 + /// matching stride 3 1 + /// returned sublayout 2:3 3:1 + /// + /// result [2:3, 3:1] + /// ``` + /// + /// Thus “neighboring modes” means adjacent shape/stride pairs inside the + /// same original layout—not two separate shapes such as `(2,3)` and + /// `(3,4)`. For a nested layout, + /// `((2,4),3):((1,2),8)` becomes `[(2,4):(1,2), 3:8]`; the parentheses + /// inside the first mode stay intact. A rank-1 layout returns a one-item + /// vector containing itself. [`Layout::get`] returns one mode on its own. + pub fn modes(&self) -> Vec { + match (&self.shape, &self.stride) { + (IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => ss + .iter() + .zip(ds) + .map(|(s, d)| Layout { + shape: s.clone(), + stride: d.clone(), + }) + .collect(), + _ => vec![self.clone()], + } + } + + /// Returns top-level [mode](crate#modes) `mode` as its own layout, as + /// CuTe's `get(layout)` and CuTeDSL's `cute.get(layout, mode=[i])`. + /// + /// ```text + /// layout ((2,4),3):((1,2),8) + /// get(0) (2,4):(1,2) + /// get(1) 3:8 + /// ``` + /// + /// The result is `modes()[mode]` without building the other modes. A + /// rank-1 layout returns itself for `mode == 0`. + /// + /// # Panics + /// + /// Panics when `mode` is not below [`Layout::rank`]. + pub fn get(&self, mode: usize) -> Layout { + match (&self.shape, &self.stride) { + (IntTuple::Tuple(ss), IntTuple::Tuple(ds)) => { + assert!( + mode < ss.len(), + "layout mode {mode} is out of range for rank {}", + ss.len() + ); + Layout { + shape: ss[mode].clone(), + stride: ds[mode].clone(), + } + } + _ => { + assert!(mode == 0, "layout mode {mode} is out of range for rank 1"); + self.clone() + } + } + } + + /// Counts the top-level [modes](crate#modes), as CuTe's `rank(layout)`. + /// + /// ```text + /// 4:1 rank 1 + /// (2,3):(3,1) rank 2 + /// ((2,4),3):((1,2),8) rank 2 nesting inside a mode does not count + /// ``` + pub fn rank(&self) -> usize { + self.shape.rank() + } + + /// Joins sublayouts as adjacent top-level [modes](crate#modes) of one layout. + /// + /// ```text + /// input sublayouts [2:1, 4:2] + /// shape entries [ 2, 4 ] -> (2,4) + /// stride entries [ 1, 2 ] -> (1,2) + /// result (2,4):(1,2) + /// ``` + /// + /// This reverses [`Layout::modes`]. One input remains unwrapped. No inputs + /// produce `1:0`: a one-cell mode has shape extent `1`, so its only valid + /// coordinate is `0`, and `0 * 0` gives offset `0`. + pub fn from_modes(modes: Vec) -> Layout { + match modes.len() { + 0 => Layout { + shape: IntTuple::Leaf(1), + stride: IntTuple::Leaf(0), + }, + 1 => modes.into_iter().next().unwrap(), + _ => Layout { + shape: IntTuple::Tuple(modes.iter().map(|m| m.shape.clone()).collect()), + stride: IntTuple::Tuple(modes.into_iter().map(|m| m.stride).collect()), + }, + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/lib.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/lib.rs new file mode 100644 index 0000000000..dab2d81290 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/lib.rs @@ -0,0 +1,204 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Describes where each logical tensor cell lives in linear memory. +//! +//! Parse `shape:stride`, then evaluate a coordinate: +//! +//! ``` +//! use cute_layout::{IntTuple, Layout}; +//! let layout: Layout = "(2,3):(3,1)".parse().unwrap(); +//! assert_eq!(layout.checked_call(&IntTuple::Leaf(1)), Some(3)); +//! ``` +//! +//! Here flat coordinate `1` selects `(row=1, column=0)`; the walk is explained below. +//! +//! # Modes +//! +//! GPU memory is a one-dimensional list of bytes, while kernels usually think +//! in coordinates such as `(row, column)`. A [`Layout`] is the bridge between +//! those two views. It pairs a **shape** with a **stride**: +//! +//! ```text +//! layout (2,3):(3,1) +//! shape stride +//! +//! offset(row, column) = row * 3 + column * 1 +//! +//! column +//! 0 1 2 +//! +----+----+----+ +//! row 0 | 0 | 1 | 2 | numbers are linear element offsets +//! +----+----+----+ +//! row 1 | 3 | 4 | 5 | +//! +----+----+----+ +//! ``` +//! +//! A coordinate identifies one logical position. In this flat 2D example it +//! has two numbers, `(row, column)`. Each number needed to identify a position +//! is called a **mode**. For this flat 2D layout, those numbers are the row and +//! column. +//! +//! Put another way, a mode is one number in the coordinate that you can change +//! while keeping the other numbers fixed: +//! +//! ```text +//! change column +//! (0,1) ------------> (0,2) +//! | +//! | change row +//! v +//! (1,1) +//! ``` +//! +//! In this flat example, each mode has one shape value and one matching stride: +//! +//! ```text +//! row mode column mode +//! shape 2 3 +//! stride 3 1 +//! ``` +//! +//! The row mode is described by shape `2` and stride `3`; the column mode is +//! described by shape `3` and stride `1`. Shape and stride entries describe +//! the modes; they are not extra modes. The row shape allows values `0` and +//! `1`, while the column shape allows values `0`, `1`, and `2`. The matching +//! stride says how far the linear offset moves when that coordinate number +//! increases by one. Here, moving to the next row adds `3` to the offset, +//! while moving to the next column adds `1`. +//! +//! Nested parentheses group coordinate parts without changing this rule. For +//! example, `((2,4),3):((1,2),8)` takes a coordinate `((a,b),c)`. It has two +//! top-level modes: +//! +//! ```text +//! mode 0: (a,b), described by (2,4):(1,2) +//! mode 1: c, described by 3:8 +//! ``` +//! +//! ## Flat-coordinate walk +//! +//! A layout can also be called with one flat coordinate instead of a tuple. +//! CuTe unpacks flat coordinates by changing the **leftmost mode first**. This +//! fixed order is called the flat-coordinate walk: +//! +//! ```text +//! shape (2,3) +//! flat coordinate 0 1 2 3 4 5 +//! unpacked coordinate (0,0) (1,0) (0,1) (1,1) (0,2) (1,2) +//! ``` +//! +//! The shape decides this walk; the strides decide the resulting offsets. For +//! `(2,3):(3,1)`, the coordinates above map to offsets `0, 3, 1, 4, 2, 5`. +//! This order is sometimes called *colexicographic*: in plain English, the +//! leftmost coordinate number changes fastest. +//! +//! Layout algebra often preserves the mapping from each **flat coordinate** +//! to its offset while changing how tuple coordinates are grouped. For +//! example, these layouts have the same flat-coordinate mapping: +//! +//! ```text +//! flat coordinate 0 1 2 3 4 5 6 7 +//! (2,4):(1,2) offset 0 1 2 3 4 5 6 7 +//! 8:1 offset 0 1 2 3 4 5 6 7 +//! ``` +//! +//! Their native tuple-shaped interfaces differ: the first names a position as +//! `(a,b)`, while the second names it with one number. For example, `(1,3)` in +//! the first layout and `7` in the second both reach offset `7`. Both layouts +//! still accept a flat coordinate through [`Layout::call`]. +//! +//! Layout offsets do not carry a unit. Most layouts count elements; layouts +//! used at a memory boundary may count bytes. Code that converts between the +//! two must state the unit explicitly. +//! +//! # Tiles and operations +//! +//! This crate also contains the operations needed to build GPU tiles: +//! +//! ```text +//! tensor layout +//! | +//! +-- divide/product --> repeated tiles +//! +-- thread/value --> cells owned by each thread +//! +-- compose --> final memory addresses +//! +-- swizzle --> shared-memory bank-friendly addresses +//! ``` +//! +//! A **tile** is a small rectangular part of a tensor handled as one unit. A +//! **thread-value (TV) layout** assigns every tile cell to a thread and to one +//! of that thread's values. The validation code checks that these pieces cover +//! the tile exactly before the compiler emits memory operations. +//! +//! # Atoms +//! +//! An **atom** is a compile-time description of one indivisible operation used +//! as a building block for a larger operation. It is not itself an intrinsic +//! and does not perform the operation. "Indivisible" means the threads and +//! values described by the atom must participate together as one unit. +//! +//! A **copy atom** is one indivisible copy transaction used as a building block +//! for a larger cooperative copy. `CopyG2S<16>` is a compile-time description +//! of one thread copying 16 bytes with `cp.async`; it does not perform the copy +//! itself. More generally, an atom can also describe a compute operation such +//! as a matrix multiply-accumulate (MMA). +//! +//! ```text +//! CopyG2S<16> marker -- compiler turns it into --> one 16-byte cp.async transaction +//! many such transactions ------------------> one cooperative tile copy +//! ``` +//! +//! The crate is `no_std` and shared across the compiler boundary: +//! +//! - `cute-rs` uses the types to describe a kernel's static layout; and +//! - the compiler evaluates the same layouts when it creates addresses. +//! +//! Keeping one implementation on both sides prevents the Rust API and emitted +//! GPU code from disagreeing. Tests compare complete mappings with a pinned +//! cutegen revision as an independent reference. +//! +//! # Module navigation +//! +//! [`mod@tuple`] owns nested integer tuples, [`layout`] constructs and evaluates +//! layouts, and [`parse`] supplies their text representation. [`algebra`] +//! groups mapping transformations, tiling, thread/value ownership and alignment. +//! [`composed`], [`swizzle`], [`predication`] and [`validation`] handle their +//! corresponding memory-mapping and safety checks. Foundational types remain +//! available at the crate root through reexports of the same definitions. + +#![no_std] + +extern crate alloc; + +pub mod algebra; +pub mod composed; +pub mod layout; +pub mod parse; +pub mod predication; +pub mod swizzle; +pub mod tuple; +pub mod validation; + +pub use composed::{ComposedLayout, ComposedLayoutError, OffsetUnit}; +pub use layout::{Layout, crd2idx}; +pub use parse::ParseLayoutError; +pub use predication::{ + AtomCoord, GmemMatrixBounds, PredicatedSource, PredicationError, PredicationResult, + SafeSourceOffset, plan_uniform_cp_async_source, +}; +pub use swizzle::Swizzle; +pub use tuple::IntTuple; +pub use validation::{ + CooperativeCopyPlan, CooperativePlanError, MAX_COOPERATIVE_THREADS, + MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE, MAX_STATIC_VALIDATION_BYTES, + MAX_STATIC_VALIDATION_ELEMENTS, ValidationError, ValidationResult, tma_phase_alignment_bytes, + validate_atom_compatibility, validate_compact_thread_layout, validate_cooperative_copy_plan, + validate_ldmatrix_source, validate_per_thread_atom_contiguity, validate_tma_encodable, + validate_tv_exact_coverage, validate_tv_layout, +}; + +/// Layout evaluation, representation and parsing contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/parse.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/parse.rs new file mode 100644 index 0000000000..8610d8cc4e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/parse.rs @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Text formatting and parsing for integer tuples and `shape:stride` layouts. + +use alloc::vec::Vec; +use core::fmt; +use core::str::FromStr; + +use crate::{IntTuple, Layout}; + +// Text form used in diagnostics and in the cutegen test oracle: +// +// shape:stride +// (8,32):(32,1) +// ((2,4),32):((128,32),1) // nested mode on the left +// +// Parentheses preserve the tuple tree; they are not merely decoration. + +impl fmt::Display for IntTuple { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + IntTuple::Leaf(v) => write!(f, "{v}"), + IntTuple::Tuple(ts) => { + write!(f, "(")?; + for (i, t) in ts.iter().enumerate() { + if i > 0 { + write!(f, ",")?; + } + write!(f, "{t}")?; + } + write!(f, ")") + } + } + } +} + +impl fmt::Display for Layout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}:{}", self.shape, self.stride) + } +} + +/// Indicates that text was not a complete, structurally valid layout or tuple. +/// +/// Examples include a missing `:`, an unclosed tuple, or shape and stride +/// trees that do not match. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParseLayoutError; + +impl fmt::Display for ParseLayoutError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "invalid CuTe layout/tuple syntax") + } +} + +impl core::error::Error for ParseLayoutError {} + +/// Parse one integer or parenthesized tuple from the remaining layout text. +fn parse_int_tuple(s: &str) -> Result<(IntTuple, &str), ParseLayoutError> { + let s = s.trim_start(); + if let Some(mut rest) = s.strip_prefix('(') { + let mut items = Vec::new(); + loop { + let (item, r) = parse_int_tuple(rest)?; + items.push(item); + let r = r.trim_start(); + if let Some(r) = r.strip_prefix(',') { + rest = r; + } else if let Some(r) = r.strip_prefix(')') { + return Ok((IntTuple::Tuple(items), r)); + } else { + return Err(ParseLayoutError); + } + } + } else { + let end = s + .char_indices() + .find(|(_, c)| !(c.is_ascii_digit() || *c == '-')) + .map(|(i, _)| i) + .unwrap_or(s.len()); + let value: i64 = s[..end].parse().map_err(|_| ParseLayoutError)?; + Ok((IntTuple::Leaf(value), &s[end..])) + } +} + +impl FromStr for IntTuple { + type Err = ParseLayoutError; + fn from_str(s: &str) -> Result { + let (t, rest) = parse_int_tuple(s)?; + if rest.trim().is_empty() { + Ok(t) + } else { + Err(ParseLayoutError) + } + } +} + +impl FromStr for Layout { + type Err = ParseLayoutError; + fn from_str(s: &str) -> Result { + let (shape, rest) = parse_int_tuple(s)?; + let rest = rest.trim_start(); + let rest = rest.strip_prefix(':').ok_or(ParseLayoutError)?; + let (stride, rest) = parse_int_tuple(rest)?; + if !rest.trim().is_empty() { + return Err(ParseLayoutError); + } + if !shape.is_congruent(&stride) { + return Err(ParseLayoutError); + } + Ok(Layout { shape, stride }) + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/predication.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/predication.rs new file mode 100644 index 0000000000..46f72851c0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/predication.rs @@ -0,0 +1,475 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Plan a safe global-memory read when a copy partly crosses a matrix edge. +//! +//! A [**copy atom**](crate#atoms) describes one indivisible, fixed-width +//! `cp.async` transaction, such as a 16-byte copy. A **predicate** is just a +//! boolean, such as `col < cols`. Branching uses that boolean to pick a +//! path; **predication** runs the same instruction on every thread and uses +//! the boolean to mask its effect: +//! +//! ```text +//! BRANCHING PREDICATION +//! +//! p = (col < cols)? p = (col < cols)? +//! │ copy 16 bytes if p, else 0 bytes +//! ┌───┴────┐ │ +//! yes no one path, all threads together +//! │ │ +//! copy skip +//! ``` +//! +//! A cooperative copy must use the right-hand form, and its predicate +//! generalizes from a boolean to the instruction's `source_bytes` operand: +//! +//! ```text +//! 16-byte atom of four f32 values +//! +//! inside a row: [ A B C D ] source_bytes = 16 +//! at row edge: [ A B | outside ] source_bytes = 8; remaining bytes become 0 +//! fully outside:[ outside ] source_bytes = 0; all 16 destination bytes become 0 +//! ``` +//! +//! `source_bytes = 0` does **not** skip the instruction. Every thread still +//! issues its cooperative copy. That copy reads zero source bytes and +//! zero-fills the complete destination atom; it is not a no-op. +//! +//! One path is required because branching around the copy fails two ways at +//! once. The op's lifecycle synchronizes inside the region a boundary branch +//! would skip, so a thread on the false path never reaches that barrier +//! (deadlock or undefined behavior). And the TV assignment gives every +//! shared-memory cell one owning thread, so a thread that skips leaves its +//! cells unwritten: a silent hole, wrong data, no error. Masking the effect +//! instead of the path fixes both. +//! +//! A row-major matrix stores each logical row left to right. The **leading +//! dimension**, also called the **pitch**, is the element distance between +//! two row starts. It can be larger than the logical column count because a +//! row may end with padding: +//! +//! ```text +//! columns = 5, leading_dim = 8 +//! row 0 offsets: [ 0 1 2 3 4 | 5 6 7 padding ] +//! row 1 offsets: [ 8 9 10 11 12 | 13 14 15 padding ] +//! ^ row starts are 8 elements apart +//! ``` +//! +//! Even a zero-byte instruction needs a source pointer operand. For an +//! out-of-range logical coordinate, the plan therefore selects `base + 0` as +//! a **safe fallback**. It never forms the imaginary out-of-range address: +//! +//! ```text +//! in bounds: base + checked(row * leading_dim + column) +//! zero read: base + 0 (aligned fallback; source_bytes = 0) +//! ``` +//! +//! This is a reference model: [`plan_uniform_cp_async_source`] evaluates +//! concrete bounds and coordinates for tests. It does not emit instructions +//! or run as a compiler pass. The exporter builds the GPU's runtime checks, +//! source offset and byte count separately in +//! `cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/copy/cooperative.rs`. + +use core::fmt; + +/// Runtime facts promised by one row-major global-memory matrix view. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct GmemMatrixBounds { + /// Number of logical rows, excluding any allocation beyond the matrix. + pub rows: u64, + /// Number of logical elements in each row, excluding row padding. + pub columns: u64, + /// Pitch: element distance from the start of one row to the next. + pub leading_dim: u64, + /// Byte alignment guaranteed for the matrix's base pointer. + pub base_alignment_bytes: u32, +} + +/// Logical `(row, column)` of the first element in one +/// [copy atom](crate#atoms). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AtomCoord { + /// Zero-based logical row. + pub row: u64, + /// Zero-based logical column where the transaction begins. + pub column: u64, +} + +/// Element offset that the compiler-generated cooperative copy may safely add +/// to the base pointer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SafeSourceOffset { + /// The [copy atom](crate#atoms) starts inside the logical matrix, so this is + /// its real offset. + InBounds { + /// Checked `row * leading_dim + column`, measured in elements. + element_offset: u64, + }, + /// The atom reads zero bytes, so use an aligned, valid pointer operand. + AlignedFallback { + /// Always zero today: `base + 0` is safe and aligned. + element_offset: u64, + }, +} + +impl SafeSourceOffset { + /// Return the element offset to add to the base pointer. + pub const fn element_offset(self) -> u64 { + match self { + Self::InBounds { element_offset } | Self::AlignedFallback { element_offset } => { + element_offset + } + } + } +} + +/// Source operands for the uniformly issued transaction described by one +/// [copy atom](crate#atoms). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PredicatedSource { + /// Number of bytes read from global memory, from zero through atom width. + /// `cp.async` zero-fills every remaining destination byte. A value of zero + /// still means “issue the instruction, read nothing, zero-fill the atom.” + pub source_bytes: u32, + /// Safe pointer offset used by that instruction. + pub source_offset: SafeSourceOffset, +} + +/// Result type returned by the predication planner. +pub type PredicationResult = core::result::Result; + +/// Reason a [copy atom](crate#atoms) cannot be planned safely. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PredicationError { + /// `cp.async` supports only 4-, 8-, or 16-byte transaction widths here. + UnsupportedAtomBytes(#[doc = "Rejected transaction width in bytes."] u32), + /// An element cannot have a zero-byte representation. + ZeroElementBytes, + /// The transaction width is not a whole number of elements. + AtomSplitsElement { + /// Requested transaction width in bytes. + atom_bytes: u32, + /// Width of one matrix element in bytes. + element_bytes: u32, + }, + /// The pitch would place the next row before the current logical row ends. + LeadingDimensionTooSmall { + /// Logical elements in one row. + columns: u64, + /// Promised distance between row starts, in elements. + leading_dim: u64, + }, + /// The base pointer is not aligned for the transaction width. + InsufficientBaseAlignment { + /// Required base alignment in bytes. + required: u32, + /// Alignment promised by the matrix view, in bytes. + guaranteed: u32, + }, + /// Computing `row * leading_dim + column` overflowed. + ElementOffsetOverflow, + /// Converting a valid element offset into bytes overflowed. + ByteOffsetOverflow { + /// Checked source position measured in elements. + element_offset: u64, + /// Width of one element in bytes. + element_bytes: u32, + }, + /// The atom begins at an address that is not transaction-aligned. + MisalignedAtomStart { + /// Source offset from the base pointer, measured in bytes. + byte_offset: u64, + /// Required transaction alignment in bytes. + atom_bytes: u32, + }, +} + +impl fmt::Display for PredicationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnsupportedAtomBytes(bytes) => { + write!(f, "cp.async atom must be 4, 8, or 16 bytes, got {bytes}") + } + Self::ZeroElementBytes => write!(f, "element width must be positive"), + Self::AtomSplitsElement { + atom_bytes, + element_bytes, + } => write!( + f, + "a {atom_bytes}-byte atom cannot be split into {element_bytes}-byte elements" + ), + Self::LeadingDimensionTooSmall { + columns, + leading_dim, + } => write!( + f, + "leading dimension {leading_dim} is smaller than the {columns} matrix columns" + ), + Self::InsufficientBaseAlignment { + required, + guaranteed, + } => write!( + f, + "a {required}-byte atom needs matching base alignment, but only {guaranteed} bytes are guaranteed" + ), + Self::ElementOffsetOverflow => { + write!(f, "row * leading_dim + column overflows u64") + } + Self::ByteOffsetOverflow { + element_offset, + element_bytes, + } => write!( + f, + "element offset {element_offset} times element width {element_bytes} overflows u64" + ), + Self::MisalignedAtomStart { + byte_offset, + atom_bytes, + } => write!( + f, + "source byte offset {byte_offset} is not aligned for a {atom_bytes}-byte atom" + ), + } + } +} + +impl core::error::Error for PredicationError {} + +/// Plan the source pointer and source size for the transaction described by one +/// [copy atom](crate#atoms). +/// +/// `atom_bytes` is the fixed destination transaction width. The returned +/// `source_bytes` is in `0..=atom_bytes` and always contains whole elements. +/// At a right edge it stops at the logical column count, never entering row +/// padding or the next row. +/// +/// For four-byte elements, a 16-byte atom, and a 12-element row pitch: +/// +/// ```text +/// atom starts at row 0, column 8; columns = 10 +/// logical row: [ 0 1 2 3 4 5 6 7 | 8 9 | padding ] +/// ^----- 8 source bytes +/// start byte: (0 * 12 + 8) * 4 = 32, which is 16-byte aligned +/// remaining 8 destination bytes are zero-filled +/// ``` +/// +/// If the starting coordinate is outside the matrix, the result uses the +/// aligned base fallback and `source_bytes = 0`. Lowering must still issue the +/// copy instruction; it reads nothing and zero-fills the entire atom. +pub fn plan_uniform_cp_async_source( + matrix: GmemMatrixBounds, + atom: AtomCoord, + atom_bytes: u32, + element_bytes: u32, +) -> PredicationResult { + if !matches!(atom_bytes, 4 | 8 | 16) { + return Err(PredicationError::UnsupportedAtomBytes(atom_bytes)); + } + if element_bytes == 0 { + return Err(PredicationError::ZeroElementBytes); + } + if !atom_bytes.is_multiple_of(element_bytes) { + return Err(PredicationError::AtomSplitsElement { + atom_bytes, + element_bytes, + }); + } + if matrix.leading_dim < matrix.columns { + return Err(PredicationError::LeadingDimensionTooSmall { + columns: matrix.columns, + leading_dim: matrix.leading_dim, + }); + } + if matrix.base_alignment_bytes == 0 || !matrix.base_alignment_bytes.is_multiple_of(atom_bytes) { + return Err(PredicationError::InsufficientBaseAlignment { + required: atom_bytes, + guaranteed: matrix.base_alignment_bytes, + }); + } + + // Check logical bounds before doing address arithmetic: + // + // outside coordinate ──► base + 0, source_bytes = 0 + // (never compute its imaginary row * pitch) + // + // Even `(u64::MAX, u64::MAX)` therefore selects the safe base pointer. + if atom.row >= matrix.rows || atom.column >= matrix.columns { + return Ok(PredicatedSource { + source_bytes: 0, + source_offset: SafeSourceOffset::AlignedFallback { element_offset: 0 }, + }); + } + + let element_offset = atom + .row + .checked_mul(matrix.leading_dim) + .and_then(|row_start| row_start.checked_add(atom.column)) + .ok_or(PredicationError::ElementOffsetOverflow)?; + let byte_offset = element_offset.checked_mul(u64::from(element_bytes)).ok_or( + PredicationError::ByteOffsetOverflow { + element_offset, + element_bytes, + }, + )?; + if !byte_offset.is_multiple_of(u64::from(atom_bytes)) { + return Err(PredicationError::MisalignedAtomStart { + byte_offset, + atom_bytes, + }); + } + + let atom_elements = u64::from(atom_bytes / element_bytes); + let available_elements = matrix.columns - atom.column; + let copied_elements = available_elements.min(atom_elements); + let source_bytes = u32::try_from(copied_elements * u64::from(element_bytes)) + .expect("source size is bounded by a 16-byte atom"); + + Ok(PredicatedSource { + source_bytes, + source_offset: SafeSourceOffset::InBounds { element_offset }, + }) +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Three-row matrix with padded row pitch and atom-aligned base for edge-copy tests. + const MATRIX: GmemMatrixBounds = GmemMatrixBounds { + rows: 3, + columns: 10, + leading_dim: 12, + base_alignment_bytes: 16, + }; + + /// An in-bounds copy atom must read its complete transaction width. + #[test] + fn full_atom_reads_the_whole_width() { + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 1, column: 4 }, 16, 4,), + Ok(PredicatedSource { + source_bytes: 16, + source_offset: SafeSourceOffset::InBounds { element_offset: 16 }, + }) + ); + } + + /// A partial edge atom must stop at logical columns and zero-fill the remaining bytes. + #[test] + fn partial_atom_stops_at_the_logical_row_edge() { + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 1, column: 8 }, 16, 4,), + Ok(PredicatedSource { + source_bytes: 8, + source_offset: SafeSourceOffset::InBounds { element_offset: 20 }, + }) + ); + } + + /// An out-of-range atom must select offset zero without computing an invalid source address. + #[test] + fn zero_atom_uses_the_aligned_base_fallback() { + for atom in [ + AtomCoord { row: 1, column: 12 }, + AtomCoord { + row: u64::MAX, + column: u64::MAX, + }, + ] { + assert_eq!( + plan_uniform_cp_async_source(MATRIX, atom, 16, 4), + Ok(PredicatedSource { + source_bytes: 0, + source_offset: SafeSourceOffset::AlignedFallback { element_offset: 0 }, + }) + ); + } + } + + /// In-bounds coordinates must still reject element-offset and byte-offset overflow. + #[test] + fn valid_offset_arithmetic_is_checked() { + let huge_pitch = GmemMatrixBounds { + rows: 3, + columns: 2, + leading_dim: u64::MAX, + base_alignment_bytes: 4, + }; + assert_eq!( + plan_uniform_cp_async_source(huge_pitch, AtomCoord { row: 1, column: 1 }, 4, 4,), + Err(PredicationError::ElementOffsetOverflow) + ); + + let byte_overflow = GmemMatrixBounds { + rows: 2, + columns: 1, + leading_dim: u64::MAX, + base_alignment_bytes: 8, + }; + assert_eq!( + plan_uniform_cp_async_source(byte_overflow, AtomCoord { row: 1, column: 0 }, 8, 8,), + Err(PredicationError::ByteOffsetOverflow { + element_offset: u64::MAX, + element_bytes: 8, + }) + ); + } + + /// Invalid transaction widths, element sizes and undersized row pitches must return explicit errors. + #[test] + fn malformed_widths_and_matrix_pitch_are_loud_errors() { + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 0, column: 0 }, 12, 4), + Err(PredicationError::UnsupportedAtomBytes(12)) + ); + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 0, column: 0 }, 16, 0), + Err(PredicationError::ZeroElementBytes) + ); + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 0, column: 0 }, 16, 3), + Err(PredicationError::AtomSplitsElement { + atom_bytes: 16, + element_bytes: 3, + }) + ); + + let bad_pitch = GmemMatrixBounds { + leading_dim: 9, + ..MATRIX + }; + assert_eq!( + plan_uniform_cp_async_source(bad_pitch, AtomCoord { row: 0, column: 0 }, 16, 4), + Err(PredicationError::LeadingDimensionTooSmall { + columns: 10, + leading_dim: 9, + }) + ); + } + + /// Both allocation alignment and each atom’s byte displacement must satisfy the copy width. + #[test] + fn base_and_atom_start_alignment_are_checked() { + let weak_base = GmemMatrixBounds { + base_alignment_bytes: 8, + ..MATRIX + }; + assert_eq!( + plan_uniform_cp_async_source(weak_base, AtomCoord { row: 0, column: 0 }, 16, 4), + Err(PredicationError::InsufficientBaseAlignment { + required: 16, + guaranteed: 8, + }) + ); + assert_eq!( + plan_uniform_cp_async_source(MATRIX, AtomCoord { row: 0, column: 1 }, 16, 4), + Err(PredicationError::MisalignedAtomStart { + byte_offset: 4, + atom_bytes: 16, + }) + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/swizzle.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/swizzle.rs new file mode 100644 index 0000000000..5fd49feb4a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/swizzle.rs @@ -0,0 +1,409 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Reorder offsets by XOR-ing some of their bits. +//! +//! A layout first turns a logical coordinate such as `(row, column)` into a +//! **tile-relative offset**. An offset counts from the start of the tile; it +//! is not a pointer. Depending on the layout, one step may mean one element +//! or one byte. A [`Swizzle`] changes that offset before the allocation's base +//! pointer is added: +//! +//! ```text +//! (row, column) ── layout ──> plain offset ── swizzle ──> physical offset +//! │ +//! allocation base pointer + physical offset ────────────────┘ +//! ``` +//! +//! Shared memory is split into hardware banks. A plain row-major layout can +//! make many threads ask the same bank for data at once. A swizzle permutes +//! chunks of a tile so those requests spread across banks. It changes where +//! values are stored; it does not copy data, fill memory, or change the +//! number of values. +//! +//! ## Reading `S` +//! +//! The offset is treated as a binary number, with bit 0 on the right: +//! +//! ```text +//! positive S: +//! +//! higher bits lower bits +//! [ source: B bits ][ gap: S-B bits ][ target: B bits ][ unchanged: M bits ] +//! └──────────── shift right S, then XOR ─────▲ +//! ``` +//! +//! - `B` is the width of the two bit fields. It is not necessarily the +//! number of row bits; which bits describe rows depends on the preceding +//! layout. +//! - `M` is the number of lowest bits that the swizzle never changes. +//! - `S` is the signed distance from the target field to the source field. +//! Positive `S` reads a higher field and moves it right. Negative `S` reads +//! a lower field and moves it left, reversing the source and target positions +//! in the diagram. +//! +//! The shifted source is XOR-ed with the existing target. XOR toggles target +//! bits where the source has a 1; it does not replace the target field. +//! +//! For example, `S<2,0,3>` copies bits 4..3 down to bits 1..0 and XORs them: +//! +//! ```text +//! input 0b0_01_0_00 = 8 +//! 01 ── move right by 3 ──> 01 +//! output 0b0_01_0_01 = 9 +//! ``` +//! +//! The source and target fields must not overlap, which is why `|S| >= B`. +//! Both fields must also fit in the supported low 32-bit swizzle window. +//! Higher offset bits pass through unchanged. +//! +//! ## What the protected bits guarantee +//! +//! If offsets count bytes, `M = 4` leaves the lowest four address-offset bits +//! unchanged. Each **aligned** 16-byte chunk therefore remains one aligned, +//! contiguous 16-byte chunk after swizzling, although that whole chunk may +//! move elsewhere: +//! +//! ```text +//! before: chunk number | byte within chunk +//! bits 3..0 +//! after: new chunk | same bits 3..0 +//! ``` +//! +//! In general, a swizzle alone preserves aligned chunks of `2^M` input +//! units. If the input counts elements, that means `2^M` elements, not bytes. +//! This is only one vectorization condition: the outer layout, composed +//! offset, and real pointer must provide compatible alignment too. + +use core::fmt; + +/// A reversible permutation of a tile-relative offset's bits. +/// +/// `Swizzle { num_bits: B, num_base: M, num_shift: S }` represents CuTe's +/// `S`; the field names are those of CuTeDSL's `Swizzle`. It +/// changes only the offset supplied to [`Swizzle::apply`]; it does not own +/// memory or move any values by itself. See the module-level bit diagram for +/// the meaning of all three fields. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Swizzle { + /// `B`: number of bits copied from the source field and XOR-ed into the + /// target field. + pub num_bits: u32, + /// `M`: number of lowest offset bits that are never changed. + /// + /// When offsets count bytes, this preserves aligned chunks of `2^M` + /// bytes. When offsets count elements, it preserves `2^M` elements. + pub num_base: u32, + /// `S`: signed distance from target field to source field. + /// + /// A positive value moves a higher source field right. A negative value + /// moves a lower source field left. + pub num_shift: i32, +} + +impl Swizzle { + /// Highest bit range that the swizzle's fields may occupy. + /// + /// This limits the source and target field positions, not the input value's + /// width: [`Swizzle::apply`] preserves any higher bits. The allocation's + /// absolute base address is never part of the swizzled offset. + pub const BIT_WIDTH: u32 = 32; + + /// Build `S` and panic if its bit fields are impossible. + /// + /// A valid swizzle has two separate `B`-bit fields, and both fields fit in + /// the supported low 32-bit swizzle window: + /// + /// ```text + /// source and target do not overlap: |S| >= B + /// highest used bit fits: M + B + |S| <= 32 + /// ``` + pub const fn new(num_bits: u32, num_base: u32, num_shift: i32) -> Self { + assert!( + Self::parameters_are_valid(num_bits, num_base, num_shift), + "invalid CuTe swizzle" + ); + Swizzle { + num_bits, + num_base, + num_shift, + } + } + + /// Build `S`, returning `None` when its fields overlap or fall + /// outside the supported low 32-bit swizzle window. + /// + /// Compiler passes use this form for values decoded from Rust types or the + /// compiler's intermediate representation (IR), because malformed input + /// must produce a diagnostic rather than panic the compiler. + pub const fn try_new(num_bits: u32, num_base: u32, num_shift: i32) -> Option { + if Self::parameters_are_valid(num_bits, num_base, num_shift) { + Some(Swizzle { + num_bits, + num_base, + num_shift, + }) + } else { + None + } + } + + /// Return whether `S` has disjoint fields inside the supported low + /// 32-bit swizzle window. + /// + /// `num_bits`, `num_base`, and `num_shift` mean `B`, `M`, and `S`, + /// respectively. + pub const fn parameters_are_valid(num_bits: u32, num_base: u32, num_shift: i32) -> bool { + let distance = num_shift.unsigned_abs(); + if distance < num_bits { + return false; + } + let Some(base_and_bits) = num_base.checked_add(num_bits) else { + return false; + }; + let Some(used_bits) = base_and_bits.checked_add(distance) else { + return false; + }; + used_bits <= Self::BIT_WIDTH + } + + /// Return whether this swizzle has disjoint fields inside the supported + /// low 32-bit swizzle window. + /// + /// The fields are public so compiler code can decode them directly. That + /// also means a caller can create an invalid struct literal, so an input + /// boundary must call this method before using such a value. + pub const fn is_valid(&self) -> bool { + Self::parameters_are_valid(self.num_bits, self.num_base, self.num_shift) + } + + /// A swizzle that leaves every offset unchanged. + /// + /// `B = 0` makes both bit fields empty. The `M` and `S` values are then + /// harmless placeholders chosen to match CuTe's canonical spelling. + pub const IDENTITY: Swizzle = Swizzle::new(0, 4, 3); + + /// Return a mask with 1s at every source-bit position. + /// + /// A mask selects bits with bitwise AND. For `S<2,0,3>`, the positive + /// shift places the source in bits 4..3: + /// + /// ```text + /// y_mask = 0b11000 + /// offset & y_mask keeps only those two source bits + /// ``` + pub fn y_mask(&self) -> i64 { + debug_assert!(self.is_valid()); + let mask = (1i64 << self.num_bits) - 1; + mask << (self.num_base as i64 + self.num_shift.max(0) as i64) + } + + /// Return a mask with 1s at every target-bit position. + /// + /// These are the positions that may be toggled after the source field is + /// shifted. For `S<2,0,3>`, the target is bits 1..0, so the result is + /// `0b11`. + pub fn z_mask(&self) -> i64 { + debug_assert!(self.is_valid()); + let mask = (1i64 << self.num_bits) - 1; + mask << (self.num_base as i64 + (-self.num_shift).max(0) as i64) + } + + /// Map one plain tile-relative offset to its physical offset. + /// + /// The method is unit-agnostic: if `x` counts bytes, the result counts + /// bytes; if `x` counts elements, the result counts elements. Do not add + /// an absolute pointer before calling this method: pointer-address bits + /// covered by the source mask would change the XOR result. + pub fn apply(&self, x: i64) -> i64 { + debug_assert!(self.is_valid()); + let masked = x & self.y_mask(); + let moved = if self.num_shift >= 0 { + masked >> self.num_shift + } else { + masked << (-self.num_shift) + }; + x ^ moved + } + + /// Return whether applying this swizzle twice restores the original + /// offset. + /// + /// This is called an *involution*. It holds when source and target fields + /// do not overlap: the first application XORs the source into the target, + /// and the second application XORs the same bits back out. + pub fn is_involution(&self) -> bool { + let y = self.y_mask(); + let z = if self.num_shift >= 0 { + y >> self.num_shift + } else { + y << (-self.num_shift) + }; + y & z == 0 + } + + /// Alignment factor encoded by this swizzle's unchanged low bits. + /// + /// `M` unchanged low bits keep every aligned `2^M`-unit chunk contiguous: + /// + /// ```text + /// S<3,4,3> with byte offsets + /// └── 2^4 = 16-byte chunks remain intact + /// ``` + /// + /// The answer uses the input offset's unit. It does not by itself prove a + /// legal vector access: the layout, composed offset, and allocation base + /// must satisfy the same alignment. For [`Swizzle::IDENTITY`], `M = 4` is + /// only a canonical placeholder; the identity mapping itself splits no + /// chunk of any size. + pub const fn max_alignment(&self) -> i64 { + debug_assert!(self.is_valid()); + 1i64 << self.num_base + } + + /// Translate CuTe's positive relative shift into `cuda-device`'s absolute + /// source-bit position. + /// + /// ```text + /// CuTe S<3,3,3> + /// + /// bit 8..6 bit 5..3 bit 2..0 + /// [ source ] [ target ] [ unchanged ] + /// + /// relative distance S = 3 + /// absolute first source bit = M + S = 3 + 3 = 6 + /// ``` + /// + /// CuTe stores the distance between the fields. `cuda-device` stores the + /// first source-bit number. `cuda-device` can only move a higher source + /// field down, so a negative CuTe shift has no representation and returns + /// `None`. Invalid swizzles also return `None`. + pub const fn cuda_device_absolute_source_bit(&self) -> Option { + if !self.is_valid() || self.num_shift < 0 { + return None; + } + self.num_base.checked_add(self.num_shift as u32) + } +} + +impl fmt::Display for Swizzle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "S<{},{},{}>", + self.num_bits, self.num_base, self.num_shift + ) + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Representative positive and negative swizzles must match manually calculated XOR offsets. + #[test] + fn hand_checked_values() { + // With B=2 and M=0, the target is bits 1..0. A positive S=3 + // selects source bits 4..3 and moves them down before the XOR. + let s = Swizzle::new(2, 0, 3); + assert_eq!(s.apply(0b01000), 0b01001); + assert_eq!(s.apply(0b10000), 0b10010); + // Inputs with no source bits set pass through unchanged. + assert_eq!(s.apply(0b00111), 0b00111); + + // A negative S reverses the direction: bits 1..0 are the source and + // move left into bits 4..3. + let n = Swizzle::new(2, 0, -3); + assert_eq!(n.y_mask(), 0b11); + assert_eq!(n.apply(0b00001), 0b01001); + } + + /// Swizzle field bounds and disjointness checks must match cutegen’s supported range. + #[test] + fn validity_matches_cutegen_boundaries() { + assert!(Swizzle::parameters_are_valid(3, 3, 3)); + assert!(Swizzle::parameters_are_valid(0, 32, 0)); + // Four-bit fields only three bits apart overlap. + assert!(!Swizzle::parameters_are_valid(4, 3, 3)); + // M + B + |S| = 33, one bit beyond the supported swizzle window. + assert!(!Swizzle::parameters_are_valid(3, 27, 3)); + assert!(!Swizzle::parameters_are_valid(1, 0, i32::MIN)); + assert!(Swizzle::try_new(2, 0, -3).is_some()); + assert!(Swizzle::try_new(3, 30, 3).is_none()); + } + + /// The panicking swizzle constructor must reject overlapping or unrepresentable fields. + #[test] + #[should_panic(expected = "invalid CuTe swizzle")] + fn unchecked_construction_fails_loudly() { + let _ = Swizzle::new(4, 3, 3); + } + + /// A zero-bit swizzle must leave every sampled offset unchanged. + #[test] + fn identity_is_identity() { + for x in 0..256 { + assert_eq!(Swizzle::IDENTITY.apply(x), x); + } + } + + /// Applying a valid disjoint-field XOR swizzle twice must restore the original offset. + #[test] + fn involution_on_disjoint_fields() { + // Because the source and target do not overlap, applying the same XOR + // permutation twice restores every offset. + let s = Swizzle::new(3, 3, 3); + assert!(s.is_involution()); + for x in 0..1024 { + assert_eq!(s.apply(s.apply(x)), x, "double apply must restore {x}"); + } + + // Every input in this complete bit window maps to one distinct output: + // the swizzle reorders offsets but never aliases two of them. + let mut seen: alloc::vec::Vec = (0..512).map(|x| s.apply(x)).collect(); + seen.sort_unstable(); + assert_eq!(seen, (0..512).collect::>()); + } + + /// Offset bits outside the swizzle’s source and target fields must remain unchanged. + #[test] + fn base_bits_pass_through() { + let s = Swizzle::new(3, 4, 3); + for x in 0..2048 { + assert_eq!( + s.apply(x) & 0b1111, + x & 0b1111, + "low base bits must not change" + ); + } + } + + /// Swizzle masks and protected alignment must agree with their bit-field definitions. + #[test] + fn masks_and_alignment_match_the_bit_fields() { + let s = Swizzle::new(3, 4, 3); + assert_eq!(s.y_mask(), 0b111 << 7); + assert_eq!(s.z_mask(), 0b111 << 4); + assert_eq!(s.max_alignment(), 16); + } + + /// Conversion to cuda-device’s absolute swizzle fields must reject unsupported negative shifts. + #[test] + fn cuda_device_conversion_is_explicit_and_positive_only() { + assert_eq!( + Swizzle::new(3, 3, 3).cuda_device_absolute_source_bit(), + Some(6) + ); + assert_eq!( + Swizzle::new(2, 0, 2).cuda_device_absolute_source_bit(), + Some(2) + ); + assert_eq!( + Swizzle::new(2, 0, -3).cuda_device_absolute_source_bit(), + None + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/tests.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/tests.rs new file mode 100644 index 0000000000..31dfe8553f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/tests.rs @@ -0,0 +1,121 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Foundation layout, representation and rejection contracts. + +use super::*; +use alloc::vec; + +/// Parse a test layout literal and fail immediately if the fixture is malformed. +fn t(v: &[i64]) -> IntTuple { + IntTuple::Tuple(v.iter().map(|x| IntTuple::Leaf(*x)).collect()) +} + +/// Tuple coordinates in a row-major layout must use row stride 32 and adjacent columns. +#[test] +fn row_major_8x32() { + // Eight rows of 32 adjacent values: moving one row adds 32 and moving + // one column adds 1. + let l = Layout::new(t(&[8, 32]), t(&[32, 1])); + assert_eq!(l.size(), 256); + assert_eq!(l.call(&t(&[0, 0])), 0); + assert_eq!(l.call(&t(&[1, 0])), 32); + assert_eq!(l.call(&t(&[3, 7])), 3 * 32 + 7); +} + +/// Flat integer coordinates must unpack with the leftmost shape mode changing fastest. +#[test] +fn integer_coord_is_colexicographic() { + // The leftmost mode changes first, so flat coordinate 5 becomes + // `(row 5, column 0)` and maps to `5 * 32`. + let l = Layout::new(t(&[8, 32]), t(&[32, 1])); + assert_eq!(l.call(&IntTuple::Leaf(5)), 160); + // After all eight rows, flat coordinate 9 becomes `(1,1)`. + assert_eq!(l.call(&IntTuple::Leaf(9)), 33); +} + +/// Layout formatting and parsing must preserve nested shape/stride structure. +#[test] +fn display_and_parse_roundtrip() { + use alloc::string::ToString; + for text in ["(8,32):(32,1)", "4:1", "((2,4),32):((128,32),1)"] { + let l: Layout = text.parse().unwrap(); + assert_eq!(l.to_string(), text); + } + assert!("(8,32):(32)".parse::().is_err()); // not congruent + assert!("(8,32)".parse::().is_err()); // missing stride +} + +/// Identity detection must inspect the complete flat-coordinate offset sequence. +#[test] +fn identity_map_detection() { + assert!(Layout::contiguous(4).is_identity_map()); + // Parentheses around a single mode do not change its mapping. + let l: Layout = "(4):(1)".parse().unwrap(); + assert!(l.is_identity_map()); + // This shape/stride order makes its colexicographic flat coordinates + // walk through adjacent offsets. + let rm: Layout = "(32,8):(1,32)".parse().unwrap(); + assert!(rm.is_identity_map()); + // A stride of two skips every other offset. + let s2: Layout = "4:2".parse().unwrap(); + assert!(!s2.is_identity_map()); + + let overflowing_size = Layout::new(t(&[i64::MAX, 2]), t(&[1, 1])); + assert!(!overflowing_size.is_identity_map()); + let overflowing_address = Layout::new(IntTuple::Leaf(3), IntTuple::Leaf(i64::MAX)); + assert_eq!(overflowing_address.checked_call(&IntTuple::Leaf(2)), None); +} + +/// `get` must return each top-level mode as `modes` lists it, and `rank` must count them. +#[test] +fn get_and_rank_follow_the_top_level_modes() { + use alloc::string::ToString; + for text in ["(8,32):(32,1)", "4:1", "((2,4),32):((128,32),1)"] { + let l: Layout = text.parse().unwrap(); + let modes = l.modes(); + assert_eq!(l.rank(), modes.len(), "{text}"); + for (mode, expected) in modes.into_iter().enumerate() { + assert_eq!(l.get(mode), expected, "{text} mode {mode}"); + } + } + let nested: Layout = "((2,4),32):((128,32),1)".parse().unwrap(); + assert_eq!(nested.get(0).to_string(), "(2,4):(128,32)"); + assert_eq!(nested.get(1).to_string(), "32:1"); + assert_eq!(Layout::contiguous(4).get(0), Layout::contiguous(4)); + assert_eq!(Layout::contiguous(4).rank(), 1); +} + +/// `get` past the rank must panic rather than return a wrong mode. +#[test] +#[should_panic(expected = "out of range")] +fn get_rejects_a_mode_past_the_rank() { + let l: Layout = "(8,32):(32,1)".parse().unwrap(); + let _ = l.get(2); +} + +/// Flattening a tuple must list its integers left to right and drop the grouping. +#[test] +fn flatten_and_is_congruent_follow_the_tuple_structure() { + let nested: IntTuple = "((2,3),4)".parse().unwrap(); + assert_eq!(nested.flatten(), vec![2, 3, 4]); + assert_eq!(IntTuple::Leaf(7).flatten(), vec![7]); + assert!(nested.is_congruent(&"((8,5),6)".parse().unwrap())); + assert!(!nested.is_congruent(&"(8,(5,6))".parse().unwrap())); +} + +/// Nested coordinate groups must retain their matching shape and stride structure. +#[test] +fn nested_shape() { + // The first top-level mode contains two nested coordinate parts; the + // second top-level mode is a single part. + let shape = IntTuple::Tuple(vec![t(&[2, 4]), IntTuple::Leaf(32)]); + let stride = IntTuple::Tuple(vec![t(&[128, 32]), IntTuple::Leaf(1)]); + let l = Layout::new(shape, stride); + assert_eq!(l.size(), 256); + // Every nested coordinate is multiplied by its matching nested stride. + let c = IntTuple::Tuple(vec![t(&[1, 2]), IntTuple::Leaf(5)]); + assert_eq!(l.call(&c), 128 + 64 + 5); +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/tuple.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/tuple.rs new file mode 100644 index 0000000000..f78ed570ce --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/tuple.rs @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Nested integer tuples shared by shapes, strides, and coordinates. + +use alloc::vec::Vec; + +/// One integer, or an ordered group of integers with preserved parentheses. +/// +/// Shapes, strides, and coordinates all use this type so their structures can +/// match: +/// +/// ```text +/// value printed form +/// Leaf(4) 4 +/// Tuple([Leaf(2), Leaf(3)]) (2,3) +/// Tuple([Tuple([Leaf(2), Leaf(3)]), Leaf(4)]) ((2,3),4) +/// ``` +/// +/// When used as a layout's shape or stride, each top-level item describes one +/// [mode](crate#modes). Thus `((2,3),4):((12,4),1)` has two top-level modes: +/// `(2,3):(12,4)` and `4:1`. The first happens to contain two nested +/// coordinate parts. Keeping the parentheses lets layout operations preserve +/// that grouping. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum IntTuple { + /// A single integer with no nested structure. + Leaf(#[doc = "One shape, stride or coordinate integer."] i64), + /// An ordered group of integers or further groups. + Tuple(#[doc = "Ordered child modes whose nested grouping is preserved."] Vec), +} + +impl IntTuple { + /// Multiplies all contained integers. + /// + /// When this tuple is a shape, that product is its logical cell count: + /// `(2,(3,4))` contains `2 * 3 * 4 = 24` cells. This low-level method does + /// not reject zero or negative values and does not protect against + /// overflow; [`Layout::checked_size`](crate::Layout::checked_size) does. + pub fn size(&self) -> i64 { + match self { + IntTuple::Leaf(v) => *v, + IntTuple::Tuple(ts) => ts.iter().map(IntTuple::size).product(), + } + } + + /// Reports whether two tuples have the same parenthesis structure, as + /// CuTe's `is_congruent`. + /// + /// The integer values do not matter: + /// + /// ```text + /// (2,(3,4)) and (8,(5,6)) -> congruent: both look like (_,(_,_)) + /// (2,(3,4)) and ((8,5),6) -> not congruent: the grouping differs + /// ``` + /// + /// A [`Layout`](crate::Layout) requires congruent shape and stride tuples so every shape + /// number has a matching stride number in the same position. + pub fn is_congruent(&self, other: &IntTuple) -> bool { + match (self, other) { + (IntTuple::Leaf(_), IntTuple::Leaf(_)) => true, + (IntTuple::Tuple(a), IntTuple::Tuple(b)) => { + a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.is_congruent(y)) + } + _ => false, + } + } + + /// Returns every integer from left to right and discards parentheses, as + /// CuTe's `flatten`. + /// + /// ```text + /// input: ((2,3),4) + /// output: [2, 3, 4] parentheses are gone + /// ``` + /// + /// A flat tuple holds nothing but integers, so it is returned as a + /// `Vec`. Algebra operations use this when they need the numbers but + /// not the nested [mode](crate#modes) structure; + /// [`algebra::flatten`](crate::algebra::flatten) flattens a whole layout. + pub fn flatten(&self) -> Vec { + let mut out = Vec::new(); + self.collect_leaves(&mut out); + out + } + + /// Append leaf integers left to right without retaining tuple parentheses. + fn collect_leaves(&self, out: &mut Vec) { + match self { + IntTuple::Leaf(v) => out.push(*v), + IntTuple::Tuple(ts) => ts.iter().for_each(|t| t.collect_leaves(out)), + } + } + + /// Counts top-level [modes](crate#modes); one ungrouped integer has rank 1. + /// + /// ```text + /// 4 -> [ 4 ] rank 1 + /// (2,3) -> [ 2 | 3 ] rank 2 + /// ((2,3),4) -> [ (2,3) | 4 ] rank 2 + /// ``` + /// + /// The inner `2` and `3` in the last example are nested inside its first + /// top-level mode, so they do not increase the top-level rank. + pub fn rank(&self) -> usize { + match self { + IntTuple::Leaf(_) => 1, + IntTuple::Tuple(ts) => ts.len(), + } + } + + /// Puts a flat leaf list back into this tuple's tree structure. + /// + /// This is an internal helper. It panics when the list does not contain + /// exactly enough values for the tree. + pub(crate) fn restructure(&self, leaves: &mut alloc::slice::Iter<'_, i64>) -> IntTuple { + match self { + IntTuple::Leaf(_) => IntTuple::Leaf(*leaves.next().expect("leaf count mismatch")), + IntTuple::Tuple(ts) => { + IntTuple::Tuple(ts.iter().map(|t| t.restructure(leaves)).collect()) + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/cooperative.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/cooperative.rs new file mode 100644 index 0000000000..c14c1ae9d8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/cooperative.rs @@ -0,0 +1,490 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Validate the complete global-to-shared cooperative copy plan. + +use super::thread_value::tv_coord; +use super::{ + MAX_COOPERATIVE_THREADS, MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE, MAX_STATIC_VALIDATION_BYTES, + MAX_STATIC_VALIDATION_ELEMENTS, validate_atom_compatibility, validate_compact_thread_layout, + validate_per_thread_atom_contiguity, validate_tv_exact_coverage, validate_tv_layout, +}; +use crate::algebra::make_layout_tv; +use crate::{ComposedLayout, IntTuple, Layout, OffsetUnit}; +use alloc::{collections::BTreeMap, format, string::String, vec::Vec}; +use core::fmt; + +/// Checked, reusable compile-time map for one block-wide global-to-shared copy. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CooperativeCopyPlan { + /// Tile shape with exactly two [modes](crate#modes): + /// `[row_count, column_count]`. + pub tiler_shape: Vec, + /// Map from `(thread ID, local value number)` to a logical tile cell. + pub tv_layout: Layout, + /// Canonical row-major global-memory tile map, measured in bytes. + pub gmem_byte_layout: ComposedLayout, + /// Complete ordinary-layout-plus-swizzle shared-memory map, in bytes. + pub smem_byte_layout: ComposedLayout, + /// Number of participating block threads: valid IDs are `0..thread_count`. + pub thread_count: i64, + /// Number of logical tile elements owned by each thread. + pub values_per_thread: i64, + /// Total cells in the tile: rows multiplied by columns. + pub tile_elements: i64, + /// Minimum `SmemTile::capacity`, measured in `T` elements. + /// + /// The unsafe `copy_g2s` caller must provide at least this capacity. + /// Static validation computes the requirement; it does not insert runtime + /// storage checks. The value includes fixed offsets and holes introduced + /// by the shared-memory mapping, not only `tile_elements`. + pub minimum_smem_capacity: i64, +} + +/// Explanation of why a complete cooperative copy plan is unsafe to lower. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CooperativePlanError( + #[doc = "Explanation of the static copy-plan constraint that failed."] String, +); + +impl fmt::Display for CooperativePlanError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl core::error::Error for CooperativePlanError {} + +/// Wrap a formatted static-copy diagnostic in the plan error type. +fn plan_error(message: impl Into) -> CooperativePlanError { + CooperativePlanError(message.into()) +} + +/// Require positive bounded extents before exhaustive static map validation. +fn checked_plan_layout_size(layout: &Layout, role: &str) -> Result { + let size = layout + .checked_size() + .ok_or_else(|| plan_error(format!("{role} layout size is invalid or overflows i64")))?; + if size > MAX_STATIC_VALIDATION_ELEMENTS { + return Err(plan_error(format!( + "{role} layout has {size} elements; static validation limit is \ + {MAX_STATIC_VALIDATION_ELEMENTS}" + ))); + } + Ok(size) +} + +/// Reject unsafe raw strides before CuTe algebra can multiply them. +/// +/// A mode with extent one always evaluates coordinate zero, so ordinary map +/// evaluation cannot reveal its stride: +/// +/// ```text +/// shape: 1 +/// coord: 0 +/// offset: 0 * huge_stride = 0 <- huge stride stays hidden +/// ``` +/// +/// Reading every raw stride first prevents later algebra from overflowing on +/// such a hidden value. +fn validate_plan_layout_arithmetic( + layout: &Layout, + role: &str, +) -> Result<(), CooperativePlanError> { + for stride in layout.stride.flatten() { + if stride.unsigned_abs() > MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE { + return Err(plan_error(format!( + "{role} layout stride {stride} exceeds static magnitude limit \ + {MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE}" + ))); + } + } + + let size = checked_plan_layout_size(layout, role)?; + for coordinate in 0..size { + let offset = layout + .checked_call(&IntTuple::Leaf(coordinate)) + .ok_or_else(|| { + plan_error(format!( + "{role} layout coordinate {coordinate} overflows i64" + )) + })?; + if offset.unsigned_abs() > MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE { + return Err(plan_error(format!( + "{role} layout coordinate {coordinate} maps to {offset}; static magnitude limit \ + is {MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE}" + ))); + } + } + Ok(()) +} + +/// Require the exact dense row-major map used by `GmemMatrix`. +/// +/// CuTe flattens a two-mode logical coordinate with the row mode changing +/// first. Row-major memory instead makes the column physically adjacent: +/// +/// ```text +/// 2 rows x 3 columns +/// flat cell: 0 1 | 2 3 | 4 5 +/// (row,col): (0,0) (1,0) | (0,1) (1,1) | (0,2) (1,2) +/// memory off: 0 3 | 1 4 | 2 5 +/// +/// offset = row * column_count + column +/// ``` +fn validate_canonical_row_major_tile_layout( + tile_layout: &Layout, + rows: i64, + columns: i64, +) -> Result<(), CooperativePlanError> { + let size = checked_plan_layout_size(tile_layout, "tile")?; + for cell in 0..size { + let row = cell % rows; + let column = cell / rows; + let expected = row + .checked_mul(columns) + .and_then(|offset| offset.checked_add(column)) + .ok_or_else(|| plan_error("canonical row-major tile offset overflows i64"))?; + let actual = tile_layout + .checked_call(&IntTuple::Leaf(cell)) + .ok_or_else(|| plan_error(format!("tile coordinate {cell} overflows i64")))?; + if actual != expected { + return Err(plan_error(format!( + "tile layout is not canonical row-major: logical cell {cell} must map to \ + {expected}, got {actual}" + ))); + } + } + Ok(()) +} + +/// Reject source atoms crossing a logical row or violating the promised pitch alignment. +fn validate_source_atom_rows( + tv: &Layout, + thread_count: i64, + values_per_thread: i64, + values_per_atom: i64, + tile_rows: i64, +) -> Result<(), CooperativePlanError> { + let step = usize::try_from(values_per_atom) + .map_err(|_| plan_error("copy atom value count does not fit this host"))?; + for thread in 0..thread_count { + for first_value in (0..values_per_thread).step_by(step) { + let mut atom_row = None; + let atom_end = first_value + .checked_add(values_per_atom) + .ok_or_else(|| plan_error("copy atom value range overflows i64"))?; + for value in first_value..atom_end { + let cell = tv + .checked_call(&tv_coord(thread, value)) + .ok_or_else(|| plan_error(format!("TV ({thread},{value}) overflows i64")))?; + // Decode the flat CuTe cell before checking the atom's row: + // + // cell = row + column * tile_rows + // row = cell % tile_rows + // + // Example for two rows: cells 0,2,4 are row 0; cells 1,3,5 + // are row 1. TileL applies the physical row-major stride later. + let row = cell.rem_euclid(tile_rows); + if let Some(first_row) = atom_row + && first_row != row + { + return Err(plan_error(format!( + "thread {thread}'s atom at value {first_value} crosses from row \ + {first_row} to row {row}" + ))); + } + atom_row = Some(row); + } + } + } + Ok(()) +} + +/// Prove that every shared byte has at most one owner and compute capacity. +/// +/// `minimum_smem_capacity` is one past the highest referenced byte, rounded +/// up to whole elements: +/// +/// ```text +/// referenced shared bytes: 8 .. 31 +/// required byte range: 0 .. 32 +/// f32 capacity: ceil(32 / 4) = 8 elements +/// ``` +fn validate_physical_smem_map( + tv: &Layout, + thread_count: i64, + values_per_thread: i64, + element_bytes: i64, + smem_bytes: &ComposedLayout, +) -> Result { + let mut owners = BTreeMap::new(); + let mut last_byte = -1i64; + for thread in 0..thread_count { + for value in 0..values_per_thread { + let cell = tv + .checked_call(&tv_coord(thread, value)) + .ok_or_else(|| plan_error(format!("TV ({thread},{value}) overflows i64")))?; + for byte_in_element in 0..element_bytes { + let offset = smem_bytes + .checked_call_with_outer_delta(&IntTuple::Leaf(cell), byte_in_element) + .ok_or_else(|| { + plan_error(format!( + "TV ({thread},{value}) byte {byte_in_element} overflows its smem offset" + )) + })?; + if offset < 0 { + return Err(plan_error(format!( + "TV ({thread},{value}) maps byte {byte_in_element} to negative smem offset \ + {offset}" + ))); + } + if offset >= MAX_STATIC_VALIDATION_BYTES { + return Err(plan_error(format!( + "smem byte offset {offset} exceeds static physical limit \ + {MAX_STATIC_VALIDATION_BYTES}" + ))); + } + if let Some((first_thread, first_value, first_byte)) = + owners.insert(offset, (thread, value, byte_in_element)) + { + return Err(plan_error(format!( + "smem byte {offset} has two owners: TV ({first_thread},{first_value}) byte \ + {first_byte} and TV ({thread},{value}) byte {byte_in_element}" + ))); + } + last_byte = last_byte.max(offset); + } + } + } + let required_bytes = last_byte + .checked_add(1) + .ok_or_else(|| plan_error("required smem byte range overflows i64"))?; + required_bytes + .checked_add(element_bytes - 1) + .map(|rounded| rounded / element_bytes) + .ok_or_else(|| plan_error("required smem capacity overflows i64")) +} + +/// Validate all compile-time facts for one block-wide global-to-shared copy. +/// +/// Inputs describe one fixed tile: +/// +/// - `atom_bytes`: transaction width described by the [copy atom](crate#atoms), +/// currently 4, 8, or 16 bytes; +/// - `thread_layout`: tile coordinates to compact block thread IDs; +/// - `value_layout`: per-thread local value arrangement; +/// - `tile_layout`: exactly two row/column [modes](crate#modes) in canonical +/// row-major order; +/// - `smem_layout`: destination element map, including any composed swizzle; +/// - `element_bytes`: byte width of one copied value. +/// +/// The validator constructs the TV assignment and proves: +/// +/// ```text +/// every (thread, local value) -> one in-range tile cell +/// every tile cell -> one owner +/// transaction described by each copy atom -> one aligned consecutive byte interval +/// every shared byte -> at most one writer +/// ``` +/// +/// Here is a complete four-cell example with two threads and two `f32` values +/// per thread: +/// +/// ```text +/// thread_layout = (2,1):(1,0) -> thread IDs 0 and 1 down the rows +/// value_layout = (1,2):(0,1) -> local values 0 and 1 across the columns +/// tile_layout = (2,2):(2,1) -> row-major element offsets 0 through 3 +/// smem_layout = identity composition of tile_layout +/// atom_bytes = 8 -> each thread copies two adjacent f32 values +/// element_bytes = 4 -> each value is four bytes +/// +/// TV map produced by make_layout_tv = (2,2):(1,2) +/// +/// value 0 value 1 +/// thread 0 cell 0/off 0 cell 2/off 1 bytes 0 through 7 +/// thread 1 cell 1/off 2 cell 3/off 3 bytes 8 through 15 +/// ``` +/// +/// `cell` is the flat tile-cell number; `off` is its row-major element offset. +/// Every cell has one owner, and each thread's two physical offsets are +/// adjacent, forming one aligned 8-byte transaction. +/// +/// The returned plan contains byte-addressed maps and the minimum shared +/// capacity. The unsafe caller guarantees enough shared storage, valid and +/// aligned pointers, the promised row spacing, and participating thread IDs. +/// The compiler computes valid source bytes from runtime matrix bounds and +/// zero-fills the remainder of each edge copy. +pub fn validate_cooperative_copy_plan( + atom_bytes: u32, + thread_layout: &Layout, + value_layout: &Layout, + tile_layout: &Layout, + smem_layout: &ComposedLayout, + element_bytes: i64, +) -> Result { + if !matches!(atom_bytes, 4 | 8 | 16) { + return Err(plan_error(format!( + "copy atom must be 4, 8, or 16 bytes, got {atom_bytes}" + ))); + } + if element_bytes <= 0 || !(element_bytes as u64).is_power_of_two() { + return Err(plan_error(format!( + "element size must be a positive power of two, got {element_bytes} bytes" + ))); + } + if i64::from(atom_bytes) % element_bytes != 0 { + return Err(plan_error( + "copy atom does not contain a whole number of elements", + )); + } + + let thread_count = checked_plan_layout_size(thread_layout, "thread")?; + let values_per_thread = checked_plan_layout_size(value_layout, "value")?; + let tile_elements = checked_plan_layout_size(tile_layout, "tile")?; + let smem_elements = checked_plan_layout_size(smem_layout.outer(), "shared-memory layout")?; + let assignments = thread_count + .checked_mul(values_per_thread) + .ok_or_else(|| plan_error("TV assignment count overflows i64"))?; + if assignments != tile_elements { + return Err(plan_error(format!( + "TV layout has {assignments} assignments for a {tile_elements}-element tile" + ))); + } + let tile_bytes = tile_elements + .checked_mul(element_bytes) + .ok_or_else(|| plan_error("copy tile byte count overflows i64"))?; + if tile_bytes > MAX_STATIC_VALIDATION_BYTES { + return Err(plan_error(format!( + "copy tile has {tile_bytes} bytes; static byte-map limit is \ + {MAX_STATIC_VALIDATION_BYTES}" + ))); + } + if smem_layout.offset().unsigned_abs() > MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE { + return Err(plan_error(format!( + "shared-memory composed offset {} exceeds static magnitude limit \ + {MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE}", + smem_layout.offset() + ))); + } + for (name, layout) in [ + ("thread", thread_layout), + ("value", value_layout), + ("tile", tile_layout), + ("shared-memory layout", smem_layout.outer()), + ] { + validate_plan_layout_arithmetic(layout, name)?; + } + + validate_compact_thread_layout(thread_layout) + .map_err(|error| plan_error(format!("invalid thread layout: {error}")))?; + if thread_count > MAX_COOPERATIVE_THREADS { + return Err(plan_error(format!( + "thread count {thread_count} exceeds CUDA block limit {MAX_COOPERATIVE_THREADS}" + ))); + } + + let (tiler_shape, tv_layout) = make_layout_tv(thread_layout, value_layout) + .map_err(|error| plan_error(format!("invalid TV algebra: {error}")))?; + validate_tv_layout(thread_layout, value_layout, &tiler_shape, &tv_layout) + .map_err(|error| plan_error(format!("invalid TV layout: {error}")))?; + + let tile_modes = tile_layout.modes(); + let tile_shape = tile_modes + .iter() + .map(|mode| { + mode.checked_size() + .ok_or_else(|| plan_error("tile mode size is invalid or overflows i64")) + }) + .collect::, _>>()?; + if tiler_shape.len() != 2 || tile_shape.len() != 2 { + return Err(plan_error( + "copy_g2s v0 requires exactly two row/column tile modes", + )); + } + if tile_shape != tiler_shape { + return Err(plan_error(format!( + "tile shape {tile_shape:?} does not match TV tiler {tiler_shape:?}" + ))); + } + + validate_tv_exact_coverage(&tv_layout, thread_count, values_per_thread, tile_elements) + .map_err(|error| plan_error(format!("invalid tile coverage: {error}")))?; + validate_canonical_row_major_tile_layout(tile_layout, tile_shape[0], tile_shape[1]) + .map_err(|error| plan_error(format!("invalid tile layout: {error}")))?; + if smem_elements != tile_elements { + return Err(plan_error(format!( + "shared-memory layout covers {smem_elements} elements but tile layout covers \ + {tile_elements}" + ))); + } + + let values_per_atom = i64::from(atom_bytes) / element_bytes; + validate_atom_compatibility(thread_count, values_per_thread, 1, values_per_atom) + .map_err(|error| plan_error(format!("invalid copy atom: {error}")))?; + validate_source_atom_rows( + &tv_layout, + thread_count, + values_per_thread, + values_per_atom, + tile_shape[0], + ) + .map_err(|error| plan_error(format!("global copy atom crosses a tile row: {error}")))?; + + let gmem_byte_layout = ComposedLayout::from_layout(tile_layout.clone(), OffsetUnit::Elements) + .to_byte_offsets(element_bytes) + .map_err(|error| plan_error(format!("invalid gmem byte layout: {error}")))?; + validate_per_thread_atom_contiguity( + &tv_layout, + thread_count, + values_per_thread, + i64::from(atom_bytes), + element_bytes, + &gmem_byte_layout, + i64::from(atom_bytes), + ) + .map_err(|error| { + plan_error(format!( + "global copy atom is not physically contiguous: {error}" + )) + })?; + + let smem_byte_layout = smem_layout + .to_byte_offsets(element_bytes) + .map_err(|error| plan_error(format!("invalid smem byte layout: {error}")))?; + validate_per_thread_atom_contiguity( + &tv_layout, + thread_count, + values_per_thread, + i64::from(atom_bytes), + element_bytes, + &smem_byte_layout, + i64::from(atom_bytes), + ) + .map_err(|error| { + plan_error(format!( + "shared-memory copy atom is not physically contiguous: {error}" + )) + })?; + let minimum_smem_capacity = validate_physical_smem_map( + &tv_layout, + thread_count, + values_per_thread, + element_bytes, + &smem_byte_layout, + ) + .map_err(|error| plan_error(format!("invalid physical shared-memory map: {error}")))?; + + Ok(CooperativeCopyPlan { + tiler_shape, + tv_layout, + gmem_byte_layout, + smem_byte_layout, + thread_count, + values_per_thread, + tile_elements, + minimum_smem_capacity, + }) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/error.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/error.rs new file mode 100644 index 0000000000..fcc41aed62 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/error.rs @@ -0,0 +1,266 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Errors returned by individual static layout proofs. + +use core::fmt; + +/// Failure reported by one reusable, low-level layout check. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ValidationError { + /// A count such as threads, values, or bytes was zero or negative. + InvalidCount { + /// Name of the invalid quantity. + name: &'static str, + /// Invalid value supplied by the caller. + value: i64, + }, + /// A count exceeds the safe exhaustive-enumeration limit. + CountTooLarge { + /// Name of the oversized quantity. + name: &'static str, + /// Count that would have been enumerated. + value: i64, + }, + /// A thread-layout coordinate produced an ID outside the block's range. + ThreadOutOfRange { + /// Logical coordinate evaluated in the thread layout. + coordinate: i64, + /// Thread ID produced by that coordinate. + thread: i64, + /// Number of threads, whose valid IDs are `0..thread_count`. + thread_count: i64, + }, + /// Two thread-layout coordinates produced the same thread ID. + DuplicateThread { + /// Repeated thread ID. + thread: i64, + /// First coordinate that produced the ID. + first_coordinate: i64, + /// Second coordinate that produced the ID. + second_coordinate: i64, + }, + /// A TV layout did not have exactly two [modes](crate#modes): thread and + /// local value. + TvRank { + /// Number of top-level modes found. + actual: usize, + }, + /// A TV mode did not have the expected number of entries. + TvModeSize { + /// Either `"thread"` or `"value"`. + mode: &'static str, + /// Required number of entries. + expected: i64, + /// Number of entries found. + actual: i64, + }, + /// The number of TV pairs differs from the number of tile cells. + TvCardinality { + /// `thread_count * values_per_thread`. + assignments: i64, + /// Number of cells in the tile. + tile_elements: i64, + }, + /// A TV pair mapped outside the tile. + TvOutOfRange { + /// Thread ID in the bad TV pair. + thread: i64, + /// Local value number in the bad TV pair. + value: i64, + /// Tile cell produced by the TV map. + cell: i64, + /// Number of valid tile cells. + tile_elements: i64, + }, + /// Two TV pairs claimed the same tile cell. + DuplicateCell { + /// Tile cell with two owners. + cell: i64, + /// First `(thread, local value)` owner. + first_owner: (i64, i64), + /// Second `(thread, local value)` owner. + second_owner: (i64, i64), + }, + /// No TV pair claimed one tile cell. + MissingCell { + /// Unowned tile cell. + cell: i64, + }, + /// The copy cannot be split into whole thread/value groups required by its + /// [copy atom](crate#atoms). + AtomDoesNotDivide { + /// Grouped dimension: `"thread"` or `"value"`. + axis: &'static str, + /// Total count in the cooperative copy. + copy_count: i64, + /// Count consumed by one atom. + atom_count: i64, + }, + /// A transaction width is not a whole number of elements. + AtomSize { + /// Transaction width in bytes. + atom_bytes: i64, + /// Width of one element in bytes. + element_bytes: i64, + }, + /// Physical contiguity was requested from an element-offset layout. + AtomNeedsByteLayout, + /// The memory base does not guarantee the transaction's alignment. + BaseAlignment { + /// Required byte alignment. + required: i64, + /// Alignment promised by the caller. + guaranteed: i64, + }, + /// The first byte of a transaction is not atom-aligned. + AtomStartMisaligned { + /// Thread that owns the atom. + thread: i64, + /// First local value in the atom. + value: i64, + /// Physical start offset in bytes. + offset: i64, + /// Required transaction alignment in bytes. + atom_bytes: i64, + }, + /// Bytes assigned to one atom do not form one consecutive interval. + AtomNotContiguous { + /// Thread that owns the atom. + thread: i64, + /// First local value in the atom. + value: i64, + /// Byte position within the atom where continuity breaks. + byte_in_atom: i64, + /// Consecutive byte offset required at that position. + expected: i64, + /// Byte offset produced by the composed memory map. + actual: i64, + }, + /// Checked integer arithmetic failed while validating the named fact. + ArithmeticOverflow( + #[doc = "Name of the layout calculation whose checked arithmetic overflowed."] &'static str, + ), +} + +impl fmt::Display for ValidationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ValidationError::InvalidCount { name, value } => { + write!(f, "{name} must be positive, got {value}") + } + ValidationError::CountTooLarge { name, value } => { + write!(f, "{name}={value} cannot be enumerated on this host") + } + ValidationError::ThreadOutOfRange { + coordinate, + thread, + thread_count, + } => write!( + f, + "thread layout coordinate {coordinate} maps to thread {thread}, outside 0..{thread_count}" + ), + ValidationError::DuplicateThread { + thread, + first_coordinate, + second_coordinate, + } => write!( + f, + "thread {thread} is produced twice, by coordinates {first_coordinate} and {second_coordinate}" + ), + ValidationError::TvRank { actual } => { + write!( + f, + "TV layout must have (thread,value) modes, got rank {actual}" + ) + } + ValidationError::TvModeSize { + mode, + expected, + actual, + } => write!(f, "TV {mode} mode has size {actual}, expected {expected}"), + ValidationError::TvCardinality { + assignments, + tile_elements, + } => write!( + f, + "TV layout has {assignments} assignments for a {tile_elements}-element tile" + ), + ValidationError::TvOutOfRange { + thread, + value, + cell, + tile_elements, + } => write!( + f, + "TV ({thread},{value}) maps to cell {cell}, outside 0..{tile_elements}" + ), + ValidationError::DuplicateCell { + cell, + first_owner, + second_owner, + } => write!( + f, + "tile cell {cell} has two owners: TV ({},{}) and TV ({},{})", + first_owner.0, first_owner.1, second_owner.0, second_owner.1 + ), + ValidationError::MissingCell { cell } => { + write!(f, "tile cell {cell} has no TV owner") + } + ValidationError::AtomDoesNotDivide { + axis, + copy_count, + atom_count, + } => write!( + f, + "copy {axis} count {copy_count} is not divisible by atom {axis} count {atom_count}" + ), + ValidationError::AtomSize { + atom_bytes, + element_bytes, + } => write!( + f, + "a {atom_bytes}-byte atom cannot be split into {element_bytes}-byte elements" + ), + ValidationError::AtomNeedsByteLayout => { + write!( + f, + "atom contiguity must be checked with a byte-offset layout" + ) + } + ValidationError::BaseAlignment { + required, + guaranteed, + } => write!( + f, + "atom needs {required}-byte base alignment, but only {guaranteed} bytes are guaranteed" + ), + ValidationError::AtomStartMisaligned { + thread, + value, + offset, + atom_bytes, + } => write!( + f, + "TV ({thread},{value}) starts at byte {offset}, not a {atom_bytes}-byte boundary" + ), + ValidationError::AtomNotContiguous { + thread, + value, + byte_in_atom, + expected, + actual, + } => write!( + f, + "TV atom ({thread},{value}) breaks at byte {byte_in_atom}: expected {expected}, got {actual}" + ), + ValidationError::ArithmeticOverflow(context) => { + write!(f, "integer overflow while checking {context}") + } + } + } +} + +impl core::error::Error for ValidationError {} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/ldmatrix.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/ldmatrix.rs new file mode 100644 index 0000000000..db98a57a5a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/ldmatrix.rs @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Check shared-memory matrix-load layout restrictions. + +use super::thread_value::{host_count, positive_count}; +use super::{ValidationError, ValidationResult}; +use crate::{ComposedLayout, IntTuple}; + +/// Prove a shared tile can feed `ldmatrix` for `f16` operands. +/// +/// `ldmatrix` reads one 8x8 `f16` matrix as eight row segments of 16 bytes, +/// each supplied by one lane's address. The hardware requires every segment +/// to be 16 consecutive, 16-byte-aligned shared bytes. This check enumerates +/// every 8-column-aligned segment of the whole tile through the composed +/// byte map (swizzle included), so any warp-tile window the compiler later +/// addresses is covered: +/// +/// ```text +/// for every row r, every column block c8 (multiples of 8): +/// bytes of elements (r, c8..c8+8) must be base, base+2, ..., base+14 +/// and base % 16 == 0 +/// ``` +/// +/// Layouts that scatter a row (column-major smem, byte-hostile swizzles) +/// fail here at compile time instead of producing garbage fragments. +pub fn validate_ldmatrix_source( + smem_byte_layout: &ComposedLayout, + rows: i64, + columns: i64, +) -> ValidationResult<()> { + let rows = positive_count("ldmatrix rows", rows)?; + let columns = positive_count("ldmatrix columns", columns)?; + if rows % 8 != 0 || columns % 8 != 0 { + return Err(ValidationError::InvalidCount { + name: "ldmatrix tile extents must be multiples of 8", + value: rows * columns, + }); + } + let total = rows + .checked_mul(columns) + .ok_or(ValidationError::CountTooLarge { + name: "ldmatrix tile elements", + value: i64::MAX, + })?; + host_count("ldmatrix tile elements", total)?; + + for row in 0..rows { + for block in 0..(columns / 8) { + let first_cell = row + rows * (block * 8); + let base = smem_byte_layout + .checked_call(&IntTuple::Leaf(first_cell)) + .ok_or(ValidationError::InvalidCount { + name: "ldmatrix segment base is unmappable", + value: first_cell, + })?; + if base % 16 != 0 { + return Err(ValidationError::AtomStartMisaligned { + thread: row, + value: block, + offset: base, + atom_bytes: 16, + }); + } + for element in 1..8 { + let cell = row + rows * (block * 8 + element); + let byte = smem_byte_layout.checked_call(&IntTuple::Leaf(cell)).ok_or( + ValidationError::InvalidCount { + name: "ldmatrix segment element is unmappable", + value: cell, + }, + )?; + if byte != base + element * 2 { + return Err(ValidationError::AtomNotContiguous { + thread: row, + value: block, + byte_in_atom: element * 2, + expected: base + element * 2, + actual: byte, + }); + } + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/mod.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/mod.rs new file mode 100644 index 0000000000..72a9f53e76 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/mod.rs @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Check a cooperative copy's compile-time layout facts before the compiler +//! turns it into GPU copy instructions. +//! +//! A **tile** is the small, fixed rectangle copied by one CUDA block. A +//! [**mode**](crate#modes) is one part of a coordinate that you can change while +//! keeping the other parts fixed. The first supported form of `copy_g2s` +//! requires two top-level tile modes: rows and columns. +//! +//! ```text +//! 2-row x 3-column tile +//! +//! column 0 column 1 column 2 +//! row 0 A B C +//! row 1 D E F +//! ``` +//! +//! A **thread-value (TV) assignment** answers: “which tile cell does local +//! value `v` of thread `t` own?” For example: +//! +//! ```text +//! TV(0,0) -> A TV(0,1) -> B TV(0,2) -> C +//! TV(1,0) -> D TV(1,1) -> E TV(1,2) -> F +//! ``` +//! +//! **Exact coverage** means every cell has exactly one TV owner. A hole means +//! an element is never copied; two owners mean two threads can write the same +//! destination. Both are rejected. +//! +//! A **compact thread layout** produces every thread ID from `0` through +//! `thread_count - 1` exactly once. The IDs may be reordered, but none may be +//! missing or repeated. A [**copy atom**](crate#atoms) describes one indivisible +//! transaction such as a 16-byte `cp.async`; all values assigned to that atom +//! must form one aligned, consecutive byte range. +//! +//! Global memory uses **row-major** order: columns are adjacent, and +//! `offset(row, column) = row * column_count + column`. Shared memory may use +//! a composed swizzle, so its complete byte map is checked separately. +//! +//! These validators exhaustively inspect small compile-time maps and return the +//! minimum shared-memory capacity. The unsafe copy caller supplies valid and +//! aligned pointers, the promised row pitch, sufficient shared storage, and +//! the required participating threads. The compiler uses runtime matrix +//! bounds to zero-fill the unread portion of each edge copy. + +mod cooperative; +mod error; +mod ldmatrix; +#[cfg(test)] +mod tests; +mod thread_value; +mod tma; + +pub use cooperative::{CooperativeCopyPlan, CooperativePlanError, validate_cooperative_copy_plan}; +pub use error::ValidationError; +pub use ldmatrix::validate_ldmatrix_source; +pub use thread_value::{ + validate_atom_compatibility, validate_compact_thread_layout, + validate_per_thread_atom_contiguity, validate_tv_exact_coverage, validate_tv_layout, +}; +pub use tma::{tma_phase_alignment_bytes, validate_tma_encodable}; + +/// Result returned by the reusable low-level layout checks in this module. +pub type ValidationResult = core::result::Result; + +/// Largest element count that an exhaustive static proof may enumerate. +/// +/// One million elements is well above a practical shared-memory tile. The +/// limit prevents a malformed compiler attribute from causing an OOM or hang. +pub const MAX_STATIC_VALIDATION_ELEMENTS: i64 = 1 << 20; + +/// Largest byte map that the exact shared-memory proof may enumerate. +pub const MAX_STATIC_VALIDATION_BYTES: i64 = 1 << 20; + +/// CUDA's maximum number of threads in one cooperative block. +pub const MAX_COOPERATIVE_THREADS: i64 = 1024; + +/// Largest absolute static stride or offset accepted before layout algebra. +/// +/// This bound also catches huge strides hidden under a size-one +/// [mode](crate#modes). +pub const MAX_STATIC_LAYOUT_OFFSET_MAGNITUDE: u64 = 1 << 20; diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tests.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tests.rs new file mode 100644 index 0000000000..f9f1f66577 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tests.rs @@ -0,0 +1,331 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression cases for the composed layout validators. + +use super::thread_value::tv_coord; +use super::*; +use crate::Swizzle; +use crate::algebra::make_layout_tv; +use crate::{ComposedLayout, IntTuple, Layout, OffsetUnit}; +use alloc::string::ToString; +use alloc::{vec, vec::Vec}; + +/// Parse a test layout literal and fail immediately if the fixture is malformed. +fn l(text: &str) -> Layout { + text.parse().unwrap() +} + +/// Construct the checked cooperative-copy plan used by the live decode canary. +fn canary_plan() -> Result { + let thread = l("(6,1):(1,0)"); + let value = l("(1,4):(0,1)"); + let tile = l("(6,4):(4,1)"); + let smem = + ComposedLayout::new(Swizzle::new(3, 4, 3), 8, tile.clone(), OffsetUnit::Elements).unwrap(); + validate_cooperative_copy_plan(16, &thread, &value, &tile, &smem, 4) +} + +/// The complete copy plan must match the canary’s thread ownership and shared capacity. +#[test] +fn high_level_plan_matches_the_live_decode_canary() { + let plan = canary_plan().unwrap(); + assert_eq!(plan.tiler_shape, vec![6, 4]); + assert_eq!(plan.thread_count, 6); + assert_eq!(plan.values_per_thread, 4); + assert_eq!(plan.tile_elements, 24); + assert_eq!(plan.minimum_smem_capacity, 32); +} + +/// The global source map must be canonical row-major even if a digit permutation covers the same cells. +#[test] +fn high_level_plan_rejects_a_digit_permuted_row_major_tile() { + let thread = l("(1,1):(1,0)"); + // The logical columns are 0,1,2,3, but this nested digit mode visits + // them as 0,2,1,3. Giving TV and TileL the same permutation makes the + // two mistakes cancel and appear contiguous: + // + // TV order: 0 2 1 3 + // TileL maps to: 0 1 2 3 + // + // A real row-major GmemMatrix still uses logical order 0,1,2,3, so + // the high-level validator must reject the permuted TileL itself. + let digit_permutation = l("(1,(2,2)):(0,(2,1))"); + let (_, tv) = make_layout_tv(&thread, &digit_permutation).unwrap(); + let apparent_offsets: Vec<_> = (0..4) + .map(|value| { + let cell = tv.checked_call(&tv_coord(0, value)).unwrap(); + digit_permutation + .checked_call(&IntTuple::Leaf(cell)) + .unwrap() + }) + .collect(); + assert_eq!(apparent_offsets, vec![0, 1, 2, 3]); + + let smem = ComposedLayout::from_layout(digit_permutation.clone(), OffsetUnit::Elements); + let error = validate_cooperative_copy_plan( + 16, + &thread, + &digit_permutation, + &digit_permutation, + &smem, + 4, + ) + .unwrap_err() + .to_string(); + assert!(error.contains("not canonical row-major"), "{error}"); +} + +/// Plan validation must reject excessive threads and huge strides hidden in size-one modes. +#[test] +fn high_level_plan_bounds_threads_and_hidden_raw_strides() { + let thread = l("(1025,1):(1,0)"); + let value = l("(1,1):(0,0)"); + let tile = l("(1025,1):(1,0)"); + let smem = ComposedLayout::from_layout(tile.clone(), OffsetUnit::Elements); + let error = validate_cooperative_copy_plan(4, &thread, &value, &tile, &smem, 4) + .unwrap_err() + .to_string(); + assert!(error.contains("CUDA block limit 1024"), "{error}"); + + // A size-one mode has only coordinate zero, which hides its stride: + // + // 0 * i64::MAX = 0 + // + // Raw-stride validation sees i64::MAX before algebra can multiply it. + let hidden_stride = l("(6,1):(1,9223372036854775807)"); + let tile = l("(6,4):(4,1)"); + let smem = ComposedLayout::from_layout(tile.clone(), OffsetUnit::Elements); + let error = + validate_cooperative_copy_plan(16, &hidden_stride, &l("(1,4):(0,1)"), &tile, &smem, 4) + .unwrap_err() + .to_string(); + assert!(error.contains("layout stride"), "{error}"); +} + +/// Thread IDs may be permuted but must cover a compact interval without holes. +#[test] +fn compact_threads_allow_permutation_but_reject_gaps() { + validate_compact_thread_layout(&l("(2,2):(2,1)")).unwrap(); + validate_compact_thread_layout(&l("(2,2):(1,2)")).unwrap(); + + assert!(matches!( + validate_compact_thread_layout(&l("4:2")), + Err(ValidationError::ThreadOutOfRange { .. }) + )); + assert!(matches!( + validate_compact_thread_layout(&l("(2,2):(0,1)")), + Err(ValidationError::DuplicateThread { .. }) + )); + + let too_large = Layout::contiguous(MAX_STATIC_VALIDATION_ELEMENTS + 1); + assert!(matches!( + validate_compact_thread_layout(&too_large), + Err(ValidationError::CountTooLarge { .. }) + )); +} + +/// Thread/value ownership must reject duplicate owners, missing cells and invalid extents. +#[test] +fn tv_layout_must_be_an_exact_partition() { + let valid = l("(4,2):(2,1)"); + validate_tv_exact_coverage(&valid, 4, 2, 8).unwrap(); + + let duplicate = l("(4,2):(1,1)"); + assert!(matches!( + validate_tv_exact_coverage(&duplicate, 4, 2, 8), + Err(ValidationError::DuplicateCell { .. }) + )); + + let out_of_range = l("(4,2):(3,1)"); + assert!(matches!( + validate_tv_exact_coverage(&out_of_range, 4, 2, 8), + Err(ValidationError::TvOutOfRange { .. }) + )); +} + +/// The constructed thread/value map must satisfy coverage and compactness proofs. +#[test] +fn make_layout_tv_result_passes_all_static_checks() { + let threads = l("(2,3):(3,1)"); + let values = l("(2,2):(2,1)"); + let (tiler, tv) = crate::algebra::make_layout_tv(&threads, &values).unwrap(); + validate_tv_layout(&threads, &values, &tiler, &tv).unwrap(); +} + +/// The tile’s value count must divide into complete instruction-width copy atoms. +#[test] +fn atom_counts_must_divide() { + validate_atom_compatibility(256, 4, 8, 2).unwrap(); + assert_eq!( + validate_atom_compatibility(250, 4, 8, 2), + Err(ValidationError::AtomDoesNotDivide { + axis: "thread", + copy_count: 250, + atom_count: 8, + }) + ); + assert!(matches!( + validate_atom_compatibility(256, 3, 8, 2), + Err(ValidationError::AtomDoesNotDivide { axis: "value", .. }) + )); +} + +/// Complete adjacent aligned per-thread atoms must pass static contiguity validation. +#[test] +fn contiguous_aligned_atoms_pass() { + // Four threads each own four f32 values: one 16-byte atom. + // + // cell: 0 1 2 3 | 4 5 6 7 | ... + // owner: T0 T0 T0 T0 |T1 T1 T1 T1 | ... + let tv = l("(4,4):(4,1)"); + let elements = + ComposedLayout::new(Swizzle::new(2, 2, 3), 0, l("16:1"), OffsetUnit::Elements).unwrap(); + let bytes = elements.to_byte_offsets(4).unwrap(); + validate_per_thread_atom_contiguity(&tv, 4, 4, 16, 4, &bytes, 16).unwrap(); +} + +/// A byte swizzle preserving four low bits must keep aligned 16-byte atoms intact. +#[test] +fn m_four_swizzle_moves_whole_sixteen_byte_atoms() { + // The map spans enough bytes to reach the high bits that feed the + // XOR. A smaller map would leave every source bit zero and make the + // swizzle look like the identity by accident. + let tv = l("(32,4):(4,1)"); + let memory = + ComposedLayout::new(Swizzle::new(3, 4, 3), 0, l("128:4"), OffsetUnit::Bytes).unwrap(); + validate_per_thread_atom_contiguity(&tv, 32, 4, 16, 4, &memory, 16).unwrap(); +} + +/// A scattered per-thread value map must be rejected as a single vector transaction. +#[test] +fn scattered_values_are_not_one_atom() { + // Coverage is complete, but each thread is raked through the tile: + // + // cell: 0 1 2 3 | 4 5 6 7 | ... + // owner: T0 T1 T2 T3 |T0 T1 T2 T3 | ... + let tv = l("(4,4):(1,4)"); + let memory = ComposedLayout::from_layout(l("16:4"), OffsetUnit::Bytes); + assert!(matches!( + validate_per_thread_atom_contiguity(&tv, 4, 4, 16, 4, &memory, 16), + Err(ValidationError::AtomNotContiguous { .. }) + )); +} + +/// Swizzles that split or scatter an instruction atom must be rejected. +#[test] +fn swizzle_must_preserve_the_whole_atom() { + let tv = l("(2,4):(4,1)"); + let memory = ComposedLayout::new( + // M=3 leaves the lowest three address bits unchanged, which + // guarantees only 2^3 = 8 adjacent bytes. Here the XOR splits a + // 16-byte transaction, and the exact byte proof catches it. + Swizzle::new(1, 3, 3), + 64, + l("8:4"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert!(matches!( + validate_per_thread_atom_contiguity(&tv, 2, 4, 16, 4, &memory, 16), + Err(ValidationError::AtomStartMisaligned { .. }) + | Err(ValidationError::AtomNotContiguous { .. }) + )); +} + +/// Overflow in an expected atom address must return a validation error instead of wrapping. +#[test] +fn expected_atom_address_overflow_is_a_validation_error() { + // The swizzle maps MAX-1 to MAX, while the next input remains valid: + // + // input: MAX-1 MAX + // output: MAX MAX-1 + let tv = l("(1,7):(7,1)"); + let memory = ComposedLayout::new( + Swizzle::new(1, 0, 1), + i64::MAX - 1, + l("7:1"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert_eq!( + validate_per_thread_atom_contiguity(&tv, 1, 7, 7, 1, &memory, 7), + Err(ValidationError::ArithmeticOverflow( + "expected atom byte address" + )) + ); +} + +/// Unswizzled u16 tiles must retain TMA’s minimum shared-base alignment. +#[test] +fn tma_identity_u16_keeps_the_hardware_minimum_alignment() { + let layout = ComposedLayout::new( + Swizzle::IDENTITY, + 0, + l("(1,256):(256,1)"), + OffsetUnit::Elements, + ) + .unwrap(); + + assert_eq!(tma_phase_alignment_bytes(&layout, 2), Ok(16)); +} + +/// B64 byte tiles must require a shared base aligned to the full 512-byte swizzle period. +#[test] +fn tma_b64_u8_needs_one_complete_512_byte_phase() { + let layout = ComposedLayout::new( + Swizzle::new(2, 4, 3), + 0, + l("(128,64):(64,1)"), + OffsetUnit::Elements, + ) + .unwrap(); + + assert_eq!(tma_phase_alignment_bytes(&layout, 1), Ok(512)); +} + +/// B128 FP16 tiles must require a shared base aligned to the full 1024-byte swizzle period. +#[test] +fn tma_b128_f16_needs_one_complete_1024_byte_phase() { + let layout = ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + l("(128,64):(64,1)"), + OffsetUnit::Elements, + ) + .unwrap(); + + assert_eq!(tma_phase_alignment_bytes(&layout, 2), Ok(1024)); +} + +/// Phase-alignment queries must reject tile/swizzle combinations outside TMA’s supported modes. +#[test] +fn tma_phase_alignment_rejects_unencodable_inputs() { + let column_major = + ComposedLayout::new(Swizzle::IDENTITY, 0, l("(2,4):(1,2)"), OffsetUnit::Elements).unwrap(); + assert!(tma_phase_alignment_bytes(&column_major, 2).is_err()); + + let b64_u8 = ComposedLayout::new( + Swizzle::new(2, 4, 3), + 0, + l("(128,64):(64,1)"), + OffsetUnit::Elements, + ) + .unwrap(); + assert!( + tma_phase_alignment_bytes(&b64_u8, 2).is_err(), + "the element width changes the byte-unit swizzle and row pitch" + ); + assert!(tma_phase_alignment_bytes(&b64_u8, 0).is_err()); + + let byte_offsets = ComposedLayout::new( + Swizzle::IDENTITY, + 0, + l("(1,256):(256,1)"), + OffsetUnit::Bytes, + ) + .unwrap(); + assert!(tma_phase_alignment_bytes(&byte_offsets, 2).is_err()); +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/thread_value.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/thread_value.rs new file mode 100644 index 0000000000..15a1dac288 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/thread_value.rs @@ -0,0 +1,366 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Thread ownership, exact coverage and copy-atom contiguity. + +use super::{MAX_STATIC_VALIDATION_ELEMENTS, ValidationError, ValidationResult}; +use crate::{ComposedLayout, IntTuple, Layout, OffsetUnit}; +use alloc::{vec, vec::Vec}; + +/// Require a compact thread layout. +/// +/// “Compact” means the layout produces every thread ID from zero through +/// `size - 1` exactly once. It describes membership, not physical ordering, so +/// a permutation is valid while a gap or duplicate is not: +/// +/// ```text +/// flat coordinate 0 1 2 3 +/// unpacked pair for tuple shapes (0,0) (1,0) (0,1) (1,1) +/// +/// layout (2,2):(2,1) -> thread ID 0 2 1 3 valid: 0,1,2,3 appear once +/// layout 4:2 -> thread ID 0 2 4 6 invalid: 1 and 3 are missing +/// layout (2,2):(1,1) -> thread ID 0 1 1 2 invalid: ID 1 appears twice +/// ``` +pub fn validate_compact_thread_layout(thread_layout: &Layout) -> ValidationResult<()> { + let thread_count = thread_layout + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("thread layout size"))?; + let thread_count_usize = host_count("thread count", thread_count)?; + let mut owners: Vec> = vec![None; thread_count_usize]; + + for coordinate in 0..thread_count { + let thread = thread_layout + .checked_call(&IntTuple::Leaf(coordinate)) + .ok_or(ValidationError::ArithmeticOverflow("thread-layout address"))?; + if !(0..thread_count).contains(&thread) { + return Err(ValidationError::ThreadOutOfRange { + coordinate, + thread, + thread_count, + }); + } + let slot = &mut owners[thread as usize]; + if let Some(first_coordinate) = *slot { + return Err(ValidationError::DuplicateThread { + thread, + first_coordinate, + second_coordinate: coordinate, + }); + } + *slot = Some(coordinate); + } + Ok(()) +} + +/// Prove exact coverage of a tile by a thread-value (TV) assignment. +/// +/// A TV layout has two top-level [modes](crate#modes): thread ID and that +/// thread's local value number. Exact coverage requires a one-to-one match: +/// +/// ```text +/// tv = (2,2):(2,1) +/// +/// value 0 value 1 +/// thread 0 0 1 +/// thread 1 2 3 +/// ^ numbers are tile cells; every cell appears once +/// ``` +/// +/// Wrong mode sizes, a mismatched assignment count, an out-of-range cell, or a +/// duplicate owner produces a specific [`ValidationError`]. When assignment +/// count equals cell count, an unowned cell necessarily comes with a duplicate +/// or out-of-range assignment and is rejected by that earlier check. +pub fn validate_tv_exact_coverage( + tv: &Layout, + thread_count: i64, + values_per_thread: i64, + tile_elements: i64, +) -> ValidationResult<()> { + let thread_count = positive_count("thread count", thread_count)?; + let values_per_thread = positive_count("values per thread", values_per_thread)?; + let tile_elements = positive_count("tile elements", tile_elements)?; + + let modes = tv.modes(); + if modes.len() != 2 { + return Err(ValidationError::TvRank { + actual: modes.len(), + }); + } + let thread_mode_size = modes[0] + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("TV thread-mode size"))?; + if thread_mode_size != thread_count { + return Err(ValidationError::TvModeSize { + mode: "thread", + expected: thread_count, + actual: thread_mode_size, + }); + } + let value_mode_size = modes[1] + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("TV value-mode size"))?; + if value_mode_size != values_per_thread { + return Err(ValidationError::TvModeSize { + mode: "value", + expected: values_per_thread, + actual: value_mode_size, + }); + } + + let assignments = thread_count + .checked_mul(values_per_thread) + .ok_or(ValidationError::ArithmeticOverflow("TV assignment count"))?; + if assignments != tile_elements { + return Err(ValidationError::TvCardinality { + assignments, + tile_elements, + }); + } + + let mut owners: Vec> = + vec![None; host_count("tile elements", tile_elements)?]; + for thread in 0..thread_count { + for value in 0..values_per_thread { + let cell = tv + .checked_call(&tv_coord(thread, value)) + .ok_or(ValidationError::ArithmeticOverflow("TV-layout address"))?; + if !(0..tile_elements).contains(&cell) { + return Err(ValidationError::TvOutOfRange { + thread, + value, + cell, + tile_elements, + }); + } + let slot = &mut owners[cell as usize]; + if let Some(first_owner) = *slot { + return Err(ValidationError::DuplicateCell { + cell, + first_owner, + second_owner: (thread, value), + }); + } + *slot = Some((thread, value)); + } + } + if let Some(cell) = owners.iter().position(Option::is_none) { + return Err(ValidationError::MissingCell { cell: cell as i64 }); + } + Ok(()) +} + +/// Check a TV layout produced by [`make_layout_tv`](crate::algebra::make_layout_tv) against its source +/// layouts. +/// +/// `tiler_shape` contains one extent per tile [mode](crate#modes). Their product +/// is the tile cell count. This function first requires compact thread IDs, +/// then requires the returned `(thread, local value) -> cell` map to cover that +/// count exactly. +pub fn validate_tv_layout( + thread_layout: &Layout, + value_layout: &Layout, + tiler_shape: &[i64], + tv: &Layout, +) -> ValidationResult<()> { + validate_compact_thread_layout(thread_layout)?; + let tile_elements = tiler_shape.iter().try_fold(1i64, |product, &extent| { + let extent = positive_count("tiler extent", extent)?; + product + .checked_mul(extent) + .ok_or(ValidationError::ArithmeticOverflow("tiler element count")) + })?; + let thread_count = thread_layout + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("thread layout size"))?; + let values_per_thread = value_layout + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("value layout size"))?; + validate_tv_exact_coverage(tv, thread_count, values_per_thread, tile_elements) +} + +/// Require the cooperative copy to contain only whole groups described by its +/// [copy atom](crate#atoms). +/// +/// A [copy atom](crate#atoms) describes how many threads participate and how +/// many local values each contributes to one hardware transaction. Both copy +/// counts must divide evenly: +/// +/// ```text +/// 256 copy threads / 8 atom threads = 32 whole groups valid +/// 250 copy threads / 8 atom threads = 31 + remainder invalid +/// ``` +/// +/// `CopyG2S` uses one thread per atom, but this helper also supports later +/// atom kinds that group several threads. +pub fn validate_atom_compatibility( + copy_threads: i64, + copy_values_per_thread: i64, + atom_threads: i64, + atom_values_per_thread: i64, +) -> ValidationResult<()> { + let copy_threads = positive_count("copy thread count", copy_threads)?; + let copy_values_per_thread = positive_count("copy values per thread", copy_values_per_thread)?; + let atom_threads = positive_count("atom thread count", atom_threads)?; + let atom_values_per_thread = positive_count("atom values per thread", atom_values_per_thread)?; + + if copy_threads % atom_threads != 0 { + return Err(ValidationError::AtomDoesNotDivide { + axis: "thread", + copy_count: copy_threads, + atom_count: atom_threads, + }); + } + if copy_values_per_thread % atom_values_per_thread != 0 { + return Err(ValidationError::AtomDoesNotDivide { + axis: "value", + copy_count: copy_values_per_thread, + atom_count: atom_values_per_thread, + }); + } + Ok(()) +} + +/// Prove that the transaction described by each one-thread +/// [copy atom](crate#atoms) maps to one aligned, consecutive byte range. +/// +/// The transaction described by a 16-byte copy atom must map like this: +/// +/// ```text +/// byte in atom: 0 1 2 ... 15 +/// physical offset: P P+1 P+2 ... P+15 +/// ^ P must also be 16-byte aligned +/// ``` +/// +/// For a smaller concrete example, let two threads each own two four-byte +/// values: +/// +/// ```text +/// tv = (2,2):(2,1) memory = 4:4 byte offsets +/// +/// owner tile cells value-start bytes complete 8-byte atom +/// thread 0 0, 1 0, 4 0 through 7 +/// thread 1 2, 3 8, 12 8 through 15 +/// ``` +/// +/// Both atom starts are 8-byte aligned, and every byte through the end of each +/// interval follows consecutively. An interleaved TV map such as +/// `(2,2):(1,2)` would instead give each thread cells `0,2` or `1,3`, so the +/// check would reject it for this simple memory layout. +/// +/// The proof evaluates the complete byte-addressed memory map, including its +/// fixed offset and swizzle. It does not reject solely because +/// [the swizzle's general alignment estimate](crate::Swizzle::max_alignment) +/// is smaller; evaluating the exact map can prove a wider interval in a +/// special case. +/// +/// This function is for `CopyG2S`, where one thread owns each atom. +/// Multi-thread atom kinds also need [`validate_atom_compatibility`] plus a +/// compiler check specific to that atom's ownership rule. +pub fn validate_per_thread_atom_contiguity( + tv: &Layout, + thread_count: i64, + values_per_thread: i64, + atom_bytes: i64, + element_bytes: i64, + memory: &ComposedLayout, + base_alignment_bytes: i64, +) -> ValidationResult<()> { + let atom_bytes = positive_count("atom bytes", atom_bytes)?; + let element_bytes = positive_count("element bytes", element_bytes)?; + if atom_bytes % element_bytes != 0 { + return Err(ValidationError::AtomSize { + atom_bytes, + element_bytes, + }); + } + if memory.unit() != OffsetUnit::Bytes { + return Err(ValidationError::AtomNeedsByteLayout); + } + let memory_size = memory + .outer() + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("memory layout size"))?; + validate_tv_exact_coverage(tv, thread_count, values_per_thread, memory_size)?; + + let values_per_atom = atom_bytes / element_bytes; + let values_per_atom_usize = host_count("atom values per thread", values_per_atom)?; + validate_atom_compatibility(thread_count, values_per_thread, 1, values_per_atom)?; + + let base_alignment_bytes = positive_count("base alignment", base_alignment_bytes)?; + if base_alignment_bytes % atom_bytes != 0 { + return Err(ValidationError::BaseAlignment { + required: atom_bytes, + guaranteed: base_alignment_bytes, + }); + } + + for thread in 0..thread_count { + for first_value in (0..values_per_thread).step_by(values_per_atom_usize) { + let first_cell = tv + .checked_call(&tv_coord(thread, first_value)) + .ok_or(ValidationError::ArithmeticOverflow("TV atom start"))?; + let first_offset = memory + .checked_call(&IntTuple::Leaf(first_cell)) + .ok_or(ValidationError::ArithmeticOverflow("memory atom start"))?; + if first_offset.rem_euclid(atom_bytes) != 0 { + return Err(ValidationError::AtomStartMisaligned { + thread, + value: first_value, + offset: first_offset, + atom_bytes, + }); + } + + for byte_in_atom in 0..atom_bytes { + let value_delta = byte_in_atom / element_bytes; + let byte_in_element = byte_in_atom % element_bytes; + let value = first_value + .checked_add(value_delta) + .ok_or(ValidationError::ArithmeticOverflow("TV atom value"))?; + let cell = tv + .checked_call(&tv_coord(thread, value)) + .ok_or(ValidationError::ArithmeticOverflow("TV atom address"))?; + let actual = memory + .checked_call_with_outer_delta(&IntTuple::Leaf(cell), byte_in_element) + .ok_or(ValidationError::ArithmeticOverflow("memory atom address"))?; + let expected = first_offset.checked_add(byte_in_atom).ok_or( + ValidationError::ArithmeticOverflow("expected atom byte address"), + )?; + if actual != expected { + return Err(ValidationError::AtomNotContiguous { + thread, + value: first_value, + byte_in_atom, + expected, + actual, + }); + } + } + } + } + Ok(()) +} + +/// Build the tuple coordinate selecting one thread and its local value. +pub(super) fn tv_coord(thread: i64, value: i64) -> IntTuple { + IntTuple::Tuple(vec![IntTuple::Leaf(thread), IntTuple::Leaf(value)]) +} + +/// Require a positive bounded count before exhaustive layout enumeration. +pub(super) fn positive_count(name: &'static str, value: i64) -> ValidationResult { + if value <= 0 { + Err(ValidationError::InvalidCount { name, value }) + } else { + Ok(value) + } +} + +/// Convert a validated count to the host indexing type without truncation. +pub(super) fn host_count(name: &'static str, value: i64) -> ValidationResult { + if value > MAX_STATIC_VALIDATION_ELEMENTS { + return Err(ValidationError::CountTooLarge { name, value }); + } + usize::try_from(value).map_err(|_| ValidationError::CountTooLarge { name, value }) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tma.rs b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tma.rs new file mode 100644 index 0000000000..31e34769f2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/src/validation/tma.rs @@ -0,0 +1,136 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Check tensor-memory-accelerator encoding and phase alignment. + +use super::thread_value::positive_count; +use super::{ValidationError, ValidationResult}; +use crate::{ComposedLayout, OffsetUnit}; + +/// Prove a shared tile is one of the four arrangements TMA can encode. +/// +/// TMA writes row-major boxes scrambled by at most one of the fixed +/// byte-unit patterns `S` (B in 0..=3), with no composed offset. +/// Swizzled boxes must span the whole row. This is the compiler-side twin +/// of the host encoder's checks, so an unencodable `SmemL` fails at KERNEL +/// compile time instead of at host run time. +pub fn validate_tma_encodable( + smem_layout: &ComposedLayout, + element_bytes: i64, +) -> ValidationResult<()> { + positive_count("TMA element bytes", element_bytes)?; + if element_bytes & (element_bytes - 1) != 0 { + return Err(ValidationError::InvalidCount { + name: "TMA element byte width must be a power of two", + value: element_bytes, + }); + } + if smem_layout.unit() != OffsetUnit::Elements { + return Err(ValidationError::InvalidCount { + name: "TMA shared layout offsets must count elements", + value: 0, + }); + } + let modes = smem_layout.outer().modes(); + if modes.len() != 2 { + return Err(ValidationError::TvRank { + actual: modes.len(), + }); + } + let rows = modes[0] + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("TMA tile rows"))?; + let columns = modes[1] + .checked_size() + .ok_or(ValidationError::ArithmeticOverflow("TMA tile columns"))?; + positive_count("TMA tile rows", rows)?; + positive_count("TMA tile columns", columns)?; + let row_leaves = modes[0].stride.flatten(); + let col_leaves = modes[1].stride.flatten(); + if row_leaves != [columns] || col_leaves != [1] { + return Err(ValidationError::InvalidCount { + name: "TMA tile must be row-major (row stride = columns, column stride = 1)", + value: rows * columns, + }); + } + if smem_layout.offset() != 0 { + return Err(ValidationError::InvalidCount { + name: "TMA cannot express a composed offset", + value: smem_layout.offset(), + }); + } + let swizzle = smem_layout.inner(); + if swizzle.num_bits == 0 { + return Ok(()); + } + let m_bytes = i64::from(swizzle.num_base) + element_bytes.trailing_zeros() as i64; + if swizzle.num_bits > 3 || m_bytes != 4 || swizzle.num_shift != 3 { + return Err(ValidationError::InvalidCount { + name: "swizzle is not one of TMA's encodable byte-unit S modes", + value: i64::from(swizzle.num_bits), + }); + } + let pitch_bytes = columns * element_bytes; + if pitch_bytes != (16 << swizzle.num_bits) { + return Err(ValidationError::InvalidCount { + name: "swizzled TMA tile pitch must equal the swizzle span", + value: pitch_bytes, + }); + } + Ok(()) +} + +/// Return the shared-base alignment required to keep one TMA swizzle phase. +/// +/// TMA always needs at least a 16-byte shared address. A swizzle adds higher +/// address bits that choose the phase of its XOR mapping: +/// +/// ```text +/// byte-unit S +/// +/// low 4 bits target B bits source B bits +/// [ unchanged ] [ XOR target ] +3 start [ XOR source ] +/// 0 ........ 3 4 .... 4+B-1 7 .... 7+B-1 +/// +/// highest participating bit + 1 = 4 + B + 3 +/// required alignment = 2^(4 + B + 3) bytes +/// ``` +/// +/// That gives the layouts used by the GEMM transport path: +/// +/// ```text +/// no swizzle 16 bytes +/// B64, B = 2 512 bytes +/// B128, B = 3 1024 bytes +/// ``` +/// +/// This is the minimum alignment promised by the unsafe TMA copy. It is not +/// a claim about the descriptor pointer or the hidden global-memory address. +/// The layout is checked with [`validate_tma_encodable`] first, so callers +/// cannot obtain an alignment for a layout the hardware cannot represent. +pub fn tma_phase_alignment_bytes( + smem_layout: &ComposedLayout, + element_bytes: i64, +) -> ValidationResult { + validate_tma_encodable(smem_layout, element_bytes)?; + let swizzle = smem_layout.inner(); + if swizzle.num_bits == 0 { + return Ok(16); + } + + // Validation above normalized the element-based swizzle to byte-unit + // S. The highest participating source bit is therefore 4+B+3-1. + let phase_bits = 4_u32 + .checked_add(swizzle.num_bits) + .and_then(|bits| bits.checked_add(3)) + .ok_or(ValidationError::ArithmeticOverflow( + "TMA swizzle phase alignment exponent", + ))?; + 1_i64 + .checked_shl(phase_bits) + .ok_or(ValidationError::ArithmeticOverflow( + "TMA swizzle phase alignment", + )) +} diff --git a/cuda-oxide/cute-rs-preview/cute-layout/tests/public_paths.rs b/cuda-oxide/cute-rs-preview/cute-layout/tests/public_paths.rs new file mode 100644 index 0000000000..7deb337ee3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-layout/tests/public_paths.rs @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Public-path compatibility across the responsibility-based layout modules. + +use cute_layout::{IntTuple, Layout, ParseLayoutError, algebra, crd2idx}; +use cute_layout::{layout, parse, tuple}; + +/// Grouped paths must expose the same types and mappings as the existing imports. +#[test] +fn grouped_and_compatibility_paths_share_layout_contracts() { + let layout: layout::Layout = "(2,3):(3,1)".parse().unwrap(); + let coordinate: tuple::IntTuple = IntTuple::Tuple(vec![IntTuple::Leaf(1), IntTuple::Leaf(2)]); + let root_layout: &Layout = &layout; + assert_eq!(root_layout.checked_call(&coordinate), Some(5)); + assert_eq!(crd2idx(&coordinate, &layout.shape, &layout.stride), 5); + assert_eq!( + layout::crd2idx(&coordinate, &layout.shape, &layout.stride), + 5 + ); + + let error: parse::ParseLayoutError = "(2,3):1".parse::().unwrap_err(); + let root_error: ParseLayoutError = error; + assert_eq!(root_error.to_string(), "invalid CuTe layout/tuple syntax"); +} + +/// Algebra families must remain interoperable with the original algebra imports. +#[test] +fn grouped_algebra_paths_preserve_coordinate_maps() { + let whole = Layout::contiguous(8); + let tile = Layout::contiguous(4); + let divided = algebra::tiling::logical_divide(&whole, &tile).unwrap(); + let divided_compat = algebra::logical_divide(&whole, &tile).unwrap(); + assert_eq!(divided, divided_compat); + assert_eq!(divided.to_string(), "(4,2):(1,4)"); + assert_eq!(algebra::mapping::coalesce(÷d), whole); + assert_eq!(algebra::coalesce(÷d), whole); + let zipped = algebra::tiling::zipped_divide(&whole, &[4]).unwrap(); + assert_eq!(zipped, algebra::zipped_divide(&whole, &[4]).unwrap()); + assert_eq!(zipped.to_string(), "(4,2):(1,4)"); + let nested: Layout = "((2,4),3):((1,2),8)".parse().unwrap(); + assert_eq!( + algebra::mapping::flatten(&nested), + algebra::flatten(&nested) + ); + assert_eq!(algebra::flatten(&nested).to_string(), "(2,4,3):(1,2,8)"); + + let thread: Layout = "(2,3):(3,1)".parse().unwrap(); + let value: Layout = "(2,2):(2,1)".parse().unwrap(); + let (shape, tv) = algebra::thread_value::make_layout_tv(&thread, &value).unwrap(); + let tv_compat = algebra::make_layout_tv(&thread, &value).unwrap(); + assert_eq!(shape, vec![4, 6]); + assert_eq!(tv.to_string(), "((3,2),(2,2)):((8,2),(4,1))"); + assert_eq!((shape, tv), tv_compat); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/Cargo.toml b/cuda-oxide/cute-rs-preview/cute-rs/Cargo.toml new file mode 100644 index 0000000000..466425a575 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "cute-rs" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "CuTe-style device library: Tensor, Layout types, partitioning, copy/mma atoms" + +[dependencies] +ctm-rs = { workspace = true } +cute-layout = { workspace = true } +cuda-device = { workspace = true } +cuda-core = { workspace = true, optional = true } + +[features] +# Host-side helpers (TMA descriptor encoding). Kernels never need this. +host = ["dep:cuda-core"] + +[dev-dependencies] +# The compile-only example also expands host metadata in mixed workspace checks. +cuda-host = { workspace = true } +trybuild = "1" diff --git a/cuda-oxide/cute-rs-preview/cute-rs/README.md b/cuda-oxide/cute-rs-preview/cute-rs/README.md new file mode 100644 index 0000000000..4e33bdedf2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/README.md @@ -0,0 +1,49 @@ +# cute-rs + +CuTe-style tensor computation for cuda-oxide kernels. Typed views, layouts, +partitions, register tiles, copies, reductions and MMA operations retain the +storage and participation contracts needed by the device compiler. + +| Domain | Contents | +| --- | --- | +| `tensor` | Core tensor views and register tiles; `partition` for strided and cyclic views | +| `copy` | `cooperative`, `tiled`, `vector` and `tma` transfers | +| `mma` | Common FP16 fragments and loads; distinct `bf16` and `block_scaled` families | +| `attention` | Checked host `metadata` and device-available `sm100` operations | +| `epilogue` | `sm120` output staging | + +The architecture ownership types in `sm100`, readiness protocols in `pipeline`, +block-scaled storage in `block_scaled`, and other focused root modules retain +their separate responsibilities. Each operation family documents its own +layout, lifetime and collective-participation requirements. + +The crate is `no_std`. The opt-in `host` feature exposes descriptor encoding, +attention metadata, decoder metadata and FFN metadata. Device SM100 attention +is available without that feature. Compile-time checking of a view does not +establish convergence, allocation readiness or safe asynchronous reuse. + +## Import paths + +Every module has one path, its defining path. Root item exports such as +`cute_rs::TensorMap` name the same items as their defining modules, here +`cute_rs::copy::tma::TensorMap`; they do not wrap or duplicate +implementations. Two parent modules re-export grouped items the same way: +`attention::{AttentionError, AttentionShape, AttentionWindow, CacheShape}` are +the `attention::metadata` items (behind the `host` feature), and the `epilogue` +items are the `epilogue::sm120` items. + +## Checks + +Run from the `cuda-oxide/` workspace with its pinned toolchain: + +```sh +cargo test -p cute-rs --no-default-features +cargo test -p cute-rs --features host +cargo clippy -p cute-rs --all-targets --features host -- -D warnings +RUSTDOCFLAGS='-D warnings' cargo doc -p cute-rs --features host --no-deps +``` + +Host checks cover layout invariants, representations, typed imports and +compile-fail contracts. Executing device operations and validating their +generated instructions additionally requires the cuda-oxide device compiler +and the appropriate GPU. diff --git a/cuda-oxide/cute-rs-preview/cute-rs/examples/global_store_compile.rs b/cuda-oxide/cute-rs-preview/cute-rs/examples/global_store_compile.rs new file mode 100644 index 0000000000..feb04b3fa9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/examples/global_store_compile.rs @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Real-Rust compile-only coverage of the aligned global-store boundary. +//! +//! From `cuda-oxide/` (no GPU execution): +//! ```text +//! CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide build --arch sm_100a \ +//! --device-codegen-crate global_store_compile -- \ +//! -p cute-rs --example global_store_compile --release +//! ``` + +#![feature(f16)] +use cuda_device::{kernel, thread}; +use cute_rs::{GlobalStoreTile, store_global_no_allocate}; + +/// # Safety +/// `out` must be 32-byte aligned and writable for `8 * blockDim.x()` elements. +/// Launch one one-dimensional block so threads write disjoint ranges. +#[kernel] +pub unsafe fn store_global_f32(out: *mut f32, bits: u64) { + let x = f32::from_bits(bits as u32); + let values = GlobalStoreTile([ + x, + x + 1.0, + x + 2.0, + x + 3.0, + x + 4.0, + x + 5.0, + x + 6.0, + x + 7.0, + ]); + unsafe { + store_global_no_allocate(out.add(thread::threadIdx_x() as usize * 8), &values); + } +} + +/// # Safety +/// `out` must be 32-byte aligned and writable for `8 * blockDim.x()` elements. +/// Launch one one-dimensional block so threads write disjoint ranges. +#[kernel] +pub unsafe fn store_global_u32(out: *mut u32, bits: u64) { + let x = bits as u32; + let values = GlobalStoreTile([x, x ^ 1, x ^ 2, x ^ 3, x ^ 4, x ^ 5, x ^ 6, x ^ 7]); + unsafe { + store_global_no_allocate(out.add(thread::threadIdx_x() as usize * 8), &values); + } +} + +/// # Safety +/// `out` must be 32-byte aligned and writable for `16 * blockDim.x()` elements. +/// Launch one one-dimensional block so threads write disjoint ranges. +#[kernel] +pub unsafe fn store_global_f16(out: *mut f16, bits: u64) { + let x = f32::from_bits(bits as u32); + let values = GlobalStoreTile([x as f16; 16]); + unsafe { + store_global_no_allocate(out.add(thread::threadIdx_x() as usize * 16), &values); + } +} + +/// Provide a host entry point so this device-store compilation probe is a valid example target. +fn main() {} // Kernels are compiled, never launched. diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/arch.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/arch.rs new file mode 100644 index 0000000000..29f869d64d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/arch.rs @@ -0,0 +1,787 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit warp-level compiler facts built from supported CUDA operations. +//! +//! This module carries the `cute.arch` names of CuTeDSL over the PTX-named +//! ctm-rs layer, mirroring `cute.arch` over the NVVM primitives in CuTeDSL. +//! One primitives-layer name lives here as well, [`fence_sync_restrict`], +//! because its order is a type and the DSL has no `cute.arch` form of it. +//! Every forwarder is an `#[inline(always)]` call to the existing ctm-rs +//! primitive: it adds no instruction, no storage and no implicit election or +//! synchronization, and cute-rs never reaches past ctm-rs to the raw device +//! operations. Argument order follows `cute.arch`, +//! the value first and the participation mask last; Rust has no default +//! arguments, so the mask is spelled out (`u32::MAX` for the full warp). +//! Safety follows the primitive: a forwarder is `unsafe` when the function it +//! calls is, under the same contract. + +use ctm_rs::mbarrier::{BarrierToken, ManagedBarrier, Ready, TmaBarrier}; +use ctm_rs::memory::{ClusterSharedAddress, SharedAddress}; + +/// Return a warp-uniform value in a form that can favor uniform-register use. +/// +/// This is CuTeDSL's `cute.arch.make_warp_uniform`: a full-warp shuffle +/// broadcasts lane zero's value. Pass the returned value to subsequent +/// calculations. The compiler may then keep those calculations in uniform +/// registers; this is an optimization opportunity, not a register-allocation +/// guarantee. +/// +/// # Safety +/// All 32 lanes must be active and execute this call convergently with equal +/// `value` arguments. Do not call from an elected-lane or partially active +/// branch. This is an identity operation only under that uniformity promise; +/// it does not check or repair a divergent value. +#[inline(always)] +pub unsafe fn make_warp_uniform(value: u32) -> u32 { + // SAFETY: the caller supplies the full-warp convergence contract above. + unsafe { ctm_rs::warp::shuffle_sync(u32::MAX, value, 0) } +} + +/// Sum one FP32 contribution from each of the 32 lanes of a warp. +/// +/// This is CuTeDSL's `cute.arch.warp_reduction_sum` for FP32. Every lane +/// receives the sum. Uses the existing CUDA butterfly reduction, with FP32 +/// additions in XOR-offset order 16, 8, 4, 2, 1. The result need not equal a +/// sequential sum. No shared memory or block barrier is required. Inactive +/// logical elements must contribute zero while their lanes still participate. +/// This is a device operation, not a host reduction. +/// +/// # Safety +/// All 32 lanes must be active and execute this call convergently. The block +/// must contain complete warps. Rust's type system cannot prove convergence; +/// do not invoke this in a lane-dependent branch or after a partial-warp exit. +#[must_use] +#[inline(always)] +pub unsafe fn warp_reduction_sum_f32(value: f32) -> f32 { + // SAFETY: every lane participates under this function's caller contract. + unsafe { ctm_rs::warp::reduce_sum_f32(value) } +} + +/// Approximate base-two exponential, preserving the device instruction semantics. +#[inline(always)] +pub fn exp2_f32(value: f32) -> f32 { + ctm_rs::float::ex2_approx_f32(value) +} + +/// Approximate reciprocal with the device instruction's flush-to-zero behavior. +#[inline(always)] +pub fn reciprocal_ftz_f32(value: f32) -> f32 { + ctm_rs::float::rcp_approx_ftz_f32(value) +} + +/// Approximate reciprocal square root without changing surrounding arithmetic. +#[inline(always)] +pub fn rsqrt_f32(value: f32) -> f32 { + ctm_rs::float::rsqrt_approx_f32(value) +} + +/// Approximate sine; callers retain responsibility for argument-range reduction. +#[inline(always)] +pub fn sin_f32(value: f32) -> f32 { + ctm_rs::float::sin_approx_f32(value) +} + +/// Maximum across one full warp, using ordered comparisons and XOR offsets. +/// +/// This is CuTeDSL's `cute.arch.warp_reduction_max` for FP32. The comparison +/// selects the right value when `left > right` is false; it deliberately +/// retains this behavior for NaNs and signed zero, so the tree stays here +/// instead of forwarding to a primitive with a different NaN rule. +/// +/// # Safety +/// All 32 lanes must execute this operation convergently. +#[inline(always)] +pub unsafe fn warp_reduction_max_f32(mut value: f32) -> f32 { + let mut offset = 16; + while offset != 0 { + let other = ctm_rs::warp::shuffle_xor_f32(value, offset); + value = if value > other { value } else { other }; + offset /= 2; + } + value +} + +/// Sum independent groups of four neighboring lanes, with XOR offsets 1 then 2. +/// +/// This tree is intentionally different from the full-warp reduction tree. +/// +/// # Safety +/// All 32 lanes must execute this operation convergently. +#[inline(always)] +pub unsafe fn quad_sum_f32(mut value: f32) -> f32 { + value += ctm_rs::warp::shuffle_xor_f32(value, 1); + value += ctm_rs::warp::shuffle_xor_f32(value, 2); + value +} + +/// A named CTA barrier with a statically fixed participant count. +/// +/// This is CuTeDSL's `pipeline.NamedBarrier(barrier_id, num_threads)` with +/// both values as const generics; [`NamedBarrier::sync`] is its +/// `arrive_and_wait`. This marker describes the synchronization operation, +/// not an owned hardware resource. Callers must keep barrier phases and their +/// participating threads consistent; distinct logical barriers must use +/// distinct IDs. +pub struct NamedBarrier; + +impl NamedBarrier { + /// Wait for all participating threads and make their shared writes visible. + /// + /// # Safety + /// Exactly `THREADS` threads in the same CTA must execute matching barrier + /// phases convergently. They must form complete warps, and no other logical + /// barrier may concurrently reuse `ID`. No participant may exit early. + #[inline(always)] + pub unsafe fn sync() { + const { + assert!(ID < 16 && THREADS > 0 && THREADS <= 1024 && THREADS.is_multiple_of(32)); + } + // SAFETY: the caller supplies the fixed complete-warp participant set. + unsafe { ctm_rs::barrier::barrier_cta_sync(ID, THREADS) }; + } +} + +// Thread, block and grid coordinates. + +/// Return the thread index within the CTA as `(x, y, z)`. +/// +/// This is `cute.arch.thread_idx()`; it reads `%tid.x`, `%tid.y` and `%tid.z` +/// through [`ctm_rs::thread`]. +#[must_use] +#[inline(always)] +pub fn thread_idx() -> (u32, u32, u32) { + ( + ctm_rs::thread::threadIdx_x(), + ctm_rs::thread::threadIdx_y(), + ctm_rs::thread::threadIdx_z(), + ) +} + +/// Return the CTA index within the grid as `(x, y, z)`. +/// +/// This is `cute.arch.block_idx()`; it reads `%ctaid.x`, `%ctaid.y` and +/// `%ctaid.z` through [`ctm_rs::thread`]. +#[must_use] +#[inline(always)] +pub fn block_idx() -> (u32, u32, u32) { + ( + ctm_rs::thread::blockIdx_x(), + ctm_rs::thread::blockIdx_y(), + ctm_rs::thread::blockIdx_z(), + ) +} + +/// Return the number of CTAs in each grid dimension as `(x, y, z)`. +/// +/// This is `cute.arch.grid_dim()`; it reads `%nctaid.x`, `%nctaid.y` and +/// `%nctaid.z` through [`ctm_rs::thread`]. +#[must_use] +#[inline(always)] +pub fn grid_dim() -> (u32, u32, u32) { + ( + ctm_rs::thread::gridDim_x(), + ctm_rs::thread::gridDim_y(), + ctm_rs::thread::gridDim_z(), + ) +} + +/// Return the number of threads in each CTA dimension as `(x, y, z)`. +/// +/// This is `cute.arch.block_dim()`; it reads `%ntid.x`, `%ntid.y` and +/// `%ntid.z` through [`ctm_rs::thread`]. +#[must_use] +#[inline(always)] +pub fn block_dim() -> (u32, u32, u32) { + ( + ctm_rs::thread::blockDim_x(), + ctm_rs::thread::blockDim_y(), + ctm_rs::thread::blockDim_z(), + ) +} + +/// Return the logical warp index within the CTA, marked warp-uniform. +/// +/// This is `cute.arch.warp_idx()`: the block-linear thread index +/// (`x + y * dim_x + z * dim_x * dim_y`) divided by 32, passed through +/// [`make_warp_uniform`]. For a one-dimensional block this is +/// `threadIdx.x / 32`. Hardware warps are formed from consecutive linear +/// thread indices, so the quotient is the same in every lane of a warp. +/// +/// # Safety +/// All 32 lanes must be active and execute this call convergently, as +/// [`make_warp_uniform`] requires. Do not call from an elected-lane or +/// partially active branch. +#[must_use] +#[inline(always)] +pub unsafe fn warp_idx() -> u32 { + let (x, y, z) = thread_idx(); + let dim_x = ctm_rs::thread::blockDim_x(); + let dim_y = ctm_rs::thread::blockDim_y(); + let linear = x + y * dim_x + z * dim_x * dim_y; + // SAFETY: the caller supplies the full-warp convergence contract above. + unsafe { make_warp_uniform(linear / 32) } +} + +/// Return this thread's lane index within its warp (`%laneid`). +/// +/// This is `cute.arch.lane_idx()` over [`ctm_rs::warp::lane_id`]. +#[must_use] +#[inline(always)] +pub fn lane_idx() -> u32 { + ctm_rs::warp::lane_id() +} + +// Warp election and synchronization. + +/// Elect one lane of the full warp and return whether this lane won. +/// +/// This is the predicate form of `cute.arch.elect_one()`, over +/// [`ctm_rs::warp::is_elected_sync`] with the full 32-lane mask (`elect.sync`, +/// sm_90+). All 32 lanes must be converged at the call. Use +/// [`ctm_rs::warp::is_elected_sync`] directly for a partial mask. +#[must_use] +#[inline(always)] +pub fn elect_one_sync() -> bool { + ctm_rs::warp::is_elected_sync(u32::MAX) +} + +/// Synchronize the warp lanes named in `mask` (`bar.warp.sync`). +/// +/// This is `cute.arch.sync_warp(mask)` over [`ctm_rs::warp::bar_warp_sync`]; pass +/// `u32::MAX` for the full warp. +/// +/// # Safety +/// The caller's lane must be included in `mask`. Every lane named in `mask` +/// must reach this call with the same mask. All accesses ordered by the barrier +/// must satisfy their allocation contracts. +#[inline(always)] +pub unsafe fn sync_warp(mask: u32) { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::bar_warp_sync(mask) } +} + +/// Synchronize `number_of_threads` CTA threads at the named barrier `barrier_id`. +/// +/// This is `cute.arch.barrier(barrier_id, number_of_threads)` over +/// [`ctm_rs::barrier::barrier_cta_sync`] (`barrier.cta.sync`). [`NamedBarrier`] +/// fixes both values at compile time; the aligned form stays in ctm-rs as +/// `barrier_cta_sync_aligned`. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +#[inline(always)] +pub unsafe fn barrier(barrier_id: u32, number_of_threads: u32) { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::barrier::barrier_cta_sync(barrier_id, number_of_threads) } +} + +/// Arrive at the named barrier `barrier_id` without waiting. +/// +/// This is `cute.arch.barrier_arrive(barrier_id, number_of_threads)` over +/// [`ctm_rs::barrier::barrier_cta_arrive`] (`barrier.cta.arrive`, the form +/// that admits a subset of the CTA's warps); the aligned form stays in ctm-rs +/// as `barrier_cta_arrive_aligned`. +/// +/// # Safety +/// Participating threads must use a compatible barrier ID and expected count. +#[inline(always)] +pub unsafe fn barrier_arrive(barrier_id: u32, number_of_threads: u32) { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::barrier::barrier_cta_arrive(barrier_id, number_of_threads) } +} + +// Mbarrier transactions. + +/// Publish preceding mbarrier initialization at cluster scope. +/// +/// This is `cute.arch.mbarrier_init_fence()` over +/// [`ctm_rs::fence::fence_mbarrier_init`] +/// (`fence.mbarrier_init.release.cluster`). +/// +/// # Safety +/// The caller must satisfy the barrier, scope, and memory-ordering contract of +/// [`ctm_rs::fence::fence_mbarrier_init`]. +#[inline(always)] +pub unsafe fn mbarrier_init_fence() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::fence::fence_mbarrier_init() } +} + +/// Arrive once on a TMA barrier and add the expected transaction bytes. +/// +/// This is `cute.arch.mbarrier_arrive_and_expect_tx(mbar, bytes)` over the typed +/// [`ctm_rs::mbarrier::arrive_expect_tx`]; the arrival decrement is one and the +/// hardware phase token is returned. +/// +/// # Safety +/// The issuing participant must be expected, and `bytes` must match the +/// asynchronous transfers completing against this barrier phase. +#[inline(always)] +pub unsafe fn mbarrier_arrive_and_expect_tx( + bar: &ManagedBarrier, + bytes: u32, +) -> BarrierToken { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::mbarrier::arrive_expect_tx(bar, bytes) } +} + +// Sync-restricted fences. + +/// Seals [`MemOrder`] to the orders the intrinsic catalog can emit. +mod sealed { + /// Implemented by the order markers of this module only. + pub trait MemOrder {} +} + +/// Memory order of a sync-restricted fence, fixed at compile time. +/// +/// This is the `order` argument of CuTeDSL's `fence_sync_restrict(order)`, +/// a `MemOrder` string enum (`"acquire"` or `"release"`) in the DSL and a +/// marker type here, so the instruction is chosen when the kernel compiles. +/// The trait is sealed and [`Acquire`] is its only implementor: the intrinsic +/// catalog has no `fence.release.sync_restrict::shared::cta.cluster` entry +/// yet, so the release order is not offered until that entry is added to the +/// catalog and bound in cuda-device. +pub trait MemOrder: sealed::MemOrder { + /// Issue the sync-restricted fence in this order at cluster scope. + /// + /// # Safety + /// Pair this with a compatible cluster-scoped synchronization protocol. + /// This does not substitute for a proxy fence or asynchronous completion + /// wait. + unsafe fn fence_sync_restrict(); +} + +/// The acquire order, `MemOrder.ACQUIRE` in CuTeDSL. +/// +/// Its fence is `fence.acquire.sync_restrict::shared::cluster.cluster`: it +/// acquires cluster-shared memory at cluster scope and leaves global memory +/// unordered. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Acquire; + +impl sealed::MemOrder for Acquire {} + +impl MemOrder for Acquire { + #[inline(always)] + unsafe fn fence_sync_restrict() { + // SAFETY: forwarded under the caller contract on the trait method. + unsafe { ctm_rs::fence::fence_sync_restrict_acquire() } + } +} + +/// Fence cluster-shared memory at cluster scope in `Order`, excluding global +/// memory. +/// +/// This is CuTeDSL's `fence_sync_restrict(order)` with the order as a type +/// parameter: `fence_sync_restrict::()` is +/// `fence.acquire.sync_restrict::shared::cluster.cluster` over +/// [`ctm_rs::fence::fence_sync_restrict_acquire`]. The DSL defines the +/// function in its primitives layer (`nvvm.fence_sync_restrict`) rather than +/// in `cute.arch`; cute-rs keeps it beside the other fences here. +/// +/// # Safety +/// Pair this with a compatible cluster-scoped synchronization protocol. This +/// does not substitute for a proxy fence or asynchronous completion wait. +#[inline(always)] +pub unsafe fn fence_sync_restrict() { + // SAFETY: forwarded under the caller contract above. + unsafe { Order::fence_sync_restrict() } +} + +// Cluster membership and barriers. + +/// Arrive at the cluster barrier with release ordering. +/// +/// This is `cute.arch.cluster_arrive()` over +/// [`ctm_rs::cluster::barrier_cluster_arrive`]; the aligned form stays in +/// ctm-rs as `barrier_cluster_arrive_aligned`. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. +#[inline(always)] +pub unsafe fn cluster_arrive() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::cluster::barrier_cluster_arrive() } +} + +/// Arrive at the cluster barrier without a release guarantee. +/// +/// This is `cute.arch.cluster_arrive_relaxed()` over +/// [`ctm_rs::cluster::barrier_cluster_arrive_relaxed`]; the aligned form stays +/// in ctm-rs as `barrier_cluster_arrive_relaxed_aligned`. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute the matching wait. Publish earlier memory accesses with the +/// required cluster-scope fence first. +#[inline(always)] +pub unsafe fn cluster_arrive_relaxed() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::cluster::barrier_cluster_arrive_relaxed() } +} + +/// Wait at the cluster barrier with acquire ordering. +/// +/// This is `cute.arch.cluster_wait()` over +/// [`ctm_rs::cluster::barrier_cluster_wait`]; the aligned form stays in ctm-rs +/// as `barrier_cluster_wait_aligned`. +/// +/// # Safety +/// Each non-exited cluster thread must arrive exactly once before completion, +/// then execute this matching wait. +#[inline(always)] +pub unsafe fn cluster_wait() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::cluster::barrier_cluster_wait() } +} + +/// Return the linear rank of this CTA within its cluster. +/// +/// This is `cute.arch.block_idx_in_cluster()` over +/// [`ctm_rs::cluster::cluster_ctarank`], the value `%cluster_ctarank` holds, in +/// `0..cluster_size`. +#[must_use] +#[inline(always)] +pub fn block_idx_in_cluster() -> u32 { + ctm_rs::cluster::cluster_ctarank() +} + +/// Return this CTA's index within its cluster as `(x, y, z)`. +/// +/// This is `cute.arch.block_in_cluster_idx()`; it reads `%cluster_ctaid.x`, +/// `%cluster_ctaid.y` and `%cluster_ctaid.z` through [`ctm_rs::cluster`]. +/// [`block_idx_in_cluster`] is the linear rank of the same CTA +/// (`%cluster_ctarank`); [`cluster_idx`] is the cluster's own `(x, y, z)` +/// within the grid. +#[must_use] +#[inline(always)] +pub fn block_in_cluster_idx() -> (u32, u32, u32) { + ( + ctm_rs::cluster::cluster_ctaidX(), + ctm_rs::cluster::cluster_ctaidY(), + ctm_rs::cluster::cluster_ctaidZ(), + ) +} + +/// Return the number of CTAs in each cluster dimension as `(x, y, z)`. +/// +/// This is `cute.arch.block_in_cluster_dim()`; it reads `%cluster_nctaid.x`, +/// `%cluster_nctaid.y` and `%cluster_nctaid.z` through [`ctm_rs::cluster`]. +/// The compiler's own `cluster_dim` vocabulary, the `#[cluster_launch]` +/// attribute and the intrinsic family behind it, means this value, CTAs per +/// cluster; CuTeDSL's `cute.arch.cluster_dim()` is [`cluster_dim`], the +/// number of clusters in the grid. +#[must_use] +#[inline(always)] +pub fn block_in_cluster_dim() -> (u32, u32, u32) { + ( + ctm_rs::cluster::cluster_nctaidX(), + ctm_rs::cluster::cluster_nctaidY(), + ctm_rs::cluster::cluster_nctaidZ(), + ) +} + +/// Return this CTA's cluster index within the grid as `(x, y, z)`. +/// +/// This is `cute.arch.cluster_idx()`; it reads `%clusterid.x`, `%clusterid.y` +/// and `%clusterid.z` through [`ctm_rs::cluster`]. The linear form of the +/// same index stays `ctm_rs::cluster::cluster_idx()`, and +/// [`block_in_cluster_idx`] is the CTA's index within that cluster. +#[must_use] +#[inline(always)] +pub fn cluster_idx() -> (u32, u32, u32) { + ( + ctm_rs::cluster::cluster_idx_x(), + ctm_rs::cluster::cluster_idx_y(), + ctm_rs::cluster::cluster_idx_z(), + ) +} + +/// Return the number of clusters in each grid dimension as `(x, y, z)`. +/// +/// This is `cute.arch.cluster_dim()`; it reads `%nclusterid.x`, +/// `%nclusterid.y` and `%nclusterid.z` through [`ctm_rs::cluster`]. The +/// compiler's own `cluster_dim` vocabulary, the `#[cluster_launch]` attribute +/// of cuda-macros and the cuda-intrinsics-gen family behind it, means CTAs +/// per cluster, which is [`block_in_cluster_dim`] here; the DSL's +/// `cluster_dim` means clusters in the grid. +#[must_use] +#[inline(always)] +pub fn cluster_dim() -> (u32, u32, u32) { + ( + ctm_rs::cluster::cluster_grid_dim_x(), + ctm_rs::cluster::cluster_grid_dim_y(), + ctm_rs::cluster::cluster_grid_dim_z(), + ) +} + +/// Map local shared storage to the same offset in cluster peer `rank`. +/// +/// This is `cute.arch.map_dsmem_ptr(ptr, rank)` over +/// [`ctm_rs::cluster::mapa`]. +/// +/// # Safety +/// `rank` must name an active CTA in this cluster with matching live storage. +/// Remote initialization, readiness and lifetime must be coordinated explicitly. +#[inline(always)] +pub unsafe fn map_dsmem_ptr( + address: SharedAddress<'_, T>, + rank: u32, +) -> ClusterSharedAddress<'_, T> { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::cluster::mapa(address, rank) } +} + +// Register budget. + +/// Decrease the warpgroup's register limit to the immediate `N`. +/// +/// This is `cute.arch.setmaxregister_decrease(n)` over +/// [`ctm_rs::thread::setmaxnreg_dec`]; `N` must be a multiple of eight in +/// `24..=256`. +/// +/// # Safety +/// Every thread in the warpgroup must execute the same operation and count +/// on supported hardware. The current register limit must be at least `N`. +/// Synchronize the warpgroup before another register-limit adjustment. +/// Subsequent code must respect the selected limit. +#[inline(always)] +pub unsafe fn setmaxregister_decrease() +where + (): ctm_rs::thread::ValidRegisterCount, +{ + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::thread::setmaxnreg_dec::() } +} + +/// Increase the warpgroup's register limit to the immediate `N`. +/// +/// This is `cute.arch.setmaxregister_increase(n)` over +/// [`ctm_rs::thread::setmaxnreg_inc`]; `N` must be a multiple of eight in +/// `24..=256`. +/// +/// # Safety +/// Every thread in the warpgroup must execute the same operation and count +/// on supported hardware. The current register limit must be at most `N`. +/// Synchronize the warpgroup before another register-limit adjustment. +/// Coordinate releases by other warpgroups to avoid waiting indefinitely +/// for unavailable registers. +#[inline(always)] +pub unsafe fn setmaxregister_increase() +where + (): ctm_rs::thread::ValidRegisterCount, +{ + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::thread::setmaxnreg_inc::() } +} + +// Bulk-copy groups. + +/// Commit this thread's pending bulk copies as one group; does not wait. +/// +/// This is `cute.arch.cp_async_bulk_commit_group()` over +/// [`ctm_rs::tma::cp_async_bulk_commit_group`]. +#[inline(always)] +pub fn cp_async_bulk_commit_group() { + ctm_rs::tma::cp_async_bulk_commit_group() +} + +/// Wait until at most `N` committed bulk groups still read their shared sources. +/// +/// This is `cute.arch.cp_async_bulk_wait_group(N, read=True)` over +/// [`ctm_rs::tma::cp_async_bulk_wait_group_read`]; `N` is an immediate in `0..=7`. It +/// is not proof that the groups' global writes are visible to other threads. +#[inline(always)] +pub fn cp_async_bulk_wait_group_read() +where + (): ctm_rs::tma::ValidWaitCount, +{ + ctm_rs::tma::cp_async_bulk_wait_group_read::() +} + +// Tensor memory completion. + +/// Complete this thread's previously issued SM100 TMEM loads. +/// +/// This is `cute.arch.fence_view_async_tmem_load()` over +/// [`ctm_rs::tcgen05::tmem::load_wait`] (`tcgen05.wait::ld`). No CTA barrier, +/// producer wait or memory/proxy fence is inserted. +/// +/// # Safety +/// All 32 warp lanes must execute convergently. The loads must target live, +/// valid TMEM allocations with correctly ordered producers. All tcgen05 +/// operations must use one CTA-group mode. +#[inline(always)] +pub unsafe fn fence_view_async_tmem_load() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::tcgen05::tmem::load_wait() } +} + +/// Complete this thread's previously issued SM100 TMEM stores. +/// +/// This is `cute.arch.fence_view_async_tmem_store()` over +/// [`ctm_rs::tcgen05::tmem::store_wait`] (`tcgen05.wait::st`). No CTA barrier, +/// consumer synchronization or memory/proxy fence is inserted. +/// +/// # Safety +/// All 32 warp lanes must execute convergently. The stores must target live, +/// valid TMEM allocations with no conflicting accesses. All tcgen05 operations +/// must use one CTA-group mode. +#[inline(always)] +pub unsafe fn fence_view_async_tmem_store() { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::tcgen05::tmem::store_wait() } +} + +// Warp votes. + +/// Ballot across the lanes in `mask`: bit `k` is set when lane `k` participates +/// and its `predicate` is true. +/// +/// This is `cute.arch.vote_ballot_sync(pred, mask)` over +/// [`ctm_rs::warp::vote_sync_ballot`] (`vote.sync.ballot`). +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn vote_ballot_sync(predicate: bool, mask: u32) -> u32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::vote_sync_ballot(mask, predicate) } +} + +/// Return whether `predicate` holds for at least one lane in `mask`. +/// +/// This is `cute.arch.vote_any_sync(pred, mask)` over +/// [`ctm_rs::warp::vote_sync_any`] (`vote.sync.any`). +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn vote_any_sync(predicate: bool, mask: u32) -> bool { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::vote_sync_any(mask, predicate) } +} + +/// Return whether `predicate` holds for every lane in `mask`. +/// +/// This is `cute.arch.vote_all_sync(pred, mask)` over +/// [`ctm_rs::warp::vote_sync_all`] (`vote.sync.all`). +/// +/// # Safety +/// Every lane named in `mask` must execute this vote with the same mask, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn vote_all_sync(predicate: bool, mask: u32) -> bool { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::vote_sync_all(mask, predicate) } +} + +/// Return the mask of the warp's currently active lanes (`activemask`). +/// +/// This is `cute.arch.activemask()` over [`ctm_rs::warp::activemask`]. Bit `k` +/// is set when lane `k` has neither exited nor diverged away at this point of +/// execution. +#[must_use] +#[inline(always)] +pub fn activemask() -> u32 { + ctm_rs::warp::activemask() +} + +// Warp shuffles. The DSL's one generic function becomes a `u32` form and an +// `_f32` form, the two widths the preview kernels exchange. + +/// Exchange a `u32` with the lane whose index is this lane's XOR `offset`. +/// +/// This is `cute.arch.shuffle_sync_bfly(value, offset, mask)` over +/// [`ctm_rs::warp::shfl_sync_bfly`] (`shfl.sync.bfly.b32`); the clamp is the +/// full warp. No memory synchronization is inserted. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. `offset` must select a source lane within the warp. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_bfly(value: u32, offset: u32, mask: u32) -> u32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_bfly(mask, value, offset) } +} + +/// FP32 form of [`shuffle_sync_bfly`] over [`ctm_rs::warp::shfl_sync_bfly_f32`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. `offset` must select a source lane within the warp. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_bfly_f32(value: f32, offset: u32, mask: u32) -> f32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_bfly_f32(mask, value, offset) } +} + +/// Read a `u32` from the lane `offset` below this lane; lanes below `offset` +/// keep their own value. +/// +/// This is `cute.arch.shuffle_sync_up(value, offset, mask)` over +/// [`ctm_rs::warp::shfl_sync_up`] (`shfl.sync.up.b32`). No memory +/// synchronization is inserted. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_up(value: u32, offset: u32, mask: u32) -> u32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_up(mask, value, offset) } +} + +/// FP32 form of [`shuffle_sync_up`] over [`ctm_rs::warp::shfl_sync_up_f32`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_up_f32(value: f32, offset: u32, mask: u32) -> f32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_up_f32(mask, value, offset) } +} + +/// Read a `u32` from the lane `offset` above this lane; the top `offset` lanes +/// keep their own value. +/// +/// This is `cute.arch.shuffle_sync_down(value, offset, mask)` over +/// [`ctm_rs::warp::shfl_sync_down`] (`shfl.sync.down.b32`). No memory +/// synchronization is inserted. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_down(value: u32, offset: u32, mask: u32) -> u32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_down(mask, value, offset) } +} + +/// FP32 form of [`shuffle_sync_down`] over [`ctm_rs::warp::shfl_sync_down_f32`]. +/// +/// # Safety +/// Every lane named in `mask` must execute with the same mask and `offset`, +/// including the caller. +#[must_use] +#[inline(always)] +pub unsafe fn shuffle_sync_down_f32(value: f32, offset: u32, mask: u32) -> f32 { + // SAFETY: forwarded under the caller contract above. + unsafe { ctm_rs::warp::shfl_sync_down_f32(mask, value, offset) } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/metadata.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/metadata.rs new file mode 100644 index 0000000000..ed60844dfd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/metadata.rs @@ -0,0 +1,345 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Checked host metadata for BF16 attention and contiguous key/value caches. +//! +//! Query heads, KV heads and head width remain separate dimensions. In +//! particular, a model's hidden width need not equal its query projection width. +//! Element counts describe dense `[token, head, dimension]` storage, with +//! dimension coordinates contiguous. Cache capacities and append ranges count +//! tokens, while allocation lengths count BF16 elements in each of K and V. +//! +//! These types neither allocate buffers nor track initialized data. A validated +//! shape does not establish allocation size, context, readiness or ownership. +//! The cache owner must check those properties, preserve initialized contents, +//! and advance its cursor only after a successful append. Calling +//! [`CacheShape::append_range`] does not reserve or initialize any storage. +//! +//! ``` +//! use cute_rs::attention::{AttentionShape, CacheShape}; +//! let attention = AttentionShape::<16, 8, 128>::new().unwrap(); +//! assert_eq!(attention.query_width(), 2048); +//! assert_eq!(attention.kv_width(), 1024); +//! let cache = CacheShape::<8, 128>::new(65).unwrap(); +//! assert_eq!(cache.append_range(17, 3).unwrap(), 17..20); +//! ``` + +use core::{fmt, mem::size_of, ops::Range}; + +use crate::numeric::BFloat16; + +/// A rejected attention dimension, token window, BF16 extent or cache append. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AttentionError { + /// A head count or head width is zero. + ZeroDimension, + /// Query heads do not form complete, equally sized groups of KV heads. + InvalidHeadGrouping, + /// Token/element arithmetic overflowed or a BF16 extent exceeds `isize::MAX` bytes. + ExtentOverflow, + /// The supplied initialized-token cursor lies beyond cache capacity. + InvalidCursor, + /// An append would exceed the cache's token capacity. + CapacityExceeded, + /// A nonempty query sequence has no keys to attend to. + MissingKeys, + /// A causal query range lies outside the available key sequence. + InvalidWindow, + /// A query-row index lies outside the window's query sequence. + InvalidQueryRow, +} + +impl fmt::Display for AttentionError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::ZeroDimension => "attention head counts and head width must be positive", + Self::InvalidHeadGrouping => "query head count must be a multiple of KV head count", + Self::ExtentOverflow => { + "attention extent overflowed or exceeds the BF16 addressable range" + } + Self::InvalidCursor => "initialized-token cursor exceeds cache capacity", + Self::CapacityExceeded => "append exceeds cache token capacity", + Self::MissingKeys => "nonempty attention queries require at least one key", + Self::InvalidWindow => "causal query range exceeds the available key sequence", + Self::InvalidQueryRow => "query row exceeds the attention window", + }) + } +} + +impl core::error::Error for AttentionError {} + +/// Validate the complete byte extent of one BF16 allocation or slice. +fn check_extent(elements: usize) -> Result { + match elements.checked_mul(size_of::()) { + Some(bytes) if bytes <= isize::MAX as usize => Ok(elements), + _ => Err(AttentionError::ExtentOverflow), + } +} + +/// Validate positive head dimensions and the BF16 extent of one token. +fn checked_width(heads: usize, head_dim: usize) -> Result { + if heads == 0 || head_dim == 0 { + return Err(AttentionError::ZeroDimension); + } + check_extent( + heads + .checked_mul(head_dim) + .ok_or(AttentionError::ExtentOverflow)?, + ) +} + +/// Validate token multiplication and the full BF16 slice byte extent. +fn checked_len(tokens: usize, width: usize) -> Result { + check_extent( + tokens + .checked_mul(width) + .ok_or(AttentionError::ExtentOverflow)?, + ) +} + +/// Validated query/KV head geometry, independent of a model's hidden width. +/// +/// `QH` must be a positive multiple of `KVH`, and `D` must be positive. No +/// architecture-specific head-width or tile-size restriction is imposed here. +/// Kernels must check any stronger specialization requirements at their boundary. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AttentionShape { + /// BF16 elements per token in a query or attention output tensor. + query_width: usize, + /// BF16 elements per token in a key or value tensor. + kv_width: usize, +} + +impl AttentionShape { + /// Check grouped-query geometry and both per-token BF16 extents. + pub fn new() -> Result { + if QH == 0 || KVH == 0 || D == 0 { + return Err(AttentionError::ZeroDimension); + } + if !QH.is_multiple_of(KVH) { + return Err(AttentionError::InvalidHeadGrouping); + } + Ok(Self { + query_width: checked_width(QH, D)?, + kv_width: checked_width(KVH, D)?, + }) + } + + /// Number of query heads. + #[must_use] + pub const fn query_heads(self) -> usize { + QH + } + + /// Number of key/value heads. + #[must_use] + pub const fn kv_heads(self) -> usize { + KVH + } + + /// Coordinates in each query, key or value head. + #[must_use] + pub const fn head_dim(self) -> usize { + D + } + + /// BF16 elements per token in a query or attention output tensor. + #[must_use] + pub const fn query_width(self) -> usize { + self.query_width + } + + /// BF16 elements per token in a key or value tensor. + #[must_use] + pub const fn kv_width(self) -> usize { + self.kv_width + } + + /// Number of consecutive query heads sharing each KV head. + #[must_use] + pub const fn groups(self) -> usize { + QH / KVH + } + + /// BF16 elements in a dense query/output tensor, checking its byte extent. + /// An empty token sequence needs zero elements. + pub fn query_len(self, tokens: usize) -> Result { + checked_len(tokens, self.query_width) + } + + /// BF16 elements in one dense key/value tensor, checking its byte extent. + /// An empty token sequence needs zero elements. + pub fn kv_len(self, tokens: usize) -> Result { + checked_len(tokens, self.kv_width) + } +} + +/// Checked query/key token counts and an optional absolute causal query offset. +/// +/// A causal row at local query index `r` can attend to the key prefix ending at +/// absolute position `query_start + r`, inclusive. Noncausal rows see all keys. +/// Empty query sequences are permitted, but every nonempty query sequence must +/// have at least one key. A causal query range must fit entirely inside the key +/// sequence, including an empty range's start position. +/// +/// This immutable metadata does not check tensor byte extents or establish that +/// key/value storage is initialized, owned or ready. Operators must validate +/// those properties separately, using initialized cache lengths rather than +/// allocation capacities as `key_tokens`. +/// +/// ``` +/// use cute_rs::attention::AttentionWindow; +/// let appended = AttentionWindow::causal(3, 19, 16).unwrap(); +/// assert_eq!(appended.visible_keys(0).unwrap(), 17); +/// assert_eq!(appended.visible_keys(2).unwrap(), 19); +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AttentionWindow { + query_tokens: usize, + key_tokens: usize, + causal_start: Option, +} + +impl AttentionWindow { + /// Validate a causal query range starting at an absolute key position. + /// + /// `query_start + query_tokens` must not overflow or exceed `key_tokens`. + pub fn causal( + query_tokens: usize, + key_tokens: usize, + query_start: usize, + ) -> Result { + let mut window = Self::noncausal(query_tokens, key_tokens)?; + let query_end = query_start + .checked_add(query_tokens) + .ok_or(AttentionError::ExtentOverflow)?; + if query_end > key_tokens { + return Err(AttentionError::InvalidWindow); + } + window.causal_start = Some(query_start); + Ok(window) + } + + /// Validate a noncausal window in which every query sees all keys. + /// Query and key sequence lengths may differ. + pub fn noncausal(query_tokens: usize, key_tokens: usize) -> Result { + if query_tokens != 0 && key_tokens == 0 { + return Err(AttentionError::MissingKeys); + } + Ok(Self { + query_tokens, + key_tokens, + causal_start: None, + }) + } + + /// Number of query rows in the window. + #[must_use] + pub const fn query_tokens(self) -> usize { + self.query_tokens + } + + /// Number of available key/value token rows. + #[must_use] + pub const fn key_tokens(self) -> usize { + self.key_tokens + } + + /// Absolute key position of the first causal query, or `None` if noncausal. + #[must_use] + pub const fn causal_start(self) -> Option { + self.causal_start + } + + /// Count the visible key prefix for one local query row. + /// + /// Rejects `query_row >= query_tokens`, including every row of an empty + /// query sequence. No causal row is fully masked by a validated window. + pub fn visible_keys(self, query_row: usize) -> Result { + if query_row >= self.query_tokens { + return Err(AttentionError::InvalidQueryRow); + } + Ok(match self.causal_start { + // The constructor checked start + query_tokens, and row is smaller. + Some(start) => start + query_row + 1, + None => self.key_tokens, + }) + } +} + +/// Validated capacity of each dense `[token, KV head, dimension]` cache tensor. +/// +/// This is immutable metadata, not an initialized-cache handle. Head dimensions +/// remain part of its Rust type so checked callers can require compatible caches: +/// +/// ```compile_fail,E0308 +/// use cute_rs::attention::CacheShape; +/// fn needs_128_wide_heads(_: CacheShape<8, 128>) {} +/// let cache = CacheShape::<8, 64>::new(65).unwrap(); +/// needs_128_wide_heads(cache); +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CacheShape { + /// Maximum number of tokens per cache tensor. + capacity: usize, + /// BF16 elements per token. + width: usize, + /// BF16 elements in each complete key or value allocation. + allocation_len: usize, +} + +impl CacheShape { + /// Check head dimensions and each cache tensor's full BF16 allocation extent. + /// Zero token capacity is allowed; head dimensions must remain valid. + pub fn new(capacity: usize) -> Result { + let width = checked_width(KVH, D)?; + Ok(Self { + capacity, + width, + allocation_len: checked_len(capacity, width)?, + }) + } + + /// Maximum number of tokens in each cache tensor. + #[must_use] + pub const fn capacity(self) -> usize { + self.capacity + } + + /// BF16 elements per token in each cache tensor. + #[must_use] + pub const fn width(self) -> usize { + self.width + } + + /// BF16 elements in each full allocation; K and V need this many separately. + #[must_use] + pub const fn allocation_len(self) -> usize { + self.allocation_len + } + + /// Validate an append and return its half-open **token** range. + /// + /// The caller supplies the number of initialized tokens. This checks its + /// bounds, not initialization itself, and neither reserves storage nor + /// advances a cursor. The owner must perform and complete the append before + /// publishing the new cursor. An empty append at capacity is valid. + pub fn append_range( + self, + valid_tokens: usize, + new_tokens: usize, + ) -> Result, AttentionError> { + if valid_tokens > self.capacity { + return Err(AttentionError::InvalidCursor); + } + let end = valid_tokens + .checked_add(new_tokens) + .ok_or(AttentionError::ExtentOverflow)?; + if end > self.capacity { + return Err(AttentionError::CapacityExceeded); + } + Ok(valid_tokens..end) + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/mod.rs new file mode 100644 index 0000000000..bf1a2a024c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/mod.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Attention metadata and architecture-specific device operations. +//! +//! [`sm100`] is available to device-only consumers. Checked BF16 shape and +//! cache metadata requires the `host` feature and is reexported here to +//! preserve the existing `attention::AttentionShape` and related imports. + +#[cfg(feature = "host")] +pub mod metadata; +pub mod sm100; + +#[cfg(feature = "host")] +pub use metadata::{AttentionError, AttentionShape, AttentionWindow, CacheShape}; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/__compiler.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/__compiler.rs new file mode 100644 index 0000000000..c8eba3ac99 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/__compiler.rs @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Exact compiler markers for independent attention operations. + +use crate::copy::tma::TensorMap; +use ctm_rs::mbarrier::Barrier; +/// Device-only shared/shared MMA marker; operand displacements count bytes from retained bases. +/// +/// # Safety +/// The public MmaSs::gemm participation, completion and storage contracts apply. +#[inline(never)] +pub unsafe fn mma_ss( + dst: u32, + a: *mut f16, + b: *mut f16, + accumulate: bool, + a_byte_offset: u32, + b_byte_offset: u32, +) { + let _ = (dst, a, b, accumulate, a_byte_offset, b_byte_offset); + unreachable!("CuTe compiler marker") +} +/// Device-only TMEM/shared MMA marker; B’s displacement counts bytes from its retained base. +/// +/// # Safety +/// The public MmaTs::gemm participation, nonoverlap and completion contracts apply. +#[inline(never)] +pub unsafe fn mma_ts( + dst: u32, + a: u32, + b: *mut f16, + accumulate: bool, + b_byte_offset: u32, +) { + let _ = (dst, a, b, accumulate, b_byte_offset); + unreachable!("CuTe compiler marker") +} +/// Device-only rank-three load marker with inner, row and outer-chunk coordinates. +/// +/// # Safety +/// The producer warp must supply matching live storage and a correctly initialized group completion barrier. +#[inline(never)] +#[allow(clippy::too_many_arguments)] +pub unsafe fn tma_load( + desc: *const TensorMap, + dst: *mut f16, + c0: i32, + c1: i32, + c2: i32, + bar: *mut Barrier, + rank: u32, +) { + let _ = (desc, dst, c0, c1, c2, bar, rank); + unreachable!("CuTe compiler marker") +} +/// Device-only rank-three store marker with inner, row and outer-chunk coordinates. +/// +/// # Safety +/// The elected issuer must publish shared writes and retain storage until the copy stops reading it. +#[inline(never)] +pub unsafe fn tma_store( + desc: *const TensorMap, + src: *mut f16, + c0: i32, + c1: i32, + c2: i32, +) { + let _ = (desc, src, c0, c1, c2); + unreachable!("CuTe compiler marker") +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/host.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/host.rs new file mode 100644 index 0000000000..1355e0f491 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/host.rs @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Host construction of the rank-three tensor maps consumed by attention tiles. +use super::SmemLayout; +use crate::copy::tma::TensorMap; +use alloc::{format, string::String}; +use core::{ffi::c_void, marker::PhantomData, mem::MaybeUninit}; +use cuda_core::sys as driver; + +/// Encode a physical row-major FP16 matrix into the selected shared operand layout. +/// +/// This is the attention form of [`crate::copy::tma::create_tensor_map_tiled`]. +/// The generic encoder describes one two-dimensional box. This one encodes a +/// rank-three tensor map whose innermost dimension is the swizzle span, so a +/// single copy fills a whole SM100 operand tile from chunks of `span` columns. +/// It also checks the SM100 tile constraints (a span of 32 or 64 elements, +/// a row count that is a multiple of 8 and at most 256, a column count that +/// is a multiple of the span), transposes the logical rows and columns of an +/// MN-major operand relative to the physical ones (global coordinates always +/// use the physical ordering), and fixes the element type to `f16` with +/// 128-byte L2 promotion. +/// +/// The descriptor does not retain the allocation or track its lifetime; callers +/// keep the allocation live until every asynchronous use completes. +#[allow(clippy::not_unsafe_ptr_arg_deref)] +pub fn create_tensor_map_tiled( + base: *mut c_void, + rows: u64, + cols: u64, + row_stride_elements: u64, +) -> Result, String> { + let () = L::VALID; + let span = L::SPAN; + let (tile_rows, tile_cols) = if L::MN_MAJOR { + (L::COLS, L::ROWS) + } else { + (L::ROWS, L::COLS) + }; + if !matches!(span, 32 | 64) + || tile_rows == 0 + || !tile_rows.is_multiple_of(8) + || tile_rows > 256 + || tile_cols == 0 + || !tile_cols.is_multiple_of(span) + { + return Err("unsupported SM100 shared TMA tile".into()); + } + if rows == 0 || cols == 0 || !cols.is_multiple_of(u64::from(span)) || row_stride_elements < cols + { + return Err("TMA matrix must be nonempty, with contiguous dimension divisible by its swizzle span and pitch at least the column extent".into()); + } + if base.is_null() || !(base as usize).is_multiple_of(16) { + return Err("TMA global base must be nonnull and 16-byte aligned".into()); + } + let pitch = row_stride_elements + .checked_mul(2) + .ok_or("TMA byte pitch overflow")?; + if pitch % 16 != 0 { + return Err("TMA byte pitch must be a multiple of 16".into()); + } + if rows > u32::MAX as u64 || cols > u32::MAX as u64 || pitch >= (1u64 << 40) { + return Err("TMA extent or pitch exceeds the driver encoding".into()); + } + let dims = [u64::from(span), rows, cols / u64::from(span)]; + let strides = [pitch, u64::from(span) * 2]; + let tile = [span, tile_rows, tile_cols / span]; + let element_strides = [1u32; 3]; + let swizzle = if span == 32 { + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_64B + } else { + driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B + }; + let mut out = MaybeUninit::::uninit(); + let status = unsafe { + driver::cuTensorMapEncodeTiled( + out.as_mut_ptr(), + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16, + 3, + base, + dims.as_ptr(), + strides.as_ptr(), + tile.as_ptr(), + element_strides.as_ptr(), + driver::CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + driver::CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + }; + if status != driver::cudaError_enum_CUDA_SUCCESS { + return Err(format!( + "SM100 rank-three cuTensorMapEncodeTiled failed: {status:?}" + )); + } + Ok(TensorMap { + bytes: unsafe { core::mem::transmute::(out.assume_init()) }, + marker: PhantomData, + }) +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mma.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mma.rs new file mode 100644 index 0000000000..d4391551af --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mma.rs @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Independent shared/shared and tensor/shared FP16 attention MMA. + +use super::{__compiler, SharedTile, SmemLayout}; +use crate::sm100::{OperandA, OperandB}; +use core::marker::PhantomData; + +/// Independent shared/shared FP16 MMA with FP32 accumulation. +/// +/// `M = 128`; `N` is a multiple of 16 in `16..=256`; `K` is a positive +/// multiple of 16; `GROUP` is 1 or 2. M/N describe the complete product. +/// Each CTA's layouts must cover A `(M / GROUP, K)` and B `(N / GROUP, K)` +/// and satisfy [`SharedLayout`](super::SharedLayout)'s restrictions. +/// +/// ``` +/// use cute_rs::attention::sm100::{MmaSs, SharedLayout}; +/// type LocalTile = SharedLayout<64, 64>; +/// let mma = MmaSs::::new(); +/// ``` +pub struct MmaSs( + #[doc = "Carries both per-CTA shared operand layouts."] PhantomData<(AL, BL)>, +); +impl + MmaSs +{ + /// Check group size, collective dimensions and per-CTA shared operand shapes. + const VALID: () = { + let () = AL::VALID; + let () = BL::VALID; + assert!(G == 1 || G == 2); + assert!(M == 128 && N >= 16 && N <= 256 && N.is_multiple_of(16)); + assert!(K >= 16 && K.is_multiple_of(16)); + assert!(AL::ROWS == M / G && BL::ROWS == N / G); + assert!(AL::COLS == K && BL::COLS == K); + }; + /// Validate the static MMA dimensions and operand layouts. + #[inline(always)] + pub const fn new() -> Self { + let () = Self::VALID; + Self(PhantomData) + } + /// Issue the complete K tile, with the first instruction honoring `accumulate`. + /// # Safety + /// One elected leader lane issues MMA after input completion. Both shared + /// views and destination TMEM must remain live until its completion barrier. + /// Every operation using the allocation must use the same CTA group. + #[inline(always)] + pub unsafe fn gemm( + &self, + dst: u32, + a: &SharedTile, + b: &SharedTile, + accumulate: bool, + ) { + unsafe { + __compiler::mma_ss::( + dst, + a.base, + b.base, + accumulate, + a.byte_offset, + b.byte_offset, + ) + } + } +} +impl Default + for MmaSs +{ + fn default() -> Self { + Self::new() + } +} + +/// TMEM/shared FP16 MMA, used for the probability/value product in FMHA. +/// +/// `GROUP = 1`, `M = 128`, and `N` is a multiple of 16 in `16..=256`. +/// The importer supports `K = 32` or a positive multiple of 64. B's shared +/// layout must cover `(N, K)` and satisfy [`SharedLayout`](super::SharedLayout). +/// +/// ``` +/// use cute_rs::attention::sm100::{MmaTs, SharedLayout}; +/// let mma = MmaTs::, 128, 128, 64, 1>::new(); +/// ``` +pub struct MmaTs( + #[doc = "Carries the shared B operand’s physical layout."] PhantomData, +); +impl MmaTs { + /// Check the group-one TMEM/shared profile, supported K span and B layout shape. + const VALID: () = { + let () = BL::VALID; + assert!(G == 1 && M == 128); + assert!(N >= 16 && N <= 256 && N.is_multiple_of(16)); + assert!(K == 32 || (K > 0 && K.is_multiple_of(64))); + assert!(BL::ROWS == N && BL::COLS == K); + }; + /// Validate the static MMA dimensions and operand layouts. + #[inline(always)] + pub const fn new() -> Self { + let () = Self::VALID; + Self(PhantomData) + } + /// # Safety + /// An elected lane issues this after packed FP16 probability stores and + /// shared input copies are visible to MMA. Source and destination TMEM + /// regions must not overlap and must remain live until completion. + #[inline(always)] + pub unsafe fn gemm( + &self, + dst: u32, + a_tmem: u32, + b: &SharedTile, + accumulate: bool, + ) { + unsafe { + __compiler::mma_ts::(dst, a_tmem, b.base, accumulate, b.byte_offset) + } + } +} +impl Default + for MmaTs +{ + fn default() -> Self { + Self::new() + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mod.rs new file mode 100644 index 0000000000..9e658e6c83 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/mod.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Composable FP16 CuTe operations for data-center Blackwell (SM100). +//! +//! QK and PV retain separate layout and MMA contracts. Control flow, online +//! softmax, buffer phases and participating warps remain explicit in Rust. +//! Layout-aware MMA/TMA use the CuTe compiler path; pipeline synchronization +//! and TMEM register transfers use CTM instruction primitives. +pub mod pipeline; + +mod mma; +mod shared; +mod tma; +mod tmem; + +pub use mma::{MmaSs, MmaTs}; +pub use shared::{SharedLayout, SharedTile, SharedTileArray, SmemLayout}; +pub use tma::TmaCopy; +pub use tmem::TmemFragment; + +#[doc(hidden)] +pub mod __compiler; + +#[cfg(feature = "host")] +mod host; +#[cfg(feature = "host")] +pub use host::create_tensor_map_tiled; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/pipeline.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/pipeline.rs new file mode 100644 index 0000000000..fdd8784c46 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/pipeline.rs @@ -0,0 +1,689 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Two-CTA SM100 barrier rings with independently typed full and empty edges. +//! +//! Policies restrict completion mechanisms and preserve memory ordering. +//! Kernels still supply correct participants, persistent positions, buffer +//! lifetimes, and completion waits for deferred operations. +//! +//! Choose each edge by what its receiver will access: +//! +//! | Handoff | Policy | +//! | --- | --- | +//! | TMA/MMA completion before MMA | [`AsyncCompletion`] | +//! | MMA completion before TMA overwrites shared memory | [`AsyncCopyReuse`] | +//! | Thread-written shared data before MMA | [`SharedToAsync`] | +//! | MMA completion before ordinary shared writes | [`SharedReuse`] | +//! | TMEM data between threads | [`TensorMemory`] | +//! | Local shared data, without / with a later TC access | [`LocalShared`] / [`LocalMemory`] | +//! +//! A [`RingCursor`] tracks one role's position; [`TwoCtaRing`] documents the +//! producer/consumer handoff. Advancing a cursor never signals a barrier. + +use core::marker::PhantomData; +use core::mem::{align_of, size_of}; +use ctm_rs::mbarrier::{self, Barrier}; +use ctm_rs::{cluster, fence, groups::CtaGroup2, memory::SharedAddress, tcgen05}; + +/// Restricts edge policies and cursor roles to the implementations in this module. +mod sealed { + /// Marks supported receiver ordering and whether later TC access needs a fence. + pub trait Wait { + // Existing policies order a subsequent TC receiver. Only explicitly + // non-TC receiver policies may omit this fence. + /// Whether the ring inserts an after-thread-sync fence following this policy’s wait. + const TC_AFTER: bool = true; + } + /// Admits per-thread remote arrivals for a policy that publishes the correct memory domain. + pub trait Signal {} + /// Admits group-two MMA completion as the signal for this edge. + pub trait Mma {} + /// Admits transaction-counted TMA completion as the signal for this edge. + pub trait Tma {} + /// Restricts cursors to producer and consumer phase conventions. + pub trait Cursor {} +} + +/// Sealed ordering policy for one barrier-ring edge. +pub trait WaitPolicy: sealed::Wait + Copy { + /// Observe the phase and acquire this policy's protected memory. + /// # Safety + /// The local shared barrier must be live and initialized, with matching + /// signalers and a current or immediately preceding phase. The ring supplies + /// an after-thread-sync fence for policies with TC receivers. Non-TC policies + /// require a separate fence before any later asynchronous TC operations. + unsafe fn wait(bar: *const Barrier, phase: u32); +} +/// Sealed policy permitting each participating thread to signal rank zero. +pub trait SignalPolicy: WaitPolicy + sealed::Signal { + /// Publish one thread's handoff to the leader's corresponding barrier. + /// # Safety + /// The local barrier must have the same offset in exactly two live CTAs. + /// Complete deferred operations first and contribute exactly the initialized + /// number of arrivals, with each participant arriving once per phase. + unsafe fn arrive(bar: *mut Barrier); +} +/// Sealed policy accepting group-two MMA multicast completion. +pub trait MmaPolicy: WaitPolicy + sealed::Mma {} +/// Sealed policy accepting transaction-counted TMA completion. +pub trait TmaPolicy: WaitPolicy + sealed::Tma {} + +/// CTA acquire for local shared data or local TMEM metadata exchanges. +/// Local arrivals and cp.async attachment use unsafe raw accessors. TMEM +/// signalers must complete their operations and fence before those arrivals. +#[derive(Clone, Copy)] +pub struct LocalMemory; +impl sealed::Wait for LocalMemory {} +impl WaitPolicy for LocalMemory { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity(bar, phase) } {} + } +} + +/// CTA acquire for ordinary shared-memory or cp.async page exchanges. +/// +/// This retains local acquire semantics and cp.async completion visibility. +/// It supplies no TC after-thread-sync fence: the receiver must use ordinary +/// memory or asynchronous copies, with no subsequent asynchronous TC access +/// before a separate TC fence. Local signaling uses the unsafe raw accessors. +#[derive(Clone, Copy)] +pub struct LocalShared; +impl sealed::Wait for LocalShared { + const TC_AFTER: bool = false; +} +impl WaitPolicy for LocalShared { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity(bar, phase) } {} + } +} + +/// Direct TMA/MMA completion without publishing an ordinary-memory payload. +/// The ring orders subsequent MMA using a TC fence. Subsequent TMA must be +/// issued after the wait. This policy is not a thread-written data handoff. +#[derive(Clone, Copy)] +pub struct AsyncCompletion; +impl sealed::Wait for AsyncCompletion {} +impl sealed::Mma for AsyncCompletion {} +impl sealed::Tma for AsyncCompletion {} +impl MmaPolicy for AsyncCompletion {} +impl TmaPolicy for AsyncCompletion {} +impl WaitPolicy for AsyncCompletion { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + } +} + +/// Observe completed MMA reads before reusing shared storage with TMA copies. +/// +/// The cluster-scoped completion poll is followed by the caller's asynchronous +/// copy issue. This policy supplies no TC fence and does not acquire an ordinary +/// memory payload. Use it for an empty edge whose receiver issues TMA copies; +/// ordinary shared overwrites require [`SharedReuse`], and subsequent TC access +/// requires a separate after-thread-sync fence. +#[derive(Clone, Copy)] +pub struct AsyncCopyReuse; +impl sealed::Wait for AsyncCopyReuse { + const TC_AFTER: bool = false; +} +impl sealed::Mma for AsyncCopyReuse {} +impl MmaPolicy for AsyncCopyReuse {} +impl WaitPolicy for AsyncCopyReuse { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + } +} + +/// TMEM handoff using cluster execution ordering and TC thread-sync fences. +/// Thread signalers must complete deferred TMEM loads/stores first. Ordinary +/// shared or global payload writes are not published by this policy. +#[derive(Clone, Copy)] +pub struct TensorMemory; +impl sealed::Wait for TensorMemory {} +impl sealed::Signal for TensorMemory {} +impl sealed::Mma for TensorMemory {} +impl MmaPolicy for TensorMemory {} +impl WaitPolicy for TensorMemory { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + } +} +impl SignalPolicy for TensorMemory { + #[inline(always)] + unsafe fn arrive(bar: *mut Barrier) { + tcgen05::fence_before_thread_sync(); + unsafe { arrive_remote(bar) }; + } +} + +/// Publish ordinary local shared stores to the two-CTA async MMA reader. +/// Every writer releases through the generic-to-async proxy fence; the wait +/// acquires cluster shared memory in the async proxy. +#[derive(Clone, Copy)] +pub struct SharedToAsync; +impl sealed::Wait for SharedToAsync {} +impl sealed::Signal for SharedToAsync {} +impl WaitPolicy for SharedToAsync { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + unsafe { fence::fence_proxy_async_acquire_sync_restrict() }; + } +} +impl SignalPolicy for SharedToAsync { + #[inline(always)] + unsafe fn arrive(bar: *mut Barrier) { + unsafe { + fence::fence_proxy_async_release_sync_restrict(); + arrive_remote(bar); + } + } +} + +/// Acquire completed MMA shared reads before ordinary shared-buffer overwrites. +/// The generic acquire keeps cluster scope and restricts its memory domain to +/// shared objects; it does not acquire unrelated global memory. +#[derive(Clone, Copy)] +pub struct SharedReuse; +impl sealed::Wait for SharedReuse {} +impl sealed::Mma for SharedReuse {} +impl MmaPolicy for SharedReuse {} +impl WaitPolicy for SharedReuse { + #[inline(always)] + unsafe fn wait(bar: *const Barrier, phase: u32) { + while !unsafe { mbarrier::mbarrier_try_wait_parity_cluster_relaxed(bar, phase) } {} + unsafe { fence::fence_sync_restrict_acquire() }; + } +} + +/// Cursor roles shared with [`crate::pipeline::PipelineState`]: a producer +/// cursor starts at empty phase one, a consumer cursor at full phase zero. +pub use crate::pipeline::{Consumer, Producer}; +impl sealed::Cursor for Producer {} +impl sealed::Cursor for Consumer {} + +/// Sealed role fixing a cursor's initial phase and permitted ring operations. +pub trait CursorRole: sealed::Cursor + Copy { + /// Starting barrier parity: one for producer empty waits, zero for consumer full waits. + #[doc(hidden)] + const INITIAL_PHASE: u32; +} +impl CursorRole for Producer { + const INITIAL_PHASE: u32 = 1; +} +impl CursorRole for Consumer { + const INITIAL_PHASE: u32 = 0; +} + +/// Persistent position in a ring, with a bounded slot and explicit wait phase. +/// +/// Advance exactly once after each stage handoff, and retain the cursor across +/// work items. Advancing wraps the slot and toggles the phase without division +/// or a monotonically increasing counter. Copies preserve position only; they +/// grant no extra ownership, participants, or permission to repeat a handoff. +/// +/// ``` +/// use cute_rs::attention::sm100::pipeline::{Producer, RingCursor}; +/// let mut cursor = RingCursor::<2, Producer>::new(); +/// cursor.advance(); +/// cursor.advance(); +/// assert_eq!((cursor.index(), cursor.phase()), (0, 0)); // Wrapped; phase flipped. +/// ``` +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RingCursor { + /// Bounded barrier-array slot in 0..STAGES. + index: u32, + /// Parity this role waits for at the current slot. + phase: u32, + /// Records the role’s initial phase and permitted handoff operations. + role: PhantomData, +} +impl RingCursor { + /// Reject zero-stage cursors before any slot arithmetic. + const VALID: () = assert!(STAGES > 0); + + /// Start at slot zero, with the producer or consumer's initial wait phase. + #[inline(always)] + pub const fn new() -> Self { + let () = Self::VALID; + Self { + index: 0, + phase: Role::INITIAL_PHASE, + role: PhantomData, + } + } + /// Current stage index, always less than STAGES. + /// + /// It is a `u32` like this ring's stage count; the TMA load pipeline's + /// [`crate::PipelineState::index`] is a `usize`. + #[inline(always)] + pub const fn index(self) -> u32 { + self.index + } + /// Phase awaited by this role at the current stage. + #[inline(always)] + pub const fn phase(self) -> u32 { + self.phase + } + /// Move to the next stage, toggling phase on wrap. + /// + /// This updates software state only; it performs no barrier operation. + #[inline(always)] + pub const fn advance(&mut self) { + self.index += 1; + if self.index == STAGES { + self.index = 0; + self.phase ^= 1; + } + } +} +impl Default for RingCursor { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +/// Non-owning view of adjacent full and empty barrier arrays in two CTAs. +/// +/// Both CTAs have identical shared-memory layouts. Thread arrivals target rank +/// zero; MMA completion multicasts to ranks zero and one. Local-only edges use +/// LocalMemory and the unsafe raw barrier accessors. +/// +/// Tickets start at zero and continue across persistent work. The initial waited +/// full and empty parities are zero and one respectively; initialization puts +/// both hardware barriers in phase zero. Do not reset or overflow tickets +/// while phases remain live. Copies grant no additional ownership, participants, +/// or permission to signal a phase twice. +/// +/// For a TMA-to-MMA B-operand ring, use `AsyncCompletion` for full and +/// `AsyncCopyReuse` for empty. One lane per CTA initializes with `init(1, 1)`; +/// publish initialization as required by [`Self::init`]. Keep both role cursors +/// across work items. The following device-only steps assume matching live +/// storage; the caller selects the B tile at the cursor's stage. +/// +/// Every producer-warp lane in each CTA acquires and copies its 64×64 FP16 +/// tile. One elected rank-zero lane expects 16384 bytes across the CTA pair: +/// +/// ```no_run +/// # #![feature(f16)] +/// # use cute_rs::{OperandB, TensorMap, attention::sm100::{SharedLayout, SharedTile, TmaCopy, pipeline::*}}; +/// type L = SharedLayout<64, 64>; +/// # type Ring = TwoCtaRing<2, AsyncCompletion, AsyncCopyReuse>; +/// # unsafe fn produce(ring: Ring, p: &mut RingCursor<2, Producer>, desc: &TensorMap, b: &mut SharedTile, rank: u32) { +/// unsafe { +/// ring.acquire_cursor(*p); +/// if rank == 0 && ctm_rs::warp::is_elected_sync(u32::MAX) { +/// ring.producer_expect_tx_cursor(*p, 16384); +/// } +/// let done = ring.full_at_cursor(*p); +/// TmaCopy::::load(desc, b, (rank * 64) as i32, 0, done, rank); +/// p.advance(); +/// } +/// # } +/// ``` +/// +/// The leader CTA's MMA warp waits; its elected lane issues the product and +/// attaches the empty signal to MMA completion. A must already be ready; A, B +/// and destination TMEM stay live until completion. `accumulate = true` +/// requires initialized TMEM. The next producer acquire waits for this release: +/// +/// ```no_run +/// # use cute_rs::{OperandA, OperandB, attention::sm100::{MmaSs, SharedLayout, SharedTile, pipeline::*}}; +/// # type L = SharedLayout<64, 64>; +/// # type Ring = TwoCtaRing<2, AsyncCompletion, AsyncCopyReuse>; +/// # unsafe fn consume(ring: Ring, c: &mut RingCursor<2, Consumer>, mma: &MmaSs, dst: u32, a: &SharedTile, b: &SharedTile, accumulate: bool) { +/// unsafe { +/// ring.wait_cursor(*c); +/// if ctm_rs::warp::is_elected_sync(u32::MAX) { +/// mma.gemm(dst, a, b, accumulate); +/// ring.mma_empty_cursor(*c); +/// } +/// c.advance(); +/// } +/// # } +/// ``` +#[derive(Clone, Copy)] +pub struct TwoCtaRing { + /// Local array of full barriers with identical offsets in both CTAs. + full: *mut Barrier, + /// Local array of empty barriers immediately after the full array. + empty: *mut Barrier, + /// Carries independent receiver-ordering policies for the two edges. + ordering: PhantomData<(Full, Empty)>, +} + +impl TwoCtaRing { + /// Reject zero-stage rings before modulo or storage-size calculations. + const VALID: () = assert!(STAGES > 0); + /// Required alignment of the barrier storage. + pub const STORAGE_ALIGN: usize = align_of::(); + /// Bytes occupied by adjacent full and empty arrays; rejects zero stages. + pub const STORAGE_BYTES: usize = { + let () = Self::VALID; + 2 * STAGES as usize * size_of::() + }; + + /// View full then empty arrays, starting at `base[first]`. + /// # Safety + /// The aligned allocation must contain at least first + 2 * STAGES barriers. + /// Before device operations, it must be local shared storage with identical + /// offsets in both CTAs, used exclusively by this pipeline until all async + /// work and waits finish. Before repurposing this storage within the kernel, + /// invalidate every initialized barrier after all participants have finished. + /// This non-owning view neither initializes nor invalidates storage. + #[inline(always)] + pub const unsafe fn from_raw(base: *mut Barrier, first: usize) -> Self { + let () = Self::VALID; + Self { + full: unsafe { base.add(first) }, + empty: unsafe { base.add(first + STAGES as usize) }, + ordering: PhantomData, + } + } + #[inline(always)] + /// Reduce a persistent ticket to its barrier-array slot. + const fn slot(ticket: u32) -> usize { + (ticket % STAGES) as usize + } + #[inline(always)] + /// Select the full barrier parity, flipping after each complete ring traversal. + const fn full_phase(ticket: u32) -> u32 { + (ticket / STAGES) & 1 + } + #[inline(always)] + /// Select the producer’s empty parity, opposite the full parity for the same ticket. + const fn empty_phase(ticket: u32) -> u32 { + Self::full_phase(ticket) ^ 1 + } + + /// Initialize every stage's counts from one thread in each CTA. + /// # Safety + /// Counts must be positive valid mbarrier counts matching the protocol. + /// Each storage location must be uninitialized or explicitly invalidated; + /// an idle but valid barrier must not be initialized again. No participant + /// may still access a previous barrier in these arrays. Publish initialization + /// and synchronize the cluster before use; constructor storage requirements + /// remain in force. + #[inline(always)] + pub unsafe fn init(self, producers: u32, consumers: u32) { + let mut stage = 0; + while stage < STAGES { + unsafe { + mbarrier::mbarrier_init(self.full.add(stage as usize), producers); + mbarrier::mbarrier_init(self.empty.add(stage as usize), consumers); + } + stage += 1; + } + } + /// Borrow the local full barrier for raw TMA or cp.async attachment. + /// # Safety + /// Storage must satisfy the constructor's allocation contract. Raw operations + /// must obey Full, counts, and ticket phase; this grants no extra arrival. + #[inline(always)] + pub unsafe fn full_at(self, ticket: u32) -> *mut Barrier { + unsafe { self.full.add(Self::slot(ticket)) } + } + /// Borrow the local empty barrier for raw local release operations. + /// # Safety + /// Storage must satisfy the constructor's allocation contract. Raw releases + /// must obey Empty, completion fences, counts, and ticket phase. + #[inline(always)] + pub unsafe fn empty_at(self, ticket: u32) -> *mut Barrier { + unsafe { self.empty.add(Self::slot(ticket)) } + } + /// Wait for reuse with the empty policy's receiver ordering. + /// # Safety + /// Use the current producer ticket with matching consumer releases. This + /// grants no additional ownership and does not order arbitrary generic + /// accesses when Empty only tracks async/TMEM completion. Non-TC policies + /// do not order subsequent asynchronous TC operations. + #[inline(always)] + pub unsafe fn acquire(self, ticket: u32) { + unsafe { Empty::wait(self.empty_at(ticket), Self::empty_phase(ticket)) }; + after_wait::(); + } + /// Wait for production with the full policy's receiver ordering. + /// # Safety + /// Use the current consumer ticket with matching producer signals. Every + /// consuming thread must wait before accessing the protected stage and obey + /// Full's receiver restrictions; non-TC policies supply no TC fence. + #[inline(always)] + pub unsafe fn wait(self, ticket: u32) { + unsafe { Full::wait(self.full_at(ticket), Self::full_phase(ticket)) }; + after_wait::(); + } + /// Set expected TMA bytes and contribute one full-barrier arrival. + /// # Safety + /// One elected rank-zero thread calls after acquire. TMA must complete + /// exactly bytes across the CTA pair on this full barrier. + #[inline(always)] + pub unsafe fn producer_expect_tx(self, ticket: u32, bytes: u32) + where + Full: TmaPolicy, + { + unsafe { mbarrier::mbarrier_arrive_expect_tx_cluster(self.full_at(ticket), 1, bytes) }; + } + /// Signal full in both CTAs after prior group-two MMA completes. + /// # Safety + /// One elected rank-zero thread calls after all MMA producing the stage. + /// Its one completion arrival must match the full barrier count. + #[inline(always)] + pub unsafe fn mma_full(self, ticket: u32) + where + Full: MmaPolicy, + { + unsafe { mma_arrive(self.full_at(ticket)) }; + } + /// Signal empty in both CTAs after prior group-two MMA completes. + /// # Safety + /// One elected rank-zero thread calls after all MMA consuming the stage. + /// Its one completion arrival must match the empty barrier count. + #[inline(always)] + pub unsafe fn mma_empty(self, ticket: u32) + where + Empty: MmaPolicy, + { + unsafe { mma_arrive(self.empty_at(ticket)) }; + } + /// Publish one thread's production to the leader's full barrier. + /// # Safety + /// Complete deferred operations first. Each counted producer contributes + /// one arrival after its last protected access, using the current ticket. + #[inline(always)] + pub unsafe fn thread_full(self, ticket: u32) + where + Full: SignalPolicy, + { + unsafe { Full::arrive(self.full_at(ticket)) }; + } + /// Publish one thread's consumption to the leader's empty barrier. + /// # Safety + /// Complete deferred operations first. Each counted consumer contributes + /// one arrival after its last protected access, using the current ticket. + #[inline(always)] + pub unsafe fn thread_empty(self, ticket: u32) + where + Empty: SignalPolicy, + { + unsafe { Empty::arrive(self.empty_at(ticket)) }; + } +} + +impl TwoCtaRing { + /// Borrow the full barrier at an already bounded cursor index. + /// # Safety + /// The storage and raw-operation requirements of [`Self::full_at`] apply. + #[inline(always)] + pub unsafe fn full_at_cursor( + self, + cursor: RingCursor, + ) -> *mut Barrier { + unsafe { self.full.add(cursor.index as usize) } + } + /// Borrow the empty barrier at an already bounded cursor index. + /// # Safety + /// The storage and raw-operation requirements of [`Self::empty_at`] apply. + #[inline(always)] + pub unsafe fn empty_at_cursor( + self, + cursor: RingCursor, + ) -> *mut Barrier { + unsafe { self.empty.add(cursor.index as usize) } + } + /// Wait for reuse using a producer cursor, without deriving slot or phase. + /// # Safety + /// The requirements of [`Self::acquire`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn acquire_cursor(self, cursor: RingCursor) { + unsafe { Empty::wait(self.empty_at_cursor(cursor), cursor.phase) }; + after_wait::(); + } + /// Wait for production using a consumer cursor. + /// # Safety + /// The requirements of [`Self::wait`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn wait_cursor(self, cursor: RingCursor) { + unsafe { Full::wait(self.full_at_cursor(cursor), cursor.phase) }; + after_wait::(); + } + /// Expect TMA bytes and arrive at the producer cursor's full barrier. + /// # Safety + /// The requirements of [`Self::producer_expect_tx`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn producer_expect_tx_cursor(self, cursor: RingCursor, bytes: u32) + where + Full: TmaPolicy, + { + unsafe { + mbarrier::mbarrier_arrive_expect_tx_cluster(self.full_at_cursor(cursor), 1, bytes) + }; + } + /// Signal MMA production at the producer cursor's full barrier. + /// # Safety + /// The requirements of [`Self::mma_full`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn mma_full_cursor(self, cursor: RingCursor) + where + Full: MmaPolicy, + { + unsafe { mma_arrive(self.full_at_cursor(cursor)) }; + } + /// Signal MMA consumption at the consumer cursor's empty barrier. + /// # Safety + /// The requirements of [`Self::mma_empty`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn mma_empty_cursor(self, cursor: RingCursor) + where + Empty: MmaPolicy, + { + unsafe { mma_arrive(self.empty_at_cursor(cursor)) }; + } + /// Publish a thread's production at the producer cursor's full barrier. + /// # Safety + /// The requirements of [`Self::thread_full`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn thread_full_cursor(self, cursor: RingCursor) + where + Full: SignalPolicy, + { + unsafe { Full::arrive(self.full_at_cursor(cursor)) }; + } + /// Publish a thread's consumption at the consumer cursor's empty barrier. + /// # Safety + /// The requirements of [`Self::thread_empty`] apply to the current cursor position. + #[inline(always)] + pub unsafe fn thread_empty_cursor(self, cursor: RingCursor) + where + Empty: SignalPolicy, + { + unsafe { Empty::arrive(self.empty_at_cursor(cursor)) }; + } +} + +#[inline(always)] +/// Insert the TC receiver fence only for policies that require it. +fn after_wait() { + if ::TC_AFTER { + tcgen05::fence_after_thread_sync(); + } +} + +#[inline(always)] +/// Signal the matching rank-zero barrier without an implicit memory release. +/// +/// # Safety +/// The two CTAs must have matching live initialized barriers and the caller must supply the policy’s publication fences. +unsafe fn arrive_remote(bar: *mut Barrier) { + unsafe { + cluster::arrive_remote_relaxed(cluster::mapa(SharedAddress::from_raw(bar), 0)); + } +} +#[inline(always)] +/// Attach prior group-two MMA completion to both CTAs’ matching barriers. +/// +/// # Safety +/// The elected leader issuer must have matching initialized barriers and expected completion counts in both CTAs. +unsafe fn mma_arrive(bar: *mut Barrier) { + unsafe { + tcgen05::commit_multicast_cg2( + SharedAddress::from_raw(bar), + cluster::CtaMask::::try_from_bits(3).unwrap_unchecked(), + ); + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + /// Ticket arithmetic must wrap stage indices and preserve opposite full/empty parities. + #[test] + fn tickets_wrap_slots_and_flip_opposite_phases() { + /// Stage-count and ordering-policy fixture for checking ticket parity arithmetic. + type Ring = TwoCtaRing<3, AsyncCompletion>; + for (ticket, slot, full, empty) in [ + (0, 0, 0, 1), + (1, 1, 0, 1), + (2, 2, 0, 1), + (3, 0, 1, 0), + (4, 1, 1, 0), + (5, 2, 1, 0), + (6, 0, 0, 1), + ] { + assert_eq!(Ring::slot(ticket), slot); + assert_eq!(Ring::full_phase(ticket), full); + assert_eq!(Ring::empty_phase(ticket), empty); + } + } + /// A new work item must retain the prior ring position instead of resetting its parity. + #[test] + fn persistent_work_keeps_phase_at_a_partial_ring_boundary() { + /// Stage-count and ordering-policy fixture for checking ticket parity arithmetic. + type Ring = TwoCtaRing<15, AsyncCompletion>; + // Two work items consume 17 stages each without restarting the ring. + for (ticket, slot, full, empty) in [(17, 2, 1, 0), (34, 4, 0, 1)] { + assert_eq!(Ring::slot(ticket), slot); + assert_eq!(Ring::full_phase(ticket), full); + assert_eq!(Ring::empty_phase(ticket), empty); + } + } + /// A one-slot ring must alternate full and empty parities on each ticket. + #[test] + fn one_stage_ring_flips_phase_on_every_ticket() { + /// Stage-count and ordering-policy fixture for checking ticket parity arithmetic. + type Ring = TwoCtaRing<1, TensorMemory>; + for (ticket, full, empty) in [(0, 0, 1), (1, 1, 0), (2, 0, 1), (3, 1, 0)] { + assert_eq!(Ring::slot(ticket), 0); + assert_eq!(Ring::full_phase(ticket), full); + assert_eq!(Ring::empty_phase(ticket), empty); + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/shared.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/shared.rs new file mode 100644 index 0000000000..3c5b4c7627 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/shared.rs @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared layouts and tile arrays for independent attention MMA. + +use core::marker::PhantomData; + +/// Physical layout of a shared FP16 operand, with CuTe's tiled byte swizzle. +mod sealed { + /// Marks only shared layouts recognized by the compiler. + pub trait Layout {} +} + +/// Implemented only by the layouts whose structure the compiler recognizes. +pub trait SmemLayout: sealed::Layout { + /// Evaluating this constant rejects an unsupported physical layout. + const VALID: (); + /// Logical MN extent. + const ROWS: u32; + /// Logical K extent. + const COLS: u32; + /// Whether the MN mode is contiguous. + const MN_MAJOR: bool; + /// Contiguous FP16 values in one swizzle atom. + const SPAN: u32; + /// Number of XOR bits in the shared-memory swizzle. + const SWIZZLE_BITS: u32; + /// Map a logical coordinate into shared storage in elements (CuTeDSL + /// `cute.crd2idx`). + fn crd2idx(row: u32, col: u32) -> usize; +} + +/// Logical `(MN,K)` layout. `MN_MAJOR` selects the contiguous logical mode. +/// Contiguous dimensions wider than 64 FP16 values are tiled in 128-byte chunks. +/// +/// The contiguous extent must be 32 or a positive multiple of 64; the other +/// extent must be a positive multiple of 8. Extents and offsets count FP16 +/// elements, not bytes. With `MN_MAJOR = false`, K is contiguous: +/// +/// ``` +/// use cute_rs::attention::sm100::SharedLayout; +/// type Tile = SharedLayout<64, 64>; +/// assert_eq!(Tile::crd2idx(0, 1), 1); +/// ``` +pub struct SharedLayout; +impl SharedLayout { + /// Reject unsupported contiguous swizzle spans and outer tile extents. + const VALID: () = { + let contiguous = if MN { R } else { C }; + let outer = if MN { C } else { R }; + assert!(contiguous >= 32 && contiguous.is_multiple_of(32)); + assert!(outer > 0 && outer.is_multiple_of(8)); + assert!(contiguous == 32 || contiguous.is_multiple_of(64)); + }; + /// Offset in FP16 elements from the start of this shared tile (CuTeDSL + /// `cute.crd2idx`). + #[inline(always)] + pub fn crd2idx(row: u32, col: u32) -> usize { + let () = Self::VALID; + let (inner, outer, outer_extent) = if MN { (row, col, C) } else { (col, row, R) }; + let span = Self::SPAN; + let linear = (inner / span) * (span * outer_extent) + outer * span + inner % span; + let mask = (1u32 << Self::SWIZZLE_BITS) - 1; + (linear ^ (((linear >> 6) & mask) << 3)) as usize + } +} +impl sealed::Layout for SharedLayout {} +impl SmemLayout for SharedLayout { + const VALID: () = SharedLayout::::VALID; + const ROWS: u32 = R; + const COLS: u32 = C; + const MN_MAJOR: bool = MN; + const SPAN: u32 = if (if MN { R } else { C }) == 32 { + 32 + } else { + 64 + }; + const SWIZZLE_BITS: u32 = if Self::SPAN == 32 { 2 } else { 3 }; + #[inline(always)] + fn crd2idx(row: u32, col: u32) -> usize { + Self::crd2idx(row, col) + } +} + +/// Non-owning shared tile, typed by its physical layout and MMA operand role. +pub struct SharedTile { + /// Original shared FP16 allocation base, before selecting an array tile. + pub(super) base: *mut f16, + /// Byte displacement of this tile from the original allocation base. + pub(super) byte_offset: u32, + /// Carries the physical layout and MMA operand role. + marker: PhantomData<(L, Role)>, +} +impl SharedTile { + /// # Safety + /// Storage must cover the layout and remain live through asynchronous use. + /// Its base must be aligned to the layout's complete swizzle period (1024 + /// bytes for B128, 512 bytes for B64). + #[inline(always)] + pub const unsafe fn from_raw(base: *mut f16) -> Self { + let () = L::VALID; + Self { + base, + byte_offset: 0, + marker: PhantomData, + } + } + /// Borrow this tile's shared storage address. + #[inline(always)] + pub const fn as_ptr(&self) -> *mut f16 { + // Forming a raw address does not access storage or require it to remain + // live. Dereferencing it retains the caller's lifetime obligations. + self.base.wrapping_add((self.byte_offset / 2) as usize) + } +} + +/// Contiguous array of complete shared tiles with one layout and operand role. +/// +/// Views retain the allocation base separately from their tile offset. MMA can +/// reuse the base descriptor; tensor copies address the selected physical tile. +/// Complete tiles preserve the XOR swizzle's phase at every array element. +pub struct SharedTileArray { + /// Shared FP16 allocation containing every complete array tile. + base: *mut f16, + /// Carries the array’s common layout and operand role. + marker: PhantomData<(L, Role)>, +} +impl SharedTileArray { + /// Bytes in one complete tile; require a whole swizzle period and a representable size. + const TILE_BYTES: u32 = { + let () = L::VALID; + let elements = (L::ROWS as u64) * (L::COLS as u64); + assert!(elements <= u32::MAX as u64 / 2); + let bytes = elements as u32 * 2; + assert!(bytes.is_multiple_of(1 << (L::SWIZZLE_BITS + 7))); + bytes + }; + /// Require a nonempty array whose total byte offset fits in u32. + const VALID: () = { + assert!(TILES > 0); + assert!(TILES <= u32::MAX / Self::TILE_BYTES); + }; + /// # Safety + /// `base` must point to live shared storage covering `TILES` consecutive + /// complete tiles. It must be aligned to the complete swizzle period (1024 + /// bytes for B128, 512 bytes for B64). The allocation and each accessed tile + /// must remain live through all asynchronous operations using their views. + #[inline(always)] + pub const unsafe fn from_raw(base: *mut f16) -> Self { + let () = Self::VALID; + Self { + base, + marker: PhantomData, + } + } + /// Select a complete tile while retaining the allocation's descriptor base. + /// # Safety + /// `index` must be less than `TILES`. All asynchronous accesses through this + /// view must obey the allocation lifetime and buffer-ownership protocol. + #[inline(always)] + pub const unsafe fn tile_unchecked(&self, index: u32) -> SharedTile { + SharedTile { + base: self.base, + byte_offset: index * Self::TILE_BYTES, + marker: PhantomData, + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tma.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tma.rs new file mode 100644 index 0000000000..9d2b26781c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tma.rs @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Rank-three copies matching independent attention operand layouts. + +use super::{__compiler, SharedTile, SmemLayout}; +use crate::copy::tma::TensorMap; +use core::marker::PhantomData; +use ctm_rs::mbarrier::Barrier; + +/// A rank-three TMA copy into the same physical layout used by MMA. +pub struct TmaCopy( + #[doc = "Retains the shared layout that must match the tensor-map descriptor."] PhantomData, +); +impl TmaCopy { + /// Copy a physical row-major global tile into a typed shared tile. + /// # Safety + /// All producer-warp lanes participate. Coordinates are in physical + /// global rows/columns; columns must be a multiple of `L::SPAN`. The leader + /// barrier must expect copies from the complete group. For group two, + /// rank is the local CTA rank and completion routes to CTA zero. + #[inline(always)] + pub unsafe fn load( + desc: &TensorMap, + dst: &mut SharedTile, + row: i32, + col: i32, + barrier: *mut Barrier, + rank: u32, + ) { + unsafe { + __compiler::tma_load::( + desc as *const _, + dst.as_ptr(), + col % L::SPAN as i32, + row, + col / L::SPAN as i32, + barrier, + rank, + ) + } + } + /// # Safety + /// One elected lane issues after stores to shared memory are visible to + /// the async proxy. Storage stays live through the bulk-copy wait. + #[inline(always)] + pub unsafe fn store(desc: &TensorMap, src: &SharedTile, row: i32, col: i32) { + unsafe { + __compiler::tma_store::( + desc as *const _, + src.as_ptr(), + col % L::SPAN as i32, + row, + col / L::SPAN as i32, + ) + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tmem.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tmem.rs new file mode 100644 index 0000000000..5717ea27ba --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/attention/sm100/tmem.rs @@ -0,0 +1,196 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Warp-collective register fragments for CTM tensor-memory loads and stores. +//! Conversion and chunking preserve the explicit completion choices of each API. + +use ctm_rs::{memory::CuSimd, tcgen05::tmem, thread}; + +/// Warp-collective 32-data-path TMEM register fragment, one row per lane. +pub struct TmemFragment { + /// This lane's FP32 register values. + pub values: [f32; REGS], +} +/// Implement one native TMEM register width with explicit load/store completion variants. +macro_rules! fragment { + ($n:literal,$load:ident,$store:ident) => { + impl TmemFragment<$n> { + /// Load a fragment and wait for the warp's preceding TMEM loads. + /// # Safety + /// All 32 lanes must execute with the same live, row-group-aligned + /// TMEM address after the producer's completion and thread fence. + #[inline(always)] + pub unsafe fn load(address: u32) -> Self { + let result = unsafe { Self::load_deferred(address) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::load_wait() }; + result + } + /// Load without a completion wait, retaining register dependencies. + /// # Safety + /// Same participation and producer obligations as `load`. Only true + /// register dependencies order subsequent operations. The caller + /// must wait before memory anti-dependencies or a consumer-release + /// handoff, and apply the required before-thread-sync fence. + /// See the [PTX register-dependency rule](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-consistency-model-canonical-sync-patterns-reg-dependency-same-thread). + #[inline(always)] + pub unsafe fn load_deferred(address: u32) -> Self { + let raw = unsafe { tmem::$load(address) }.to_array(); + let mut values = [0.0; $n]; + let mut i = 0; + while i < $n { + thread::__unroll_config::<0>(); + values[i] = f32::from_bits(raw[i]); + i += 1; + } + Self { values } + } + /// Store a fragment and wait for the warp's preceding TMEM stores. + /// # Safety + /// All 32 lanes must participate with the same live TMEM address. + /// Complete consumer handoffs only after the returned store wait. + #[inline(always)] + pub unsafe fn store(&self, address: u32) { + unsafe { self.store_deferred(address) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::store_wait() }; + } + /// Store without waiting, allowing disjoint TMEM chunks to overlap. + /// # Safety + /// Same participation and lifetime requirements as `store`. The + /// caller must wait for stores before conflicting memory operations + /// or completion handoffs, with the required thread-sync fence. + #[inline(always)] + pub unsafe fn store_deferred(&self, address: u32) { + let mut raw = [0u32; $n]; + let mut i = 0; + while i < $n { + thread::__unroll_config::<0>(); + raw[i] = self.values[i].to_bits(); + i += 1; + } + unsafe { Self::store_packed_deferred(address, raw) } + } + /// Store packed words unchanged (low FP16 lane in the low 16 bits). + /// # Safety + /// Same collective and lifetime requirements as `store`. + #[inline(always)] + pub unsafe fn store_packed(address: u32, raw: [u32; $n]) { + unsafe { Self::store_packed_deferred(address, raw) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::store_wait() }; + } + /// Store packed words without waiting for completion. + /// # Safety + /// Same participation and explicit completion obligations as + /// `store_deferred`. + #[inline(always)] + pub unsafe fn store_packed_deferred(address: u32, raw: [u32; $n]) { + unsafe { tmem::$store(address, CuSimd::new(raw)) }; + } + } + }; +} +fragment!(2, load_32x32b_x2, store_32x32b_x2); +fragment!(4, load_32x32b_x4, store_32x32b_x4); +fragment!(8, load_32x32b_x8, store_32x32b_x8); +fragment!(16, load_32x32b_x16, store_32x32b_x16); +fragment!(32, load_32x32b_x32, store_32x32b_x32); +// The attention references transfer larger logical fragments as repeated x32 +// instructions. Wide native transfers pass standalone probes but fail in the +// full attention compilation context with the pinned compiler. +/// Build a wider logical fragment from fixed x32 TMEM instruction chunks. +macro_rules! chunked_fragment { + ($n:literal; $($offset:literal => $chunk:ident),+) => { + impl TmemFragment<$n> { + /// Load the logical fragment in x32 chunks with one completion wait. + /// # Safety + /// All lanes participate with the same live row-group-aligned TMEM + /// address after producer completion and the thread-sync fence. + #[inline(always)] + pub unsafe fn load(address: u32) -> Self { + let result = unsafe { Self::load_deferred(address) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::load_wait() }; + result + } + /// Load x32 chunks without a completion wait. + /// # Safety + /// Same participation and producer requirements as `load`. Only true + /// register dependencies provide ordering. Wait before memory + /// anti-dependencies or a consumer-release handoff, and apply the + /// required before-thread-sync fence. + #[inline(always)] + pub unsafe fn load_deferred(address: u32) -> Self { + $(let $chunk = unsafe { tmem::load_32x32b_x32(address + $offset) };)+ + let mut values = [0.0; $n]; + $( + let raw = $chunk.to_array(); + let mut i = 0; + while i < 32 { + thread::__unroll_config::<0>(); + values[$offset + i] = f32::from_bits(raw[i]); + i += 1; + } + )+ + Self { values } + } + /// Store the logical fragment in x32 chunks with one completion wait. + /// # Safety + /// All lanes participate with the same live TMEM address. Publish + /// completion only after a before-thread-sync fence by every lane. + #[inline(always)] + pub unsafe fn store(&self, address: u32) { + unsafe { self.store_deferred(address) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::store_wait() }; + } + /// Store x32 chunks without waiting for completion. + /// # Safety + /// Same participation and lifetime requirements as `store`. Wait + /// before conflicting memory operations or completion handoffs, + /// with the required thread-sync fence. + #[inline(always)] + pub unsafe fn store_deferred(&self, address: u32) { + let mut raw = [0; $n]; + let mut i = 0; + while i < $n { + thread::__unroll_config::<0>(); + raw[i] = self.values[i].to_bits(); + i += 1; + } + unsafe { Self::store_packed_deferred(address, raw) } + } + /// Store packed FP16 words in x32 chunks, preserving every bit. + /// # Safety + /// Same participation, lifetime, and fence obligations as `store`. + #[inline(always)] + pub unsafe fn store_packed(address: u32, raw: [u32; $n]) { + unsafe { Self::store_packed_deferred(address, raw) }; + // SAFETY: the caller supplies the full-warp participation contract. + unsafe { tmem::store_wait() }; + } + /// Store packed x32 chunks without waiting for completion. + /// # Safety + /// Same participation and explicit completion obligations as + /// `store_deferred`. + #[inline(always)] + pub unsafe fn store_packed_deferred(address: u32, raw: [u32; $n]) { + $( + let mut chunk = [0; 32]; + let mut i = 0; + while i < 32 { + thread::__unroll_config::<0>(); + chunk[i] = raw[$offset + i]; + i += 1; + } + unsafe { tmem::store_32x32b_x32(address + $offset, CuSimd::new(chunk)) }; + )+ + } + } + } +} +chunked_fragment!(64; 0 => a, 32 => b); +chunked_fragment!(128; 0 => a, 32 => b, 64 => c, 96 => d); diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/__compiler.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/__compiler.rs new file mode 100644 index 0000000000..6e712768f7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/__compiler.rs @@ -0,0 +1,259 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Exact compiler boundaries for block-scaled GEMV. + +use super::*; +use crate::copy::vector::{assume, load_tile}; +use crate::numeric::{E4M3x2, Float4E2M1FN, Float4E2M1FNx2, Float8E4M3FN, Float8E8M0FNU, UE8M0x2}; +use crate::tensor::Tensor; +use core::marker::PhantomData; + +#[inline(never)] +/// Compiler boundary constructing packed value and canonical scale views for K16 scaling. +pub(super) const fn block_scaled_make<'a, Mode: Copy>( + values: &'a [u8], + scales: &'a [u8], + rows: usize, + k: usize, +) -> BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, 16, KMajor> { + let scale_layout = Sm1xxBlockScaleKMajor { + rest_m: rows.div_ceil(CANONICAL_SCALE_M), + rest_k: k.div_ceil(16 * CANONICAL_SCALE_K_GROUPS), + }; + BlockScaledTensor { + values: Tensor::from_storage( + values, + KMajor { + rows, + k, + mode: PhantomData, + }, + ), + scales: Tensor::from_storage(scales, scale_layout), + } +} + +#[inline(never)] +/// Compiler boundary selecting a batch row and computing its value/scale byte bases. +pub(super) fn block_scaled_thread_row<'a, Mode: Copy>( + tensor: BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, 16, KMajor>, + batch: usize, + row: usize, +) -> BlockScaledThreadRow<'a, Mode> { + let packed_k = tensor.values.layout.k / 2; + let value_row_base = (batch * tensor.values.layout.rows + row) * packed_k; + let scale_row_base = sm1xx_scale_offset!(tensor.scales.layout, batch, row, 0); + + BlockScaledThreadRow { + values: tensor.values.storage, + scales: tensor.scales.storage, + value_row_base, + scale_row_base, + mode: PhantomData, + } +} + +#[inline(never)] +/// Compiler boundary selecting one K64 tile: 32 packed bytes and four scale bytes. +pub(super) fn block_scaled_k_tile<'a, Mode: Copy>( + row: BlockScaledThreadRow<'a, Mode>, + tile: usize, +) -> BlockScaledTile64<'a, Mode> { + BlockScaledTile64 { + values: row.values, + scales: row.scales, + value_base: row.value_row_base + tile * 32, + scale_base: row.scale_row_base + tile * CANONICAL_SCALE_ATOM_BYTES, + mode: PhantomData, + } +} + +#[inline(never)] +/// Load one K64 tile and decode its four scales. +/// +/// # Safety +/// The value and scale ranges must be live, in bounds and aligned as required by the tile’s public load operation. +pub(super) unsafe fn block_scaled_load_k64<'a, Mode: Copy>( + tile: BlockScaledTile64<'a, Mode>, +) -> LoadedBlockScaledTile64 { + // State each alignment fact immediately before its load. The device + // compiler reads that direct link; putting the value in a struct would + // hide it. + let value_base = unsafe { assume::<16>(tile.value_base) }; + let scale_base = unsafe { assume::<4>(tile.scale_base) }; + let values_lo = unsafe { load_tile::(tile.values, value_base) }; + let values_hi = unsafe { load_tile::(tile.values, value_base + 16) }; + let scale_bytes = unsafe { load_tile::(tile.scales, scale_base) }; + + let (scale0, scale1) = UE8M0x2::from_bytes(scale_bytes[0], scale_bytes[1]).to_f32x2(); + let (scale2, scale3) = UE8M0x2::from_bytes(scale_bytes[2], scale_bytes[3]).to_f32x2(); + + LoadedBlockScaledTile64 { + values_lo, + values_hi, + scale0, + scale1, + scale2, + scale3, + mode: PhantomData, + } +} + +#[inline(never)] +/// Accumulate all 64 scaled FP4 products in fixed K-group order. +pub(super) fn block_scaled_dot_k64( + lhs: LoadedBlockScaledTile64, + rhs: LoadedBlockScaledTile64, + mut acc: f32, +) -> f32 { + /// Expand one K16 scale group into scalar FP4-pair products in fixed order. + macro_rules! accumulate_group { + ($group:literal; $($pair:literal),+ $(,)?) => {{ + let scale_a = lhs.scale::<$group>(); + let scale_b = rhs.scale::<$group>(); + $( + { + let (a_low, a_high) = lhs.value_pair::<$pair>().to_f32x2(); + let (b_low, b_high) = rhs.value_pair::<$pair>().to_f32x2(); + acc += ((a_low * scale_a) * b_low) * scale_b; + acc += ((a_high * scale_a) * b_high) * scale_b; + } + )+ + }}; + } + + accumulate_group!(0; 0, 1, 2, 3, 4, 5, 6, 7); + accumulate_group!(1; 8, 9, 10, 11, 12, 13, 14, 15); + accumulate_group!(2; 16, 17, 18, 19, 20, 21, 22, 23); + accumulate_group!(3; 24, 25, 26, 27, 28, 29, 30, 31); + acc +} + +#[inline(never)] +/// Compiler boundary constructing packed E2M1 value and canonical E4M3 scale views for K16 scaling. +pub(super) const fn nvfp4_make<'a, Mode: Copy>( + values: &'a [u8], + scales: &'a [u8], + rows: usize, + k: usize, +) -> BlockScaledTensor<'a, Float4E2M1FN, Float8E4M3FN, 16, KMajor> { + let scale_layout = Sm1xxBlockScaleKMajor { + rest_m: rows.div_ceil(CANONICAL_SCALE_M), + rest_k: k.div_ceil(16 * CANONICAL_SCALE_K_GROUPS), + }; + BlockScaledTensor { + values: Tensor::from_storage( + values, + KMajor { + rows, + k, + mode: PhantomData, + }, + ), + scales: Tensor::from_storage(scales, scale_layout), + } +} + +#[inline(never)] +/// Compiler boundary selecting an NVFP4 batch row and computing its value/scale byte bases. +pub(super) const fn nvfp4_thread_row<'a, Mode: Copy>( + tensor: BlockScaledTensor<'a, Float4E2M1FN, Float8E4M3FN, 16, KMajor>, + batch: usize, + row: usize, +) -> Nvfp4ThreadRow<'a, Mode> { + let packed_k = tensor.values.layout.k / 2; + let value_row_base = (batch * tensor.values.layout.rows + row) * packed_k; + let scale_row_base = sm1xx_scale_offset!(tensor.scales.layout, batch, row, 0); + Nvfp4ThreadRow { + values: tensor.values.storage, + scales: tensor.scales.storage, + value_row_base, + scale_row_base, + mode: PhantomData, + } +} + +#[inline(never)] +/// Compiler boundary selecting one K64 NVFP4 tile: 32 packed bytes and four E4M3 scale bytes. +pub(super) const fn nvfp4_k_tile<'a, Mode: Copy>( + row: Nvfp4ThreadRow<'a, Mode>, + tile: usize, +) -> Nvfp4Tile64<'a, Mode> { + Nvfp4Tile64 { + values: row.values, + scales: row.scales, + value_base: row.value_row_base + tile * 32, + scale_base: row.scale_row_base + tile * CANONICAL_SCALE_ATOM_BYTES, + mode: PhantomData, + } +} + +#[inline(never)] +/// Load one K64 NVFP4 tile and decode its four E4M3 scales. +pub(super) unsafe fn nvfp4_load_k64( + tile: Nvfp4Tile64<'_, Mode>, +) -> LoadedNvfp4Tile64 { + let value_base = unsafe { assume::<16>(tile.value_base) }; + let scale_base = unsafe { assume::<4>(tile.scale_base) }; + let values_lo = unsafe { load_tile::(tile.values, value_base) }; + let values_hi = unsafe { load_tile::(tile.values, value_base + 16) }; + let scale_bytes = unsafe { load_tile::(tile.scales, scale_base) }; + let (scale0, scale1) = E4M3x2::from_bytes(scale_bytes[0], scale_bytes[1]).to_f32x2(); + let (scale2, scale3) = E4M3x2::from_bytes(scale_bytes[2], scale_bytes[3]).to_f32x2(); + LoadedNvfp4Tile64 { + values_lo, + values_hi, + scale0, + scale1, + scale2, + scale3, + mode: PhantomData, + } +} + +#[inline(never)] +/// Accumulate all 64 E4M3-scaled FP4 products in fixed K-group order. +pub(super) fn nvfp4_dot_k64( + lhs: LoadedNvfp4Tile64, + rhs: LoadedNvfp4Tile64, + mut acc: f32, +) -> f32 { + let mut scale_group = |a_bytes: &[u8], b_bytes: &[u8], scale_a: f32, scale_b: f32| { + let mut i = 0usize; + while i < 8 { + let (a0, a1) = Float4E2M1FNx2::from_bits(a_bytes[i]).to_f32x2(); + let (b0, b1) = Float4E2M1FNx2::from_bits(b_bytes[i]).to_f32x2(); + acc += ((a0 * scale_a) * b0) * scale_b; + acc += ((a1 * scale_a) * b1) * scale_b; + i += 1; + } + }; + scale_group( + &lhs.values_lo[0..8], + &rhs.values_lo[0..8], + lhs.scale0, + rhs.scale0, + ); + scale_group( + &lhs.values_lo[8..16], + &rhs.values_lo[8..16], + lhs.scale1, + rhs.scale1, + ); + scale_group( + &lhs.values_hi[0..8], + &rhs.values_hi[0..8], + lhs.scale2, + rhs.scale2, + ); + scale_group( + &lhs.values_hi[8..16], + &rhs.values_hi[8..16], + lhs.scale3, + rhs.scale3, + ); + acc +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/gemv.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/gemv.rs new file mode 100644 index 0000000000..9f0707f0f5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/gemv.rs @@ -0,0 +1,233 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Per-thread GEMV rows, K64 tiles, and loaded register values. + +use super::{__compiler, Mkl, Nkl}; +use crate::numeric::Float4E2M1FNx2; +use core::marker::PhantomData; + +/// One thread's view of a row whose K values are adjacent. +/// +/// It stores slice references plus the first value and scale offsets. It does +/// not load the row. +#[derive(Clone, Copy, Debug)] +pub struct BlockScaledThreadRow<'a, Mode> { + /// Packed FP4 value storage, two logical values per byte. + pub(super) values: &'a [u8], + /// UE8M0 scale bytes in canonical Blackwell block order. + pub(super) scales: &'a [u8], + /// Byte offset of this row’s first packed value. + pub(super) value_row_base: usize, + /// Byte offset of this row’s first scale group. + pub(super) scale_row_base: usize, + /// Records whether the row belongs to the M or N operand. + pub(super) mode: PhantomData, +} + +impl<'a, Mode: Copy> BlockScaledThreadRow<'a, Mode> { + /// Select 64 neighboring K values and their four scales. + /// + /// ```text + /// row ── k_tile(0) ─► K 0..63 + /// └─ k_tile(1) ─► K 64..127 + /// ``` + /// + /// This calculates offsets only. Call [`BlockScaledTile64::load`] to read + /// global memory. + #[must_use] + #[inline(always)] + pub fn k_tile(self, tile: usize) -> BlockScaledTile64<'a, Mode> { + __compiler::block_scaled_k_tile(self, tile) + } +} + +/// A view of 64 packed FP4 values and four UE8M0 scales. +/// +/// Each scale covers 16 values. The view stores offsets, not loaded values. +#[derive(Clone, Copy, Debug)] +pub struct BlockScaledTile64<'a, Mode> { + /// Borrowed packed value bytes for the complete tensor. + pub(super) values: &'a [u8], + /// Borrowed scale bytes in canonical Blackwell block order. + pub(super) scales: &'a [u8], + /// Byte offset of the selected 64-value tile. + pub(super) value_base: usize, + /// Byte offset of the tile’s four adjacent scale bytes. + pub(super) scale_base: usize, + /// Retains the operand coordinate convention without runtime data. + pub(super) mode: PhantomData, +} + +impl<'a, Mode: Copy> BlockScaledTile64<'a, Mode> { + /// Load the packed values and scales into this thread's registers. + /// + /// Each scale is converted to `f32` once and then reused by its group of 16 + /// FP4 values. + /// + /// # Safety + /// + /// - `values[value_base..value_base + 32]` must be readable. + /// - `values[value_base]` must start at a 16-byte-aligned address. + /// - `scales[scale_base..scale_base + 4]` must be readable. + /// - `scales[scale_base]` must start at a 4-byte-aligned address. + #[must_use] + #[inline(always)] + pub unsafe fn load(self) -> LoadedBlockScaledTile64 { + unsafe { __compiler::block_scaled_load_k64(self) } + } +} + +/// One loaded K=64 tile held in a thread's registers. +/// +/// The 64 FP4 values remain packed as 32 bytes. The four scales have already +/// been converted to `f32`. +#[derive(Clone, Copy, Debug)] +pub struct LoadedBlockScaledTile64 { + /// First 32 FP4 values, packed two per byte. + pub(super) values_lo: [u8; 16], + /// Last 32 FP4 values, packed two per byte. + pub(super) values_hi: [u8; 16], + /// Decoded scale for logical K positions 0 through 15. + pub(super) scale0: f32, + /// Decoded scale for logical K positions 16 through 31. + pub(super) scale1: f32, + /// Decoded scale for logical K positions 32 through 47. + pub(super) scale2: f32, + /// Decoded scale for logical K positions 48 through 63. + pub(super) scale3: f32, + /// Retains the operand role while values reside in registers. + pub(super) mode: PhantomData, +} + +impl LoadedBlockScaledTile64 { + /// Return packed FP4 pair `PAIR` in increasing K order. + #[must_use] + #[inline(always)] + pub fn value_pair(self) -> Float4E2M1FNx2 { + const { assert!(PAIR < 32) }; + let bits = if PAIR < 16 { + self.values_lo[PAIR] + } else { + self.values_hi[PAIR - 16] + }; + Float4E2M1FNx2::from_bits(bits) + } + + /// Return the `f32` scale for one of the four 16-value groups. + #[must_use] + #[inline(always)] + pub fn scale(self) -> f32 { + const { assert!(GROUP < 4) }; + match GROUP { + 0 => self.scale0, + 1 => self.scale1, + 2 => self.scale2, + _ => self.scale3, + } + } +} + +impl LoadedBlockScaledTile64 { + /// Add the dot product of two K=64 tiles to `acc`. + /// + /// The types require matrix data ([`Mkl`]) on the left and vector data + /// ([`Nkl`]) on the right: + /// + /// ```text + /// acc + sum((A[k] × A_scale[k]) × B[k] × B_scale[k]) + /// ``` + /// + /// Each scale is reused for 16 values. Each packed byte contributes two + /// FP4 values in increasing K order. + #[must_use] + #[inline(always)] + pub fn dot_accumulate(self, rhs: LoadedBlockScaledTile64, acc: f32) -> f32 { + __compiler::block_scaled_dot_k64(self, rhs, acc) + } +} + +/// One thread's view of an NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) row. +/// +/// It stores slice references plus the row's starting value and scale offsets. +#[derive(Clone, Copy, Debug)] +pub struct Nvfp4ThreadRow<'a, Mode> { + /// Packed FP4 value storage, two logical values per byte. + pub(super) values: &'a [u8], + /// E4M3 scale bytes in canonical Blackwell block order. + pub(super) scales: &'a [u8], + /// Byte offset of this row's first packed value. + pub(super) value_row_base: usize, + /// Byte offset of this row's first scale group. + pub(super) scale_row_base: usize, + /// Records whether the row belongs to the M or N operand. + pub(super) mode: PhantomData, +} + +impl<'a, Mode: Copy> Nvfp4ThreadRow<'a, Mode> { + /// Select 64 neighboring K values and their four `Float8E4M3FN` scales. + #[must_use] + #[inline(always)] + pub const fn k_tile(self, tile: usize) -> Nvfp4Tile64<'a, Mode> { + __compiler::nvfp4_k_tile(self, tile) + } +} + +/// A view of 64 packed `Float4E2M1FN` values and four `Float8E4M3FN` scales (`SF_VEC = 16`). +#[derive(Clone, Copy, Debug)] +pub struct Nvfp4Tile64<'a, Mode> { + /// Borrowed packed value bytes for the complete tensor. + pub(super) values: &'a [u8], + /// Borrowed E4M3 scale bytes in canonical Blackwell block order. + pub(super) scales: &'a [u8], + /// Byte offset of the selected 64-value tile. + pub(super) value_base: usize, + /// Byte offset of the tile's four adjacent E4M3 scale bytes. + pub(super) scale_base: usize, + /// Retains the operand coordinate convention without runtime data. + pub(super) mode: PhantomData, +} + +impl<'a, Mode: Copy> Nvfp4Tile64<'a, Mode> { + /// Load the 64 packed `Float4E2M1FN` values and decode the four `Float8E4M3FN` scales to `f32`. + /// + /// # Safety + /// + /// - `values[value_base..value_base + 32]` must be readable and 16-byte aligned. + /// - `scales[scale_base..scale_base + 4]` must be readable and 4-byte aligned. + #[must_use] + #[inline(always)] + pub unsafe fn load(self) -> LoadedNvfp4Tile64 { + unsafe { __compiler::nvfp4_load_k64(self) } + } +} + +/// 64 packed `Float4E2M1FN` values and four decoded `Float8E4M3FN` (`f32`) scales for one K=64 tile. +#[derive(Clone, Copy, Debug)] +pub struct LoadedNvfp4Tile64 { + /// First 16 packed bytes, covering K values 0..31. + pub(super) values_lo: [u8; 16], + /// Second 16 packed bytes, covering K values 32..63. + pub(super) values_hi: [u8; 16], + /// Decoded E4M3 scale for K values 0..15. + pub(super) scale0: f32, + /// Decoded E4M3 scale for K values 16..31. + pub(super) scale1: f32, + /// Decoded E4M3 scale for K values 32..47. + pub(super) scale2: f32, + /// Decoded E4M3 scale for K values 48..63. + pub(super) scale3: f32, + /// Retains the operand coordinate convention without runtime data. + pub(super) mode: PhantomData, +} + +impl LoadedNvfp4Tile64 { + /// Add the dot product of two K=64 NVFP4 tiles (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) to `acc`. + #[must_use] + #[inline(always)] + pub fn dot_accumulate(self, rhs: LoadedNvfp4Tile64, acc: f32) -> f32 { + __compiler::nvfp4_dot_k64(self, rhs, acc) + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/layouts.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/layouts.rs new file mode 100644 index 0000000000..86b40e2bc4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/layouts.rs @@ -0,0 +1,221 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Value and scale coordinate layouts for block-scaled tensors. + +use crate::copy::tiled::SharedTensor; +use crate::numeric::UE8M0x4; +use core::marker::PhantomData; + +/// Rows whose K values are adjacent in memory. +/// +/// ```text +/// row 0: [k0 k1 k2 ...] +/// row 1: [k0 k1 k2 ...] +/// ``` +/// +/// `Mode` is either [`Mkl`](super::Mkl) or [`Nkl`](super::Nkl). It prevents A rows and B rows from +/// being swapped even though both have the same memory layout. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct KMajor { + /// Logical M or N rows in each batch. + pub(super) rows: usize, + /// Logical FP4 values along one row’s K dimension. + pub(super) k: usize, + /// Selects the M/K/L or N/K/L coordinate convention. + pub(super) mode: PhantomData, +} + +impl KMajor { + /// Describe `rows` rows, each containing `k` adjacent logical values. + #[must_use] + #[inline(always)] + pub const fn new(rows: usize, k: usize) -> Self { + Self { + rows, + k, + mode: PhantomData, + } + } + + /// Return the number of rows in each batch: `M` or `N`. + #[must_use] + #[inline(always)] + pub const fn rows(self) -> usize { + self.rows + } + + /// Return the number of logical K values in each row. + #[must_use] + #[inline(always)] + pub const fn k(self) -> usize { + self.k + } +} + +/// Scales stored one row after another with no padding. +/// +/// One scale covers `SF_VEC` neighboring K values. For `SF_VEC = 4`: +/// +/// ```text +/// values: [0 1 2 3] [4 5 6 7] +/// scales: [ s0 ] [ s1 ] +/// memory: row 0 scales, then row 1 scales, ... +/// ``` +/// +/// [`Sm1xxBlockScaleKMajor`] stores the same logical scales in Blackwell's +/// padded hardware layout instead. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DenseBlockScaleKMajor { + /// Logical M or N rows in the scale matrix. + pub(super) rows: usize, + /// Scale groups per row, rounded up to include a partial final group. + pub(super) groups: usize, + /// Carries the operand’s row-coordinate convention. + mode: PhantomData, +} + +impl DenseBlockScaleKMajor { + /// Describe scales for `rows` rows of `k` logical values. + #[must_use] + #[inline(always)] + pub const fn new(rows: usize, k: usize) -> Self { + const { assert!(SF_VEC > 0) }; + Self { + rows, + groups: k.div_ceil(SF_VEC), + mode: PhantomData, + } + } + + /// Map the coordinate `(row, K group)` to its byte index (CuTeDSL + /// `cute.crd2idx` on this layout). + #[must_use] + #[inline(always)] + pub const fn crd2idx(&self, row: usize, group: usize) -> usize { + row * self.groups + group + } + + /// Return the number of scale groups in each row. + #[must_use] + #[inline(always)] + pub const fn groups(self) -> usize { + self.groups + } + + /// Return the total number of scale bytes. + #[must_use] + #[inline(always)] + pub const fn storage_len(self) -> usize { + self.rows * self.groups + } +} + +/// Number of rows in one Blackwell scale block. +pub const CANONICAL_SCALE_M: usize = 128; +/// Number of neighboring K scale groups in one Blackwell scale block. +pub const CANONICAL_SCALE_K_GROUPS: usize = 4; +/// Bytes occupied by one Blackwell scale block. +pub const CANONICAL_SCALE_ATOM_BYTES: usize = 512; + +/// Shared-memory layout for one Blackwell scale block with `SF = 32`. +/// +/// The block contains 128 rows × 4 scale bytes = 512 bytes. TMA moves all 512 +/// bytes without interpreting them. Each row's four scales form one `u32`: +/// +/// ```text +/// one row: [scale 0][scale 1][scale 2][scale 3] +/// +/// physical word order: +/// row 0, row 32, row 64, row 96, +/// row 1, row 33, row 65, row 97, ... +/// ``` +/// +/// This hardware order differs from the simple row-after-row order in +/// [`DenseBlockScaleKMajor`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Sm120ScaleAtom; + +impl Sm120ScaleAtom { + /// Number of logical rows in the block. + pub const ROWS: usize = CANONICAL_SCALE_M; + /// Number of `u32` words in the block; one word per row. + pub const WORDS: usize = CANONICAL_SCALE_ATOM_BYTES / size_of::(); + /// Number of bytes copied by TMA. + pub const BYTES: usize = CANONICAL_SCALE_ATOM_BYTES; + + /// Return the `u32` word holding one logical row's scales. + /// + /// The layout repeats every 128 rows, so only the lowest seven row bits are + /// used. + #[must_use] + #[inline(always)] + pub const fn word_offset(row: usize) -> usize { + (row & 31) * 4 + ((row >> 5) & 3) + } +} + +/// A typed view of one Blackwell scale block in shared memory. +/// +/// Creating this view with [`SharedTensor::from_raw_parts`] does not allocate, +/// convert, or copy memory. Its address must be aligned for `u32`, even when +/// TMA wrote the same bytes using `u16` storage boxes. `Role` keeps A and B +/// scale blocks distinct. +pub type SharedScaleAtom = SharedTensor; + +/// Blackwell's required memory layout for block scales. +/// +/// Logically, each `(batch, row, K group)` points to one byte. Physically, +/// scales are arranged in 512-byte blocks: +/// +/// ```text +/// logical scales +/// │ +/// ├─ rows in groups of 128 +/// └─ K groups in groups of 4 +/// │ +/// ▼ +/// 512-byte scale blocks +/// ``` +/// +/// Partial row or K groups still reserve a complete block. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Sm1xxBlockScaleKMajor { + /// Number of complete or padded 128-row blocks per batch. + pub(super) rest_m: usize, + /// Number of complete or padded four-scale-group blocks along K. + pub(super) rest_k: usize, +} + +impl Sm1xxBlockScaleKMajor { + /// Describe Blackwell-formatted scales for `rows × k` logical values. + #[must_use] + #[inline(always)] + pub const fn new(rows: usize, k: usize) -> Self { + assert!(SF_VEC > 0, "scale-vector size must be positive"); + Self { + rest_m: rows.div_ceil(CANONICAL_SCALE_M), + rest_k: k.div_ceil(SF_VEC * CANONICAL_SCALE_K_GROUPS), + } + } + + /// Map the coordinate `(batch, row, K group)` to its byte index in the + /// 512-byte block order (CuTeDSL `cute.crd2idx` on this layout). + #[must_use] + #[inline(always)] + pub const fn crd2idx(self, batch: usize, row: usize, group: usize) -> usize { + sm1xx_scale_offset!(self, batch, row, group) + } + + /// Return the required byte count, including complete padded blocks. + #[must_use] + #[inline(always)] + pub const fn storage_len(self, batches: usize) -> usize { + batches * self.rest_m * self.rest_k * CANONICAL_SCALE_ATOM_BYTES + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/mod.rs new file mode 100644 index 0000000000..a9ff2e98bc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/mod.rs @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor views for values that share scales in small blocks. +//! +//! Block scaling stores narrow values, such as FP4, plus one scale for each +//! short K group: +//! +//! ```text +//! values: [v0 v1 ... v15] [v16 ... v31] +//! scales: [ s0 ] [ s1 ] +//! result: value × its group's scale +//! ``` +//! +//! The views connect the value memory, scale memory, layout, and group size in +//! one Rust type: +//! +//! ```text +//! value slice + scale slice +//! │ +//! ▼ +//! BlockScaledTensor +//! ├── GEMV: row ─► K64 tile ─► registers +//! └── GEMM: global tile ─► shared tile ─► MMA +//! ``` +//! +//! Creating or reshaping a view does not allocate, copy, or convert data. The +//! small wrapper methods are inlined during device compilation. + +/// Marks an operand whose logical coordinates are `(M, K, L)`. +/// +/// `M` is the matrix row, `K` is the dot-product position, and `L` is batch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Mkl; + +/// Marks an operand whose logical coordinates are `(N, K, L)`. +/// +/// `N` is the output column, `K` is the dot-product position, and `L` is batch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Nkl; + +// Map one logical scale coordinate into Blackwell's 512-byte block order. +/// Map batch, row and K-group coordinates into the canonical 512-byte scale-block order. +macro_rules! sm1xx_scale_offset { + ($layout:expr, $batch:expr, $row:expr, $group:expr) => {{ + let layout = $layout; + let batch = $batch; + let row = $row; + let group = $group; + batch * layout.rest_m * layout.rest_k * CANONICAL_SCALE_ATOM_BYTES + + (row >> 7) * layout.rest_k * CANONICAL_SCALE_ATOM_BYTES + + (group >> 2) * CANONICAL_SCALE_ATOM_BYTES + + ((row & 31) << 4) + + (((row >> 5) & 3) << 2) + + (group & 3) + }}; +} + +// Compiler boundaries keep their original module paths. Storage fields are +// visible only within this family, where the boundary builds typed views. +mod gemv; +mod layouts; +mod scales; +mod tensor; + +pub use gemv::{ + BlockScaledThreadRow, BlockScaledTile64, LoadedBlockScaledTile64, LoadedNvfp4Tile64, + Nvfp4ThreadRow, Nvfp4Tile64, +}; +pub use layouts::{ + CANONICAL_SCALE_ATOM_BYTES, CANONICAL_SCALE_K_GROUPS, CANONICAL_SCALE_M, DenseBlockScaleKMajor, + KMajor, SharedScaleAtom, Sm1xxBlockScaleKMajor, Sm120ScaleAtom, +}; +pub use scales::{MmaScalePair, MmaScaleQuad, ScalePack4}; +pub use tensor::{BlockScaledTensor, Mxf4BlockScaledTensor, Nvfp4BlockScaledTensor}; + +/// Stable calls recognized by the CuTe importer. +mod __compiler; +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/scales.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/scales.rs new file mode 100644 index 0000000000..d277172842 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/scales.rs @@ -0,0 +1,138 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Packed scales passed to block-scaled MMA. + +use core::marker::PhantomData; + +/// The two [`Float8E8M0FNU`](crate::numeric::Float8E8M0FNU) scales used by one K=64 MXFP4 MMA. +/// +/// Each scale covers 32 K values: +/// +/// ```text +/// K: 0 ................ 31 | 32 ............... 63 +/// scale: s0 | s1 +/// ``` +/// +/// `Mode` keeps A scales ([`Mkl`](super::Mkl)) separate from B scales ([`Nkl`](super::Nkl)) so the MMA +/// call cannot swap them accidentally. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MmaScalePair { + /// Two UE8M0 scale bytes in the low 16 bits, in increasing K order. + pub(super) bits: u32, + /// Distinguishes A-side and B-side scale providers. + pub(super) mode: PhantomData, +} + +impl MmaScalePair { + /// Return the two packed bytes passed to the MMA instruction. + /// + /// Bits 7..0 scale K=0..31. Bits 15..8 scale K=32..63. + #[must_use] + #[inline(always)] + #[allow(dead_code)] + pub(crate) const fn bits(self) -> u32 { + self.bits + } +} + +/// Four neighboring K-group scales packed in one register. +/// +/// ```text +/// [s0][s1][s2][s3] +/// └ pair 0 ┘└ pair 1 ┘ +/// ``` +/// +/// [`Self::pair`] selects the two scales needed by one K=64 MMA. `Mode` keeps +/// A and B scale packs distinct. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ScalePack4 { + /// Four UE8M0 bytes in increasing K-group order. + bits: u32, + /// Distinguishes M-row and N-row scale packs. + mode: PhantomData, +} + +impl ScalePack4 { + /// Wrap four packed scale bytes without changing them. + #[inline(always)] + pub(crate) const fn from_bits(bits: u32) -> Self { + Self { + bits, + mode: PhantomData, + } + } + + #[cfg(test)] + /// Return the packed bytes for host-side layout tests. + #[inline(always)] + pub(crate) const fn bits(self) -> u32 { + self.bits + } + + /// Select scales 0 and 1 for `HALF = 0`, or 2 and 3 for `HALF = 1`. + #[must_use] + #[inline(always)] + pub const fn pair(self) -> MmaScalePair { + const { assert!(HALF < 2) }; + MmaScalePair { + bits: (self.bits >> (HALF * 16)) & 0xffff, + mode: PhantomData, + } + } + + /// Select either two-scale pair using a runtime index and no bounds check. + /// + /// # Safety + /// + /// `half` must be exactly `0` or `1`. + #[inline(always)] + pub(crate) const unsafe fn pair_at_unchecked(self, half: usize) -> MmaScalePair { + MmaScalePair { + bits: (self.bits >> (half * 16)) & 0xffff, + mode: PhantomData, + } + } +} + +/// Four [`Float8E4M3FN`](crate::numeric::Float8E4M3FN) scales used by one K=64 NVFP4 MMA (`SF_VEC = 16`). +/// +/// Each scale covers 16 consecutive K values: +/// +/// ```text +/// K: 0..15 | 16..31 | 32..47 | 48..63 +/// scale: s0 | s1 | s2 | s3 +/// ``` +/// +/// `Mode` keeps A scales ([`Mkl`](super::Mkl)) separate from B scales ([`Nkl`](super::Nkl)). +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MmaScaleQuad { + /// Four E4M3 scale bytes packed in increasing K-group order. + pub(super) bits: u32, + /// Distinguishes A-side and B-side scale providers. + pub(super) mode: PhantomData, +} + +impl MmaScaleQuad { + /// Wrap four packed E4M3 scale bytes in increasing K order. + #[must_use] + #[inline(always)] + pub const fn from_bits(bits: u32) -> Self { + Self { + bits, + mode: PhantomData, + } + } + + /// Return the 32-bit packed scale word passed to the NVFP4 MMA instruction. + #[must_use] + #[inline(always)] + pub const fn bits(self) -> u32 { + self.bits + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tensor.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tensor.rs new file mode 100644 index 0000000000..0bcebbe40c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tensor.rs @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Global tensor views joining packed values and scale storage. + +use super::{ + __compiler, BlockScaledThreadRow, CANONICAL_SCALE_K_GROUPS, DenseBlockScaleKMajor, KMajor, + MmaScalePair, MmaScaleQuad, Nvfp4ThreadRow, Sm1xxBlockScaleKMajor, +}; +use crate::copy::cooperative::GmemMatrix; +use crate::copy::tiled::GlobalCopyTensor; +use crate::markers::LeadingDim; +use crate::numeric::{E4M3x4, Float4E2M1FN, Float8E4M3FN, Float8E8M0FNU, Mxf4E2M1, UE8M0x4}; +use crate::tensor::{Tensor, TensorElement}; +use core::marker::PhantomData; + +/// Packed MXFP4 values and their dense UE8M0 scales for one batch (`L = 1`). +/// +/// ```text +/// logical K values: [fp4 fp4 fp4 fp4] ... 32 values ... +/// value storage: [ one f16 box ] +/// scale storage: [ one UE8M0 scale ] +/// ``` +/// +/// The `f16` is only a 16-bit storage box for four FP4 nibbles; it is not an +/// FP16 numeric value. Scales are stored densely, one byte per 32 K values. +#[derive(Clone, Copy, Debug)] +pub struct Mxf4BlockScaledTensor<'a, Mode: Copy> { + /// Packed value storage viewed as four FP4 values per f16 carrier. + values: Tensor<'a, Mxf4E2M1, KMajor, f16>, + /// One UE8M0 scale per logical K32 group, stored densely by row. + pub(super) scales: Tensor<'a, Float8E8M0FNU, DenseBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> Mxf4BlockScaledTensor<'a, Mode> { + /// View packed FP4 values and dense scale bytes as one block-scaled tensor. + /// + /// This does not allocate, copy, or validate memory. Later copy and load + /// operations still require valid buffer sizes and alignment. + #[must_use] + #[inline(always)] + pub const fn from_slices( + values: &'a [f16], + scales: &'a [u8], + rows: usize, + logical_k: usize, + ) -> Self { + Self { + values: Tensor::from_storage(values, KMajor::new(rows, logical_k)), + scales: Tensor::from_storage(scales, DenseBlockScaleKMajor::new(rows, logical_k)), + } + } + + /// Create the global-memory view used by a block-wide value copy. + /// + /// Logical K counts FP4 values. The returned physical view counts `f16` + /// storage boxes, so its width is `K / 4`: + /// + /// ```text + /// 128 FP4 values ──► 32 f16 storage boxes + /// ``` + #[must_use] + #[inline(always)] + pub fn values_for_copy(&self) -> GlobalCopyTensor, Mode> { + let carrier_k = self.values.layout.k / 4; + GlobalCopyTensor::from_carrier(GmemMatrix { + base: self.values.storage.as_ptr(), + rows: self.values.layout.rows, + cols: carrier_k, + leading_dim: LeadingDim::<8> { + elements: carrier_k, + }, + }) + } + + /// View each four neighboring scale bytes as one `u32` copy word. + /// + /// ```text + /// [s0][s1][s2][s3] ──► one u32 + /// ``` + /// + /// This does not copy or convert the bytes. The tiled-copy layout decides + /// which rows and K=128 sections are moved together. + /// + /// # Safety + /// + /// - The scale slice must contain at least `rows * groups` bytes. + /// - Its first address must be aligned for `u32`. + /// - Each row's number of scale groups must be divisible by four. + /// - The returned view must not outlive this scale slice. + #[must_use] + #[inline(always)] + pub unsafe fn scale_words_for_copy( + &self, + ) -> GlobalCopyTensor, Mode> { + let words = self.scales.layout.groups / 4; + GlobalCopyTensor::from_carrier(GmemMatrix { + base: self.scales.storage.as_ptr().cast::(), + rows: self.scales.layout.rows, + cols: words, + leading_dim: LeadingDim::<1> { elements: words }, + }) + } + + /// Load the two neighboring scale bytes used by one K=64 MMA. + /// + /// The result keeps `Mode`, so the MMA API can distinguish A scales from B + /// scales. + #[must_use] + #[inline(always)] + pub fn load_scale_pair(&self, row: usize, pair: usize) -> MmaScalePair { + let offset = self.scales.layout.crd2idx(row, pair * 2); + MmaScalePair { + bits: (self.scales.storage[offset] as u32) + | ((self.scales.storage[offset + 1] as u32) << 8), + mode: PhantomData, + } + } +} + +/// One view that keeps packed values and their scales together. +/// +/// The Rust type records the data format, scale format, number of values per +/// scale, and value layout. This prevents mixing incompatible pieces later. +#[derive(Clone, Copy, Debug)] +pub struct BlockScaledTensor< + 'a, + Data: TensorElement, + Scale: TensorElement, + const SF_VEC: usize, + DataLayout, +> { + /// Packed byte storage with the tensor’s logical value layout. + pub(super) values: Tensor<'a, Data, DataLayout, u8>, + /// Scale bytes arranged in canonical Blackwell scale blocks. + pub(super) scales: Tensor<'a, Scale, Sm1xxBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, 16, KMajor> { + /// View packed values and Blackwell-formatted scales as one tensor. + /// + /// This does not validate, allocate, convert, or copy memory. The later + /// unsafe load still requires correctly sized and aligned slices. + #[must_use] + #[inline(always)] + pub const fn from_slices(values: &'a [u8], scales: &'a [u8], rows: usize, k: usize) -> Self { + __compiler::block_scaled_make::(values, scales, rows, k) + } + + /// Return the number of complete K=64 tiles. + /// + /// Each tile has four groups of 16 values. `k` must be divisible by 64 + /// before any tile is loaded. + #[must_use] + #[inline(always)] + pub fn k_tile_count(self) -> usize { + self.values.layout.k / (16 * CANONICAL_SCALE_K_GROUPS) + } + + /// Select one matrix or vector row in `batch` for the current thread. + /// + /// This calculates starting offsets only; it does not read memory. + #[must_use] + #[inline(always)] + pub fn thread_row(self, batch: usize, row: usize) -> BlockScaledThreadRow<'a, Mode> { + __compiler::block_scaled_thread_row(self, batch, row) + } +} + +impl<'a, Mode: Copy> BlockScaledTensor<'a, Float4E2M1FN, Float8E4M3FN, 16, KMajor> { + /// View packed `Float4E2M1FN` values and canonical `Float8E4M3FN` (`SF_VEC = 16`) scales as one NVFP4 tensor. + #[must_use] + #[inline(always)] + pub const fn from_slices(values: &'a [u8], scales: &'a [u8], rows: usize, k: usize) -> Self { + __compiler::nvfp4_make(values, scales, rows, k) + } + + /// Return the number of complete K=64 tiles (`4` groups of 16 values per tile). + #[must_use] + #[inline(always)] + pub const fn k_tile_count(self) -> usize { + self.values.layout.k / (16 * CANONICAL_SCALE_K_GROUPS) + } + + /// Select one matrix or vector row in `batch` for the current thread. + #[must_use] + #[inline(always)] + pub const fn thread_row(self, batch: usize, row: usize) -> Nvfp4ThreadRow<'a, Mode> { + __compiler::nvfp4_thread_row(self, batch, row) + } +} + +/// Packed NVFP4 (`Float4E2M1FN`) values and their dense `Float8E4M3FN` scales (`SF_VEC = 16`, `L = 1`). +/// +/// ```text +/// logical K values: [fp4 fp4 fp4 fp4] ... 16 values ... +/// value storage: [ one f16 box ] (4 boxes per 16 values) +/// scale storage: [ one E4M3 scale byte ] +/// ``` +#[derive(Clone, Copy, Debug)] +pub struct Nvfp4BlockScaledTensor<'a, Mode: Copy> { + /// Packed value storage viewed as four FP4 values per f16 carrier. + values: Tensor<'a, Mxf4E2M1, KMajor, f16>, + /// One E4M3 scale per logical K16 group, stored densely by row. + pub(super) scales: Tensor<'a, Float8E4M3FN, DenseBlockScaleKMajor, u8>, +} + +impl<'a, Mode: Copy> Nvfp4BlockScaledTensor<'a, Mode> { + /// View packed FP4 values and dense E4M3 scale bytes as one NVFP4 block-scaled tensor. + #[must_use] + #[inline(always)] + pub const fn from_slices( + values: &'a [f16], + scales: &'a [u8], + rows: usize, + logical_k: usize, + ) -> Self { + Self { + values: Tensor::from_storage(values, KMajor::new(rows, logical_k)), + scales: Tensor::from_storage(scales, DenseBlockScaleKMajor::new(rows, logical_k)), + } + } + + /// Create the global-memory view used by a block-wide value copy. + #[must_use] + #[inline(always)] + pub fn values_for_copy(&self) -> GlobalCopyTensor, Mode> { + let carrier_k = self.values.layout.k / 4; + GlobalCopyTensor::from_carrier(GmemMatrix { + base: self.values.storage.as_ptr(), + rows: self.values.layout.rows, + cols: carrier_k, + leading_dim: LeadingDim::<8> { + elements: carrier_k, + }, + }) + } + + /// View each four neighboring E4M3 scale bytes (covering K=64) as one `u32` copy word. + /// + /// # Safety + /// + /// - The scale slice must contain at least `rows * groups` bytes and be `u32` aligned. + /// - Each row's number of scale groups must be divisible by four. + #[must_use] + #[inline(always)] + pub unsafe fn scale_words_for_copy( + &self, + ) -> GlobalCopyTensor, Mode> { + let words = self.scales.layout.groups / 4; + GlobalCopyTensor::from_carrier(GmemMatrix { + base: self.scales.storage.as_ptr().cast::(), + rows: self.scales.layout.rows, + cols: words, + leading_dim: LeadingDim::<1> { elements: words }, + }) + } + + /// Load the four neighboring `Float8E4M3FN` scale bytes used by one K=64 NVFP4 MMA. + #[must_use] + #[inline(always)] + pub fn load_scale_quad(&self, row: usize, k64_step: usize) -> MmaScaleQuad { + let offset = self.scales.layout.crd2idx(row, k64_step * 4); + let bits = (self.scales.storage[offset] as u32) + | ((self.scales.storage[offset + 1] as u32) << 8) + | ((self.scales.storage[offset + 2] as u32) << 16) + | ((self.scales.storage[offset + 3] as u32) << 24); + MmaScaleQuad::from_bits(bits) + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tests.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tests.rs new file mode 100644 index 0000000000..8f59153ea1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/block_scaled/tests.rs @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks packed-value carriers and dense versus Blackwell scale layouts."] +use super::*; +use crate::numeric::{Float4E2M1FN, Float8E8M0FNU}; + +/// Canonical scale coordinates must match known Blackwell offsets, including block boundaries. +#[test] +fn canonical_scale_layout_matches_blackwell_golden_offsets() { + let layout = Sm1xxBlockScaleKMajor::<16>::new(256, 128); + + assert_eq!(layout.crd2idx(0, 0, 0), 0); + assert_eq!(layout.crd2idx(0, 31, 0), 496); + assert_eq!(layout.crd2idx(0, 32, 0), 4); + assert_eq!(layout.crd2idx(0, 127, 3), 511); + assert_eq!(layout.crd2idx(0, 0, 4), 512); + assert_eq!(layout.crd2idx(0, 128, 0), 1_024); + assert_eq!(layout.crd2idx(1, 0, 0), 2_048); + assert_eq!(layout.storage_len(2), 4_096); +} + +/// M/K/L and N/K/L scale/value views must retain distinct type identities. +#[test] +fn matrix_and_vector_modes_remain_distinct_types() { + let values = [0u8; 128]; + let scales = [0u8; CANONICAL_SCALE_ATOM_BYTES]; + + let matrix = BlockScaledTensor::>::from_slices( + &values, &scales, 4, 64, + ); + let vector = BlockScaledTensor::>::from_slices( + &values, &scales, 1, 64, + ); + + assert_eq!(matrix.k_tile_count(), 1); + assert_eq!(vector.k_tile_count(), 1); + assert_eq!(matrix.thread_row(0, 3).k_tile(0).value_base, 96); + assert_eq!(vector.thread_row(0, 0).k_tile(0).value_base, 0); +} + +/// Logical K extents must survive conversion to packed value and scale copy carriers. +#[test] +fn mxf4_tensor_keeps_logical_k_and_materializes_packed_carriers() { + let values = [f16::from_bits(0); 128]; + let scales = [127u8; 16]; + let matrix = Mxf4BlockScaledTensor::::from_slices(&values, &scales, 4, 64); + + let copy = matrix.values_for_copy(); + assert_eq!(copy.carrier().rows, 4); + assert_eq!(copy.carrier().cols, 16); + assert_eq!(copy.carrier().leading_dim.elements, 16); + assert_eq!(matrix.scales.layout.groups(), 2); + assert_eq!(matrix.scales.layout.storage_len(), 8); + assert_eq!(matrix.load_scale_pair(3, 0).bits(), 0x7f7f); +} + +/// Packed scale words must retain low-byte-first K-group order when selecting either half. +#[test] +fn staged_scale_pack_preserves_little_endian_k_order() { + let pack = ScalePack4::::from_bits(0x8382_8180); + assert_eq!(pack.pair::<0>().bits(), 0x8180); + assert_eq!(pack.pair::<1>().bits(), 0x8382); + + /// Four-byte-aligned scale storage for testing u32 scale carriers. + #[repr(align(4))] + struct AlignedScales(#[doc = "Scale bytes arranged in increasing K-group order."] [u8; 32]); + let values = [f16::from_bits(0); 256]; + let scales = AlignedScales([127; 32]); + let matrix = Mxf4BlockScaledTensor::::from_slices(&values, &scales.0, 4, 256); + // Safe because the wrapper gives the bytes u32 alignment and each row + // has 8 groups, which can be split into groups of 4. + let copy = unsafe { matrix.scale_words_for_copy() }; + assert_eq!(copy.carrier().rows, 4); + assert_eq!(copy.carrier().cols, 2); + assert_eq!(copy.carrier().leading_dim.elements, 2); +} + +/// Dense row scales and padded Blackwell scale blocks must report their different storage requirements. +#[test] +fn dense_and_canonical_scale_layouts_are_distinct() { + let dense = DenseBlockScaleKMajor::::new(128, 128); + let canonical = Sm1xxBlockScaleKMajor::<32>::new(128, 128); + + assert_eq!(dense.crd2idx(32, 0), 128); + assert_eq!(dense.storage_len(), 512); + assert_eq!(canonical.crd2idx(0, 32, 0), 4); + assert_eq!(canonical.storage_len(1), 512); +} + +/// Per-row shared scale words must match the canonical 512-byte atom permutation. +#[test] +fn shared_scale_atom_words_match_blackwell_golden_layout() { + let canonical = Sm1xxBlockScaleKMajor::<32>::new(128, 128); + let mut seen = [false; Sm120ScaleAtom::WORDS]; + + for row in 0..Sm120ScaleAtom::ROWS { + let word = Sm120ScaleAtom::word_offset(row); + assert_eq!(word * 4, canonical.crd2idx(0, row, 0)); + for group in 0..CANONICAL_SCALE_K_GROUPS { + assert_eq!(word * 4 + group, canonical.crd2idx(0, row, group)); + } + assert!(!seen[word]); + seen[word] = true; + } + + assert!(seen.into_iter().all(|present| present)); + assert_eq!(Sm120ScaleAtom::BYTES, 512); + assert_eq!( + [ + Sm120ScaleAtom::word_offset(0), + Sm120ScaleAtom::word_offset(31), + Sm120ScaleAtom::word_offset(32), + Sm120ScaleAtom::word_offset(63), + Sm120ScaleAtom::word_offset(64), + Sm120ScaleAtom::word_offset(95), + Sm120ScaleAtom::word_offset(96), + Sm120ScaleAtom::word_offset(127), + ], + [0, 124, 1, 125, 2, 126, 3, 127] + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/cooperative.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/cooperative.rs new file mode 100644 index 0000000000..1aacdf55e6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/cooperative.rs @@ -0,0 +1,279 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Cooperative copies from global memory to shared memory. +//! +//! CUDA groups its workers like this: +//! +//! ```text +//! thread = one worker +//! warp = 32 threads that run together +//! CTA = one block of threads that can share memory +//! ``` +//! +//! [`copy_g2s`] lets every thread in a CTA copy one piece of a tile: +//! +//! ```text +//! global matrix ── all CTA threads ──► shared tile +//! ``` +//! +//! Types hold facts fixed before the kernel starts, such as layouts and +//! alignment. Function arguments hold values known only while it runs, such +//! as pointers, sizes, the tile coordinate, and the thread number. + +use core::marker::PhantomData; + +use crate::markers::LeadingDimMarker; + +/// Commit each participating thread's asynchronous copies and expose the tile. +/// +/// This single-stage protocol waits for all committed copy groups before the +/// typed CTA barrier. It does not implement a multistage producer pipeline. +/// +/// # Safety +/// Every thread participating in the preceding cooperative copies must call +/// this operation convergently, with exactly `THREADS` complete-warp +/// participants. Destination storage must remain live through the wait, and +/// `BARRIER` must satisfy `NamedBarrier::sync`'s phase and exclusivity rules. +#[inline(always)] +pub unsafe fn finish_async_copies() { + // SAFETY: every participating thread retains its copy destinations and + // executes the matching completion protocol under the caller contract. + unsafe { + ctm_rs::async_copy::cp_async_commit_group(); + ctm_rs::async_copy::cp_async_wait_group::<0>(); + crate::arch::NamedBarrier::::sync(); + } +} + +/// A row-major matrix stored in global GPU memory. +/// +/// `leading_dim.elements` is the row pitch: the number of elements from the +/// start of one row to the start of the next. It can be larger than `cols` +/// when rows contain padding. +/// +/// ```text +/// row 0: [ data data data | padding ] +/// row 1: [ data data data | padding ] +/// <---- row pitch ----> +/// ``` +/// +/// `LD` records a compile-time divisibility promise about that pitch. Build +/// this plain data value with a struct literal. [`copy_g2s`] checks the static +/// layout contract; its caller must satisfy the documented pointer, capacity, +/// alignment, and pitch requirements. +#[repr(C)] +pub struct GmemMatrix { + /// Start of the global-memory allocation. + pub base: *const T, + /// Logical matrix rows. + pub rows: usize, + /// Logical matrix columns, excluding row padding. + pub cols: usize, + /// Row pitch plus its compile-time divisibility marker. + pub leading_dim: LD, +} + +/// A tile in memory shared by all threads in one CTA. +/// +/// `Layout` has no runtime bytes. It tells the compiler how logical rows and +/// columns map to physical addresses: +/// +/// ```text +/// base pointer + capacity + zero-byte Layout marker +/// │ +/// └── visible at the copy call +/// ``` +/// +/// Build this plain data value with a struct literal. +#[repr(C)] +pub struct SmemTile { + /// Start of the shared-memory allocation. + pub base: *mut T, + /// Number of `T` elements available from `base`. + pub capacity: usize, + /// Zero-byte marker that keeps the shared layout in the type. + pub layout: PhantomData, +} + +/// A borrowed fixed-size logical row in a typed shared-memory tile. +/// +/// This view contains coordinates and a borrow, not loaded values. Reductions +/// can consume one element at a time without retaining a full register tile. +pub struct SharedRow<'a, T: Copy, Layout, const N: usize> { + tile: &'a SmemTile, + row: usize, + column: usize, +} + +impl SharedRow<'_, T, Layout, N> { + /// Reduce corresponding shared/global logical coordinates in increasing order. + /// + /// Each pair is loaded immediately before its callback. The operation does + /// not create intermediate register tiles or reassociate the accumulation. + /// The other tensor may use a different physical layout and must contain + /// at least `N` logical elements. Empty rows never load or call `f`. + #[inline(always)] + pub fn zip_fold( + self, + other: crate::tensor::partition::StridedTensor<'_, Element, OtherLayout, Storage>, + initial: Accumulator, + mut f: impl FnMut(Accumulator, T, Storage) -> Accumulator, + ) -> Accumulator + where + Element: crate::tensor::TensorElement, + OtherLayout: crate::tensor::partition::IndexLayout, + Storage: Copy, + { + assert!(N <= other.len()); + let mut accumulator = initial; + let mut i = 0; + while i < N { + ctm_rs::thread::__unroll_config::<0>(); + // SAFETY: the other extent was checked above; row_view's unsafe + // constructor established the shared row's full-lifetime contract. + let (global, shared) = unsafe { + ( + other.get_unchecked(i), + self.tile.get(self.row, self.column + i), + ) + }; + accumulator = f(accumulator, shared, global); + i += 1; + } + accumulator + } +} + +impl SmemTile { + /// Borrow `N` consecutive logical row coordinates without loading them. + /// + /// # Safety + /// For the entire returned view's lifetime, every `(row, column + i)`, + /// `i < N`, must be in the layout domain, map below `capacity`, and refer + /// to live, aligned, initialized readable storage. All asynchronous writes + /// must have completed and the necessary synchronization must have occurred. + /// No write may conflict with the borrowed row during that lifetime. Empty + /// rows contain no coordinates and do not require an accessible pointer. + #[inline(always)] + pub unsafe fn row_view( + &self, + row: usize, + column: usize, + ) -> SharedRow<'_, T, Layout, N> { + SharedRow { + tile: self, + row, + column, + } + } + + /// Read one logical shared-memory coordinate through its declared layout. + /// + /// This is the single-coordinate read `tensor[(row, column)]` of CuTe's + /// `Tensor.__getitem__`. + /// + /// # Safety + /// The coordinate must lie within the layout, its mapped offset must be + /// below `capacity`, and the element must be initialized and readable. + /// Pending asynchronous writes must be completed and appropriate thread + /// synchronization performed; no active write may conflict with this read. + #[inline(always)] + pub unsafe fn get(&self, row: usize, column: usize) -> T { + let offset = crate::static_layout::crd2idx_2d::(row, column); + unsafe { self.base.add(offset).read() } + } + + /// Gather N consecutive logical row coordinates into a register tile. + /// + /// Consecutive logical coordinates may be physically swizzled; the layout + /// determines each address rather than assuming contiguous physical bytes. + /// This is CuTe's `Tensor.load` over the `N` coordinates of one row. + /// + /// # Safety + /// Every `(row, column + i)` for `i < N` must satisfy `get`'s bounds, + /// initialization, completion and synchronization requirements. + #[inline(always)] + pub unsafe fn load( + &self, + row: usize, + column: usize, + ) -> crate::tensor::RegisterTile { + crate::tensor::RegisterTile::from_fn(|i| unsafe { self.get(row, column + i) }) + } +} + +/// Copy one global-memory tile into shared memory using the whole CTA. +/// +/// `tile_coord = (r, c)` counts tiles, not elements. For a 16x8 tile, +/// `(2, 3)` starts at element `(32, 24)`. +/// +/// The compiler retains this call as one typed cooperative-copy operation. +/// The selected backend chooses its realization. On a matrix edge, every +/// thread still follows the same path; a copy reports how many source bytes +/// are valid, and missing bytes become zero. The caller must still complete +/// the copy protocol and synchronize the CTA. +/// +/// # Safety +/// +/// The CTA must contain exactly the threads described by `ThrL`. Every one of +/// those threads must call this function on the same control-flow path, and +/// their `tidx` values must cover `ThrL` exactly once. Both raw pointers must +/// be valid. Each `Atom`, meaning one fixed-size copy piece, must be aligned. +/// The destination must have the capacity required by `SmemL`. +/// +/// `src.leading_dim.elements` must be at least `src.cols` and divisible by +/// the divisor carried by its `LeadingDim` type. The compiler checks that +/// the thread/value layout covers the tile exactly, that each atom maps to a +/// valid destination, and that `D * size_of::()` is at least one atom. +/// Edge tiles are safe because out-of-bounds source bytes are replaced by +/// zero instead of being read. +#[inline(never)] +#[allow(clippy::too_many_arguments)] +pub unsafe fn copy_g2s( + src: &GmemMatrix, + tile_coord: (usize, usize), + dst_smem_tile: &mut SmemTile, + tidx: u32, +) { + let _ = ( + src, + tile_coord, + dst_smem_tile, + tidx, + PhantomData::<(Atom, ThrL, ValL, TileL)>, + ); + unreachable!("cute-rs `copy_g2s` executed outside device compilation") +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Layout markers must add no bytes to global and shared copy carriers. + #[test] + fn runtime_carriers_are_only_the_visible_abi_fields() { + // The layout marker occupies no bytes: + // + // SmemTile = [ pointer ][ capacity ][ zero-byte static marker ] + assert_eq!( + core::mem::size_of::>>(), + 4 * core::mem::size_of::() + ); + assert_eq!( + core::mem::size_of::>(), + 2 * core::mem::size_of::() + ); + + let tile = SmemTile:: { + base: core::ptr::null_mut(), + capacity: 24, + layout: PhantomData, + }; + assert!(tile.base.is_null()); + assert_eq!(tile.capacity, 24); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/mod.rs new file mode 100644 index 0000000000..1e829dce8b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/mod.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Data movement between global memory, shared memory and thread registers. +//! +//! [`cooperative`] moves shared tiles with block participation; [`tiled`] +//! describes typed copy atoms and operand roles. [`vector`] moves one thread's +//! short vectors, while [`tma`] describes asynchronous whole-tile transfers. +//! Each family retains its own layout and participation requirements. + +pub mod cooperative; +pub mod tiled; +pub mod tma; +pub mod vector; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tiled.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tiled.rs new file mode 100644 index 0000000000..fe8a28c64a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tiled.rs @@ -0,0 +1,295 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed, block-wide copies from global memory to shared memory. +//! +//! Every thread in a block helps move part of one tile: +//! +//! ```text +//! global tile: [thread 0 part][thread 1 part][...] +//! │ +//! TiledCopy::get_slice(thread_idx) ──► ThrCopy::copy +//! │ +//! ▼ +//! shared tile: [thread 0 part][thread 1 part][...] +//! ``` +//! +//! [`GlobalCopyTensor`] describes the source. [`SharedTensor`] describes the +//! destination. [`TiledCopy`] describes how the work is split across threads, +//! and [`TiledCopy::get_slice`] selects one thread's part of that plan as a +//! [`ThrCopy`], following CuTeDSL's `TiledCopy.get_slice` and `ThrCopy`. +//! These wrappers add no runtime storage beyond the source and destination +//! views and the thread index. The compiler inlines them down to +//! [`crate::copy_g2s`]. + +use core::marker::PhantomData; + +use crate::copy::cooperative::{GmemMatrix, SmemTile, copy_g2s}; +use crate::markers::LeadingDimMarker; +use crate::tensor::TensorElement; + +/// A typed global-memory source for a block-wide copy. +/// +/// `Element` says what the bits mean. `T` says how those bits are physically +/// stored and copied. For example: +/// +/// ```text +/// Element = Mxf4E2M1 four logical FP4 values +/// T = f16 one 16-bit storage box +/// ``` +/// +/// `Role` can mark the source as A or B so typed operations cannot swap them. +pub struct GlobalCopyTensor, T: Copy, LD, Role = ()> { + /// Global pointer, matrix bounds and row-pitch promise used by cooperative lowering. + carrier: GmemMatrix, + /// Logical element encoding and A/B role attached to the global storage. + element: PhantomData<(Element, Role)>, +} + +impl, T: Copy, LD, Role> GlobalCopyTensor { + /// Attach an element meaning and role to a global-memory matrix view. + #[must_use] + #[inline(always)] + pub(crate) const fn from_carrier(carrier: GmemMatrix) -> Self { + Self { + carrier, + element: PhantomData, + } + } + + /// Return the underlying matrix view used by the compiler copy operation. + #[inline(always)] + pub(crate) const fn carrier(&self) -> &GmemMatrix { + &self.carrier + } +} + +/// A typed view of an existing shared-memory tile. +/// +/// `Element` and `T` have the same meaning as in [`GlobalCopyTensor`]. `Layout` +/// maps tile coordinates to shared-memory offsets. `Role` can distinguish A +/// storage from B storage. +/// +/// This view does not allocate shared memory. +pub struct SharedTensor, T: Copy, Layout, Role = ()> { + /// Shared allocation base, capacity and physical layout marker. + carrier: SmemTile, + /// Logical element encoding and A/B role attached to the shared storage. + element: PhantomData<(Element, Role)>, +} + +impl, T: Copy, Layout, Role> SharedTensor { + /// View an existing shared-memory allocation as a typed tile. + /// + /// # Safety + /// + /// - `base` must point to shared memory and be aligned for `T`. + /// - At least `capacity` elements must remain live while the view is used. + /// - Every offset produced by `Layout` must be below `capacity`. + /// - Threads must synchronize before reading data written by other + /// threads or by an asynchronous copy. + /// - Reads and writes must not break Rust's aliasing rules. + #[must_use] + #[inline(always)] + pub const unsafe fn from_raw_parts(base: *mut T, capacity: usize) -> Self { + unsafe { __compiler::shared_tensor_overlay(base, capacity) } + } + + /// Return the read-only shared-memory view used by compiler operations. + #[inline(always)] + pub(crate) const fn carrier(&self) -> &SmemTile { + &self.carrier + } + + /// Return the writable shared-memory view used by compiler operations. + #[inline(always)] + const fn carrier_mut(&mut self) -> &mut SmemTile { + &mut self.carrier + } +} + +/// Stable compiler boundary for attaching tensor meaning to shared storage. +/// +/// Kernel code uses [`SharedTensor::from_raw_parts`]. The importer recognizes +/// this exact helper when rustc inlines that small public wrapper. +#[doc(hidden)] +pub mod __compiler { + use super::*; + + /// Compiler boundary attaching element and operand-role meaning to a shared-memory carrier. + #[inline(never)] + pub const unsafe fn shared_tensor_overlay, T: Copy, Layout, Role>( + base: *mut T, + capacity: usize, + ) -> SharedTensor { + SharedTensor { + carrier: SmemTile { + base, + capacity, + layout: PhantomData, + }, + element: PhantomData, + } + } +} + +/// A complete plan for how one block copies one tile. +/// +/// The type parameters describe the plan: +/// +/// ```text +/// Atom smallest hardware copy +/// ThreadLayout which thread copies which part +/// ValueLayout which values that thread copies +/// TileLayout shape of the complete source tile +/// ``` +/// +/// The value contains no runtime data. [`Self::get_slice`] selects one +/// thread's part of the plan; the destination layout comes from the +/// [`SharedTensor`] passed to [`ThrCopy::copy`]. +pub struct TiledCopy { + /// Records the atom and thread/value/tile maps without runtime storage. + config: PhantomData<(Atom, ThreadLayout, ValueLayout, TileLayout)>, +} + +impl + TiledCopy +{ + /// Create the zero-sized copy plan. + #[must_use] + #[inline(always)] + pub const fn new() -> Self { + Self { + config: PhantomData, + } + } + + /// Select the part of the plan one thread carries out. + /// + /// This is CuTeDSL's `TiledCopy.get_slice(thr_idx)`. Device use must pass + /// the calling thread's own index in the block, as `ThreadLayout` numbers + /// the threads; [`ThrCopy::copy`] then issues that thread's part. + #[must_use] + #[inline(always)] + pub const fn get_slice( + &self, + thread_idx: u32, + ) -> ThrCopy { + ThrCopy { + thread_idx, + plan: PhantomData, + } + } +} + +impl Default + for TiledCopy +{ + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +/// One thread's slice of a [`TiledCopy`] plan. +/// +/// This is CuTeDSL's `ThrCopy`, returned by [`TiledCopy::get_slice`]. The +/// plan has no runtime data, so the slice holds only the thread index and +/// the plan's type parameters: +/// +/// ```text +/// TiledCopy +/// │ get_slice(thread_idx) +/// ▼ +/// ThrCopy { thread_idx } one per thread, copies its own part +/// ``` +pub struct ThrCopy { + /// Index of the thread this slice belongs to, as `ThreadLayout` numbers the threads. + thread_idx: u32, + /// Records the atom and thread/value/tile maps without runtime storage. + plan: PhantomData<(Atom, ThreadLayout, ValueLayout, TileLayout)>, +} + +impl + ThrCopy +{ + /// Copy this thread's part of one global tile into shared memory. + /// + /// `tile_coord` selects the source tile. The thread index given to + /// [`TiledCopy::get_slice`] selects this thread's assigned part of that + /// tile. In CuTeDSL this is `cute.copy(tiled_copy, + /// thr_copy.partition_S(src), thr_copy.partition_D(dst))`; the Rust form + /// covers one two-dimensional global tile and one shared destination + /// tile, with both partitions derived from the plan's layouts inside the + /// compiler's cooperative-copy operation, [`crate::copy_g2s`]. + /// + /// # Safety + /// + /// - Every thread named by `ThreadLayout` must call this operation. + /// - The launched thread count must match `ThreadLayout`. + /// - Source and destination pointers, sizes, layouts, row pitches, and + /// copy alignment must all match their typed descriptions. + /// - Threads must wait for the copy before reading the destination. + #[inline(always)] + pub unsafe fn copy< + Element: TensorElement, + T: Copy, + LD: LeadingDimMarker, + SharedLayout, + Role, + >( + &self, + source: &GlobalCopyTensor, + tile_coord: (usize, usize), + destination: &mut SharedTensor, + ) { + unsafe { + copy_g2s::( + source.carrier(), + tile_coord, + destination.carrier_mut(), + self.thread_idx, + ) + } + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Copy plans must be zero-sized and thread slices must hold only the thread index, + /// while typed source/destination views retain their carrier sizes. + #[test] + fn tiled_copy_is_static_and_typed_views_preserve_carrier_size() { + /// Marker-only copy plan used to check zero runtime size. + type Copy = TiledCopy<(), (), (), ()>; + assert_eq!(core::mem::size_of::(), 0); + /// One thread's slice of the marker-only plan. + type Slice = ThrCopy<(), (), (), ()>; + assert_eq!(core::mem::size_of::(), core::mem::size_of::()); + /// Distinct logical element marker used to test typed copy carrier size. + #[derive(Clone, Copy)] + struct Element; + impl TensorElement for Element {} + assert_eq!( + core::mem::size_of::>(), + core::mem::size_of::>() + ); + assert_eq!( + core::mem::size_of::>(), + core::mem::size_of::>() + ); + } + + /// The thread slice must carry the index it was selected with. + #[test] + fn get_slice_records_the_thread_index() { + let plan = TiledCopy::<(), (), (), ()>::new(); + assert_eq!(plan.get_slice(0).thread_idx, 0); + assert_eq!(plan.get_slice(127).thread_idx, 127); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/host.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/host.rs new file mode 100644 index 0000000000..9fc49ac151 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/host.rs @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Host-side descriptor encoding and element contracts. + +use super::{ + ReifySmem2D, TensorMap, TmaEncodeOptions, TmaL2Promotion, TmaSwizzleMode, tma_swizzle_mode_for, + tma_swizzle_span_bytes, +}; +use core::marker::PhantomData; +use core::mem::MaybeUninit; +use cuda_core::sys as driver; + +/// An element type TMA can encode in a descriptor. +/// +/// `DATA_TYPE` is the CUDA driver value. `BYTES` is one element's size. +pub trait TmaElement { + /// CUDA driver code for this element type. + const DATA_TYPE: driver::CUtensorMapDataType; + /// Bytes in one element. + const BYTES: u32; +} + +/// Raw bytes, including one byte containing two packed FP4 values. +impl TmaElement for u8 { + const DATA_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT8; + const BYTES: u32 = 1; +} + +impl TmaElement for f32 { + const DATA_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + const BYTES: u32 = 4; +} + +impl TmaElement for f16 { + const DATA_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + const BYTES: u32 = 2; +} + +/// Raw 16-bit values, including BF16 bit patterns. TMA only moves bytes. +impl TmaElement for u16 { + const DATA_TYPE: driver::CUtensorMapDataType = + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT16; + const BYTES: u32 = 2; +} + +/// Build a descriptor that loads or stores one `SmemL` tile. +/// +/// This is the CuTeDSL host name `create_tensor_map_tiled`; the element type, +/// box and swizzle come from `T` and `SmemL` instead of explicit arguments. +/// The global matrix is row-major. Types provide the tile size, swizzle, +/// and element format. Arguments provide the global pointer, matrix size, +/// and row pitch in elements. Unsupported layouts return an error; they +/// are never rounded to a nearby TMA mode. +pub fn create_tensor_map_tiled( + global_base: *mut core::ffi::c_void, + rows: u64, + cols: u64, + ld_elements: u64, +) -> Result, alloc::string::String> { + create_tensor_map_tiled_with_options( + global_base, + rows, + cols, + ld_elements, + TmaEncodeOptions::DEFAULT, + ) +} + +/// Build a row-major descriptor with an explicit L2 cache policy. +/// +/// Tile shape and validation match [`create_tensor_map_tiled`]. Use +/// [`TmaEncodeOptions::L2_128B`] to match the usual CUTLASS/CuTeDSL input +/// policy. The shorter function keeps the original no-promotion default. +// The driver stores this address in the descriptor. It does not read from +// `global_base` while encoding it. +#[allow(clippy::not_unsafe_ptr_arg_deref)] +pub fn create_tensor_map_tiled_with_options( + global_base: *mut core::ffi::c_void, + rows: u64, + cols: u64, + ld_elements: u64, + options: TmaEncodeOptions, +) -> Result, alloc::string::String> { + use alloc::format; + + if SmemL::SHAPE.0 <= 0 || SmemL::SHAPE.1 <= 0 { + return Err(format!( + "TMA tile extents must be positive, got {} x {}", + SmemL::SHAPE.0, + SmemL::SHAPE.1 + )); + } + if SmemL::STRIDE.0 != SmemL::SHAPE.1 || SmemL::STRIDE.1 != 1 { + return Err(format!( + "TMA v0 requires a row-major shared tile, got strides ({}, {})", + SmemL::STRIDE.0, + SmemL::STRIDE.1 + )); + } + if SmemL::OFFSET != 0 { + return Err(format!( + "TMA cannot express a composed offset (got {})", + SmemL::OFFSET + )); + } + let Some(mode) = tma_swizzle_mode_for::(T::BYTES) else { + return Err(format!( + "swizzle S<{},{},{}> (elements) is not one of TMA's encodable modes", + SmemL::NUM_BITS, + SmemL::NUM_BASE, + SmemL::NUM_SHIFT + )); + }; + if let Some(span_bytes) = tma_swizzle_span_bytes(mode) { + let pitch_bytes = SmemL::SHAPE + .1 + .checked_mul(i64::from(T::BYTES)) + .ok_or_else(|| alloc::string::String::from("TMA shared-tile pitch overflows i64"))?; + if pitch_bytes != span_bytes { + return Err(format!( + "swizzled TMA tile pitch must equal the {span_bytes}-byte swizzle span, got {pitch_bytes} bytes" + )); + } + } + let swizzle = match mode { + TmaSwizzleMode::None => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_NONE, + TmaSwizzleMode::B32 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_32B, + TmaSwizzleMode::B64 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_64B, + TmaSwizzleMode::B128 => driver::CUtensorMapSwizzle_enum_CU_TENSOR_MAP_SWIZZLE_128B, + }; + if ld_elements < cols { + return Err(format!( + "pitch {ld_elements} is smaller than the column count {cols}" + )); + } + + // TMA lists the innermost dimension first. Dimension 0 is columns, + // which are contiguous in a row-major matrix. + let global_dim: [u64; 2] = [cols, rows]; + let global_strides: [u64; 1] = [ld_elements * u64::from(T::BYTES)]; + let box_dim: [u32; 2] = [SmemL::SHAPE.1 as u32, SmemL::SHAPE.0 as u32]; + let element_strides: [u32; 2] = [1, 1]; + let l2_promotion = match options.l2_promotion { + TmaL2Promotion::None => driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_NONE, + TmaL2Promotion::B64 => { + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_64B + } + TmaL2Promotion::B128 => { + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_128B + } + TmaL2Promotion::B256 => { + driver::CUtensorMapL2promotion_enum_CU_TENSOR_MAP_L2_PROMOTION_L2_256B + } + }; + + let mut tensor_map = MaybeUninit::::uninit(); + let result = unsafe { + driver::cuTensorMapEncodeTiled( + tensor_map.as_mut_ptr(), + T::DATA_TYPE, + 2, + global_base, + global_dim.as_ptr(), + global_strides.as_ptr(), + box_dim.as_ptr(), + element_strides.as_ptr(), + driver::CUtensorMapInterleave_enum_CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + l2_promotion, + driver::CUtensorMapFloatOOBfill_enum_CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + }; + if result != driver::cudaError_enum_CUDA_SUCCESS { + return Err(format!("cuTensorMapEncodeTiled failed: {result:?}")); + } + let map = unsafe { tensor_map.assume_init() }; + /// Check that driver descriptor storage remains exactly 128 bytes before transmuting. + const _: () = assert!(core::mem::size_of::() == 128); + Ok(TensorMap { + bytes: unsafe { core::mem::transmute::(map) }, + marker: PhantomData, + }) +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/mod.rs new file mode 100644 index 0000000000..cb3a38c33f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/mod.rs @@ -0,0 +1,307 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed tile copies with the Tensor Memory Accelerator (TMA). +//! +//! TMA is hardware that moves a whole tile between global memory and shared +//! memory. A thread is one GPU worker. A block, also called a CTA, is a group +//! of threads that can use the same shared memory. +//! +//! Before launch, the host builds a 128-byte *descriptor*. It contains the +//! global address, sizes, element type, and tile layout. The kernel gives that +//! descriptor to TMA. Here, Rust also puts the element and shared layout in +//! the descriptor type: +//! +//! ```text +//! host: TensorMap -> kernel: &TensorMap +//! descriptor value #[grid_constant] parameter +//! same element and layout types +//! ``` +//! +//! If host and kernel layouts differ, the program does not compile. A +//! *swizzle* is a fixed address rearrangement that reduces shared-memory bank +//! conflicts: several threads requesting the same memory bank and having to +//! wait. Both sides read sizes, strides, and the swizzle from [`ReifySmem2D`]. +//! +//! MXFP4 is a block-scaled 4-bit floating-point format. It uses ordinary byte +//! and 16-bit transport views: +//! +//! ```text +//! A or B: 128x128 FP4 values -> 128x64 bytes -> 8192 bytes +//! two FP4 values share each byte; use a 64-byte (B64) swizzle +//! +//! scales: 128x4 UE8M0 bytes -> 1x256 u16 values -> 512 bytes +//! UE8M0 is an 8-bit power-of-two scale +//! u16 is only the unit TMA uses to move the same bytes +//! ``` +//! +//! The scale layout expects the standard SM120 packing. After the copy, the +//! kernel may read the same 512 bytes as 128 packed `u32` scale words. + +use core::marker::PhantomData; + +use crate::markers::ReifySmem2D; + +/// One 128-byte TMA descriptor typed by element and shared layout. +/// +/// Its runtime bytes have the same layout as the CUDA driver's `CUtensorMap`. +/// A `#[grid_constant] desc: &TensorMap` kernel parameter takes this +/// value at launch. Every thread borrows the same read-only descriptor bytes. +/// The tensor data stays in device memory and must remain live until the GPU +/// work completes. +#[repr(C, align(64))] +pub struct TensorMap { + /// Descriptor bytes produced by the CUDA driver. + pub bytes: [u8; 128], + /// Zero-byte marker that keeps `T` and `SmemL` in the type. + pub marker: PhantomData<(T, SmemL)>, +} + +impl TensorMap { + /// Prefetch this descriptor into the device's tensor-map cache. + /// + /// This is CuTeDSL's `cpasync.prefetch_descriptor`. Issuer selection + /// remains with the caller, so a kernel can prefetch its descriptors once + /// before starting TMA transfers. + /// + /// # Safety + /// + /// Call only on a device supporting TMA. `self` must contain a valid encoded + /// CUDA tensor map in device-accessible storage, such as a grid-constant + /// kernel argument, and remain live through the prefetch. + #[inline(always)] + pub unsafe fn prefetch_descriptor(&self) { + unsafe { + ctm_rs::tma::prefetch_tensormap( + core::ptr::from_ref(self).cast::(), + ) + }; + } +} + +// The type parameters describe the encoded bytes; they are not stored values. +impl Copy for TensorMap {} + +impl Clone for TensorMap { + fn clone(&self) -> Self { + *self + } +} + +/// Shared-memory swizzles supported by TMA. +/// +/// TMA supports only these byte-based patterns. Each keeps 16-byte chunks +/// together while rearranging the surrounding addresses: +/// +/// ```text +/// mode byte-based pattern +/// None no rearrangement +/// B32 S<1,4,3> +/// B64 S<2,4,3> +/// B128 S<3,4,3> +/// ``` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TmaSwizzleMode { + /// Keep shared-memory addresses in their original order. + None, + /// Rearrange addresses within a 32-byte span. + B32, + /// Rearrange addresses within a 64-byte span. + B64, + /// Rearrange addresses within a 128-byte span. + B128, +} + +/// Requested level-2 (L2) cache fetch size for a TMA descriptor. +/// +/// [`Self::None`] is the default used by `create_tensor_map_tiled`. Select +/// [`Self::B128`] to request a 128-byte fetch for GEMM input tiles. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TmaL2Promotion { + /// Do not request a larger L2 fetch. + None, + /// Fetch 64-byte L2 sectors. + B64, + /// Fetch 128-byte L2 sectors. + B128, + /// Fetch 256-byte L2 sectors. + B256, +} + +/// Host options for `create_tensor_map_tiled_with_options`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TmaEncodeOptions { + /// L2 fetch size requested for this descriptor. + pub l2_promotion: TmaL2Promotion, +} + +impl TmaEncodeOptions { + /// Do not request a larger L2 fetch. + pub const DEFAULT: Self = Self { + l2_promotion: TmaL2Promotion::None, + }; + + /// Request 128-byte L2 sectors. + pub const L2_128B: Self = Self { + l2_promotion: TmaL2Promotion::B128, + }; +} + +impl Default for TmaEncodeOptions { + fn default() -> Self { + Self::DEFAULT + } +} + +/// Convert a byte-based `S` swizzle to a TMA mode. +/// +/// Returns `None` when hardware cannot represent the exact pattern. It never +/// rounds to a nearby mode. +pub const fn tma_swizzle_mode(bits: u32, m_bytes: u32, shift: i32) -> Option { + if bits == 0 { + return Some(TmaSwizzleMode::None); + } + match (bits, m_bytes, shift) { + (1, 4, 3) => Some(TmaSwizzleMode::B32), + (2, 4, 3) => Some(TmaSwizzleMode::B64), + (3, 4, 3) => Some(TmaSwizzleMode::B128), + _ => None, + } +} + +/// Convert `SmemL`'s element-based swizzle to a byte-based TMA mode. +/// +/// `element_bytes` shifts the address fields from elements to bytes. Returns +/// `None` if the size is zero, is not a power of two, or TMA cannot represent +/// the result exactly. The function never rounds the size. +pub const fn tma_swizzle_mode_for( + element_bytes: u32, +) -> Option { + if !element_bytes.is_power_of_two() { + return None; + } + let log2 = element_bytes.trailing_zeros(); + tma_swizzle_mode(SmemL::NUM_BITS, SmemL::NUM_BASE + log2, SmemL::NUM_SHIFT) +} + +/// Bytes covered by one repeat of the swizzle, or `None` without a swizzle. +#[cfg(any(feature = "host", test))] +const fn tma_swizzle_span_bytes(mode: TmaSwizzleMode) -> Option { + match mode { + TmaSwizzleMode::None => None, + TmaSwizzleMode::B32 => Some(32), + TmaSwizzleMode::B64 => Some(64), + TmaSwizzleMode::B128 => Some(128), + } +} + +/// Bytes covered by one `SmemL` tile of `T`: the layout's cosize times the +/// element size, as CuTeDSL's `size_in_bytes(dtype, layout)` computes it. +/// +/// The cosize is the offset of the tile's last element plus one, taken from +/// the affine part of the layout (its shape and strides). A swizzle only +/// permutes addresses inside that span and the composed offset is not +/// counted, so a padded row stride enlarges the result while a swizzle does +/// not. For the dense row-major tiles TMA encodes, this equals rows times +/// columns times the element size. +/// +/// Pass this number to the load pipeline's expected-transaction operation for +/// [`copy_tma_2d`]. In a kernel, compute it as a constant so it creates no +/// runtime function call: +/// +/// ```text +/// const N: u32 = size_in_bytes::(); +/// ``` +#[inline(always)] +pub const fn size_in_bytes() -> u32 { + let (rows, cols) = SmemL::SHAPE; + if rows <= 0 || cols <= 0 { + return 0; + } + let (row_stride, col_stride) = SmemL::STRIDE; + let cosize = (rows - 1) * row_stride + (cols - 1) * col_stride + 1; + cosize as u32 * size_of::() as u32 +} + +/// Start one TMA copy from global memory into a shared tile. +/// +/// The compiler recognizes this call. `desc` and the destination must carry +/// the same `SmemL`, which Rust checks at the call. The compiler also checks +/// that TMA supports the layout. +/// +/// `coord = (r, c)` counts tiles, not elements. The element origin is: +/// +/// ```text +/// row = r * SmemL::SHAPE.0 +/// col = c * SmemL::SHAPE.1 +/// ``` +/// +/// The caller must initialize the `mbarrier`, attach [`size_in_bytes`] as its +/// expected byte count, and wait for it. +/// +/// # Safety +/// +/// Exactly one thread in the block must call this function. `barrier` and +/// `dst_smem_tile` must be valid objects in this block's shared memory. Use +/// the lower-level raw TMA APIs for remote or multicast cluster copies. +/// +/// The destination address must meet the alignment required by the layout's +/// swizzle phase. `desc` must have been built by `create_tensor_map_tiled` for the +/// same `SmemL`. The packed-MXFP4 B64 layout described above requires a +/// 512-byte-aligned buffer base; the CuTeDSL comparison kernel uses 1024. +#[inline(never)] +pub unsafe fn copy_tma_2d( + desc: *const TensorMap, + coord: (usize, usize), + dst_smem_tile: &mut crate::copy::cooperative::SmemTile, + barrier: *mut cuda_device::barrier::Barrier, +) { + let _ = (desc, coord, dst_smem_tile, barrier); + unreachable!("cute-rs `copy_tma_2d` executed outside device compilation") +} + +/// Start one TMA copy from a shared tile into global memory. +/// +/// This is the reverse of [`copy_tma_2d`]. The descriptor and source carry +/// the same element and layout types, so Rust rejects mismatches. `coord` +/// counts tiles; `(r, c)` starts at element +/// `(r * SmemL::SHAPE.0, c * SmemL::SHAPE.1)`. +/// +/// TMA stores are tracked in copy groups rather than `mbarrier` objects. TMA +/// observes shared memory through an asynchronous *proxy*, a hardware view of +/// memory separate from ordinary thread accesses. The caller must publish +/// earlier ordinary shared writes to that proxy, issue the stores, commit the +/// group, and wait before reusing the source. [`crate::PipelineTmaStore`] +/// provides the commit and wait steps. +/// +/// The small `SmemTile` view is passed by value. This consumes only its base +/// pointer and capacity; it does not move or copy the shared-memory contents. +/// +/// # Safety +/// +/// Exactly one thread in the block must call this function. +/// `src_smem_tile` must be a live, correctly aligned tile in this block's +/// shared memory and use `SmemL`. `desc` must have been built for the same `T` +/// and `SmemL`. Both must remain valid until the committed copy group has +/// finished reading shared memory. Before this call, a proxy fence must make +/// all earlier ordinary shared-memory writes visible to TMA's async proxy. +#[inline(never)] +pub unsafe fn copy_tma_s2g_2d( + desc: *const TensorMap, + coord: (usize, usize), + src_smem_tile: crate::copy::cooperative::SmemTile, +) { + let _ = (desc, coord, src_smem_tile); + unreachable!("cute-rs `copy_tma_s2g_2d` executed outside device compilation") +} + +#[cfg(feature = "host")] +mod host; +#[cfg(feature = "host")] +pub use host::{TmaElement, create_tensor_map_tiled, create_tensor_map_tiled_with_options}; + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/tests.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/tests.rs new file mode 100644 index 0000000000..a26142999a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/tma/tests.rs @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks TMA descriptor ABI, swizzle modes, element units and host encoding options."] +/// TMA descriptors must retain their opaque ABI and Copy behavior even with non-Copy markers. +#[test] +fn descriptor_value_abi_does_not_require_copy_marker_types() { + // Layout and element markers need not be Copy: only descriptor bytes move. + /// Non-Copy marker proving descriptor Copy does not require marker trait bounds. + struct Marker; + /// Require Copy at compile time without constructing or reading a descriptor. + fn assert_copy() {} + assert_copy::>(); + assert_eq!( + core::mem::size_of::>(), + 128 + ); + assert_eq!( + core::mem::align_of::>(), + 64 + ); +} + +use super::*; + +/// Swizzle classification must accept exact hardware modes and reject nearby incompatible patterns. +#[test] +fn only_the_hardware_modes_are_encodable() { + assert_eq!(tma_swizzle_mode(0, 9, 9), Some(TmaSwizzleMode::None)); + assert_eq!(tma_swizzle_mode(1, 4, 3), Some(TmaSwizzleMode::B32)); + assert_eq!(tma_swizzle_mode(2, 4, 3), Some(TmaSwizzleMode::B64)); + assert_eq!(tma_swizzle_mode(3, 4, 3), Some(TmaSwizzleMode::B128)); + // Similar but unsupported patterns return None instead of rounding. + assert_eq!(tma_swizzle_mode(3, 5, 3), None); + assert_eq!(tma_swizzle_mode(3, 4, 4), None); + assert_eq!(tma_swizzle_mode(4, 4, 3), None); +} + +/// Element-unit swizzles must convert to the exact byte-unit TMA encoding. +#[test] +fn element_recast_reaches_the_byte_modes() { + use crate::markers::{ComposedLayout, ReifySmem2D, RowMajor, Swizzle}; + // For f16, element pattern S<3,3,3> becomes byte pattern S<3,4,3>, + // which is TMA B128. + /// FP16 element-unit layout that must recast to TMA’s B128 byte pattern. + type Smem = ComposedLayout, 0, RowMajor<8, 64>>; + assert_eq!(Smem::NUM_BITS, 3); + assert_eq!(tma_swizzle_mode_for::(2), Some(TmaSwizzleMode::B128)); + // For f32, it becomes S<3,5,3>, which TMA cannot encode. + assert_eq!(tma_swizzle_mode_for::(4), None); + // Invalid element sizes fail instead of being rounded. + assert_eq!(tma_swizzle_mode_for::(0), None); + assert_eq!(tma_swizzle_mode_for::(3), None); +} + +/// Named TMA modes must report their exact supported row-byte spans. +#[test] +fn swizzle_modes_name_their_exact_row_spans() { + assert_eq!(tma_swizzle_span_bytes(TmaSwizzleMode::None), None); + assert_eq!(tma_swizzle_span_bytes(TmaSwizzleMode::B32), Some(32)); + assert_eq!(tma_swizzle_span_bytes(TmaSwizzleMode::B64), Some(64)); + assert_eq!(tma_swizzle_span_bytes(TmaSwizzleMode::B128), Some(128)); +} + +/// The packed MXFP4 value tile must encode an 8192-byte footprint using B64 swizzling. +#[test] +fn flattened_mxfp4_ab_tile_is_8192_bytes_with_b64_swizzle() { + use crate::markers::{ComposedLayout, RowMajor, Swizzle}; + + // Two FP4 values share each byte: 128 K values become 64 bytes. + // Byte pattern S<2,4,3> is TMA B64. + /// Packed-byte value tile using the exact B64 TMA swizzle. + type ByteSmem = ComposedLayout, 0, RowMajor<128, 64>>; + assert_eq!(size_in_bytes::(), 8_192); + assert_eq!( + tma_swizzle_mode_for::(1), + Some(TmaSwizzleMode::B64) + ); +} + +/// Canonical scale transport must encode 512 bytes with no shared swizzle. +#[test] +fn canonical_mxfp4_scale_tile_is_512_bytes_without_swizzle() { + use crate::markers::RowMajor; + + // Standard SM120 packing: 128 rows x 4 UE8M0 bytes = 512 bytes. + // Viewing those bytes as u16 creates one legal 1x256 TMA tile. + /// One unswizzled row of 256 u16 scale carriers. + type ScaleSmem = RowMajor<1, 256>; + assert_eq!(size_in_bytes::(), 512); + assert_eq!( + tma_swizzle_mode_for::(2), + Some(TmaSwizzleMode::None) + ); +} + +/// `size_in_bytes` follows CuTeDSL: the cosize of the affine layout times the element size. +#[test] +fn size_in_bytes_is_cosize_times_element_size() { + use crate::markers::{C, ComposedLayout, Layout, RowMajor, Swizzle}; + + // A padded row stride enlarges the cosize beyond rows x columns. + /// Eight 64-element rows with a 72-element pitch. + type Padded = Layout<(C<8>, C<64>), (C<72>, C<1>)>; + assert_eq!(size_in_bytes::(), (7 * 72 + 63 + 1) * 2); + // A swizzle permutes addresses inside the span and adds no bytes. + /// B128-swizzled 8x64 FP16 tile. + type Swizzled = ComposedLayout, 0, RowMajor<8, 64>>; + assert_eq!(size_in_bytes::(), 8 * 64 * 2); +} + +/// Default descriptor options must retain no promotion while exposing the CuTeDSL L2 policy. +#[test] +fn descriptor_options_preserve_default_and_offer_cutedsl_policy() { + assert_eq!( + TmaEncodeOptions::default().l2_promotion, + TmaL2Promotion::None + ); + assert_eq!(TmaEncodeOptions::L2_128B.l2_promotion, TmaL2Promotion::B128); +} + +/// The packed-byte host carrier must use the driver’s UINT8 tensor-map element code. +#[cfg(feature = "host")] +#[test] +fn byte_carrier_has_uint8_tma_format() { + use cuda_core::sys as driver; + + assert_eq!(::BYTES, 1); + assert_eq!( + ::DATA_TYPE, + driver::CUtensorMapDataType_enum_CU_TENSOR_MAP_DATA_TYPE_UINT8 + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/copy/vector.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/vector.rs new file mode 100644 index 0000000000..0f2663d306 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/copy/vector.rs @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! One-thread vector loads and stores. +//! +//! These functions move one short, adjacent tile in a single transaction: +//! +//! ```text +//! load_tile::(src, 8) +//! │ +//! ▼ +//! src: [... | 8 | 9 | 10 | 11 | ...] +//! └──── one 16-byte load ────┘ +//! ``` +//! +//! During device compilation, the compiler replaces each function call with +//! a CuTe vector-copy operation. The Rust bodies are never run on the GPU and +//! intentionally panic if called on the host. +//! +//! # Size and alignment +//! +//! The first version supports a total transfer size that is a power of two +//! from 4 through 16 bytes. The caller must guarantee: +//! +//! ```text +//! bytes = N * size_of::() +//! idx % N == 0 +//! buffer address % bytes == 0 +//! idx + N <= buffer length +//! ``` +//! +//! GPU allocations used for kernel slices are 256-byte aligned. A tile index +//! built as `tile_number * N` therefore meets the remaining alignment rule. +//! A short final tile must use checked scalar loads or stores instead. + +/// Load `src[idx..idx + N]` as one vector transaction. +/// +/// # Safety +/// +/// - The full range must be in bounds. +/// - `idx` must be divisible by `N`. +/// - The first byte must be aligned to `N * size_of::()` bytes. +/// - The total byte size must be a supported power of two from 4 through 16. +#[inline(never)] +pub unsafe fn load_tile(src: &[T], idx: usize) -> [T; N] { + let _ = (src, idx); + unreachable!("cute-rs `load_tile` executed outside device compilation") +} + +/// Store `vals` at `dst[idx..idx + N]` as one vector transaction. +/// +/// `dst` is a raw device pointer, such as one returned by +/// `DisjointSlice::as_mut_ptr`. +/// +/// # Safety +/// +/// - The full destination range must be live, writable, and in bounds. +/// - No other active access may conflict with that range. +/// - `idx` must be divisible by `N`. +/// - The first byte must be aligned to `N * size_of::()` bytes. +/// - The total byte size must be a supported power of two from 4 through 16. +#[inline(never)] +pub unsafe fn store_tile(dst: *mut T, idx: usize, vals: &[T; N]) { + let _ = (dst, idx, vals); + unreachable!("cute-rs `store_tile` executed outside device compilation") +} + +/// Return `x` while promising the compiler that it is divisible by `D`. +/// +/// This is CuTeDSL's `cute.assume(x, divby=D)` with the divisor as a const +/// generic. +/// +/// ```text +/// x = 3 * 4 = 12 +/// assume::<4>(x) ──► 12, with the promise 12 % 4 == 0 +/// ``` +/// +/// Pass the returned value directly to the operation that needs the promise. +/// The compiler does not search through stored structs to find it. Struct +/// fields such as matrix row pitch use `LeadingDim` to keep the same fact +/// in their type. +/// +/// Prefer building the value so the promise is obvious; `thread * D` is +/// always divisible by `D`. +/// +/// # Safety +/// +/// `D` must be greater than zero and `x % D` must equal zero. A false promise +/// may make a later memory instruction misaligned, which is undefined +/// behavior. +#[inline(never)] +pub unsafe fn assume(x: usize) -> usize { + let _ = x; + unreachable!("cute-rs `assume` executed outside device compilation") +} + +// Short final tiles are handled by ordinary checked Rust code: +// +// if base + N <= len { /* vectorized tile calls */ } +// else { /* scalar loop over the tail */ } +// +// Only boundary threads use the scalar path. A block-wide `copy_g2s` follows +// a different rule because every thread must reach nearby barriers: +// +// every thread calls copy_g2s -> its assigned part copies 16 / fewer / 0 bytes + +/// Restricts cache-aware stores to supported scalar carrier types. +mod global_store_element { + /// Implemented only for f16, f32 and packed u32 store elements. + pub trait Sealed {} + impl Sealed for f16 {} + impl Sealed for f32 {} + impl Sealed for u32 {} +} + +/// Scalar storage supported by a cache-aware register-to-global copy. +/// +/// Implemented for `f16`, `f32`, and `u32` (including packed FP16 pairs). +pub trait GlobalStoreElement: Copy + global_store_element::Sealed {} +impl GlobalStoreElement for f16 {} +impl GlobalStoreElement for f32 {} +impl GlobalStoreElement for u32 {} + +/// A naturally owned, 32-byte-aligned carrier for a global register-tile store. +/// +/// The alignment is part of the Rust type, so passing a reference to this +/// carrier never strengthens the source pointer's alignment by assumption. +/// Construct it directly as `GlobalStoreTile([a, b, c, d, e, f, g, h])`. +#[repr(C, align(32))] +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct GlobalStoreTile( + #[doc = "Ordered scalar carriers to write in one aligned global transaction."] pub [T; N], +); + +/// Store one adjacent register tile without allocating its data in L1. +/// +/// This is an ordinary weak, non-atomic store with a cache eviction hint; it +/// supplies no synchronization. The copy width is `N * size_of::()` bytes. +/// Supported widths are 4, 8, 16, and 32 bytes. The 32-byte form requires SM100 +/// or newer; the selected backend checks target support. Existing memory +/// contents are replaced without conversion, including packed `u32` values. +/// +/// # Safety +/// +/// `dst` must point to a live, writable global-memory range of `N` elements, +/// aligned to the full copy width. No conflicting access may race this store. +/// The aligned source carrier must remain live during the copy. The transfer +/// width must be supported on the selected target. This function is a device compiler +/// boundary and panics if executed on the host. +#[inline(never)] +pub unsafe fn store_global_no_allocate( + dst: *mut T, + values: &GlobalStoreTile, +) { + const { + assert!(matches!( + N.checked_mul(core::mem::size_of::()), + Some(4 | 8 | 16 | 32) + )); + } + let _ = (dst, values); + unreachable!("cute-rs `store_global_no_allocate` executed outside device compilation") +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/decoder.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/decoder.rs new file mode 100644 index 0000000000..300851a6df --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/decoder.rs @@ -0,0 +1,573 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Checked BF16 decoder arguments, embedding indices and normalization roles. +//! +//! Hidden width and attention width are independent: a decoder with hidden width +//! 1024 can produce 16 query heads of width 128. Projection roles, normalization +//! roles, BF16 storage and fixed widths stay in Rust types across composition. +//! Matrices retain checkpoint orientation `[output, input]`. +//! +//! These descriptors borrow existing allocations and do not launch work. Shapes +//! cannot prove checkpoint contents, device context or stream readiness. Operators +//! must check those properties and retain input borrows until launches complete. + +use core::{fmt, marker::PhantomData}; +use cuda_core::DeviceBuffer; + +use crate::attention::metadata::{AttentionError, AttentionShape}; +use crate::ffn::{ + AttentionOutput, FfnError, FfnShape, FfnWeights, Key, ProjectionRole, ProjectionWeights, Query, + RowMajorShape, Value, +}; +use crate::numeric::BFloat16; + +/// A rejected decoder shape, vocabulary index or projection width. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DecoderError { + /// A BF16 matrix or allocation extent is invalid. + Matrix(FfnError), + /// Attention heads or their element extents are invalid. + Attention(AttentionError), + /// Typed projection widths do not equal their head counts times head width. + ProjectionWidthMismatch, + /// An embedding vocabulary must contain at least one row. + EmptyVocabulary, + /// A token index is outside the embedding vocabulary. + InvalidToken { + /// Position of the invalid token in the supplied host sequence. + index: usize, + /// Supplied token index. + token: u32, + /// Exclusive vocabulary bound. + vocabulary: usize, + }, +} + +impl fmt::Display for DecoderError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Matrix(error) => error.fmt(f), + Self::Attention(error) => error.fmt(f), + Self::ProjectionWidthMismatch => { + f.write_str("projection widths do not match attention head geometry") + } + Self::EmptyVocabulary => f.write_str("embedding vocabulary must be nonempty"), + Self::InvalidToken { + index, + token, + vocabulary, + } => write!( + f, + "token {token} at index {index} exceeds vocabulary bound {vocabulary}" + ), + } + } +} + +impl core::error::Error for DecoderError {} + +impl From for DecoderError { + fn from(error: FfnError) -> Self { + Self::Matrix(error) + } +} + +impl From for DecoderError { + fn from(error: AttentionError) -> Self { + Self::Attention(error) + } +} + +/// Checked activation geometry for a decoder layer with fixed model dimensions. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DecoderShape< + const H: usize, + const I: usize, + const QH: usize, + const KVH: usize, + const D: usize, +> { + /// Hidden and intermediate activation shapes. + ffn: FfnShape, + /// Independent query and key/value head geometry. + attention: AttentionShape, + /// Dense token-major query projection. + query: RowMajorShape, + /// Dense token-major key or value projection. + key_value: RowMajorShape, + /// Query heads flattened into independent normalization rows. + query_norm: RowMajorShape, + /// Key heads flattened into independent normalization rows. + key_norm: RowMajorShape, +} + +impl + DecoderShape +{ + /// Check all activation extents before allocation; input rows may be padded. + pub fn new(tokens: usize, input_stride: usize) -> Result { + let attention = AttentionShape::new()?; + let ffn = FfnShape::new(tokens, input_stride)?; + let query = RowMajorShape::dense(tokens, attention.query_width())?; + let key_value = RowMajorShape::dense(tokens, attention.kv_width())?; + let query_rows = tokens.checked_mul(QH).ok_or(FfnError::ExtentOverflow)?; + let key_rows = tokens.checked_mul(KVH).ok_or(FfnError::ExtentOverflow)?; + Ok(Self { + ffn, + attention, + query, + key_value, + query_norm: RowMajorShape::dense(query_rows, D)?, + key_norm: RowMajorShape::dense(key_rows, D)?, + }) + } + + /// Number of input tokens in the invocation. + #[must_use] + pub const fn tokens(self) -> usize { + self.ffn.rows() + } + /// Checked input matrix, including row padding. + #[must_use] + pub const fn input(self) -> RowMajorShape { + self.ffn.input() + } + /// Dense hidden-width activation shape. + #[must_use] + pub const fn hidden(self) -> RowMajorShape { + self.ffn.hidden() + } + /// Dense intermediate-width FFN shape. + #[must_use] + pub const fn intermediate(self) -> RowMajorShape { + self.ffn.intermediate() + } + /// Dense `[token, query_head * head_dim]` projection shape. + #[must_use] + pub const fn query(self) -> RowMajorShape { + self.query + } + /// Dense `[token, kv_head * head_dim]` key or value shape. + #[must_use] + pub const fn key_value(self) -> RowMajorShape { + self.key_value + } + /// Dense `[token * query_head, head_dim]` per-head normalization view. + #[must_use] + pub const fn query_norm(self) -> RowMajorShape { + self.query_norm + } + /// Dense `[token * kv_head, head_dim]` per-head normalization view. + #[must_use] + pub const fn key_norm(self) -> RowMajorShape { + self.key_norm + } + /// Query head count. + #[must_use] + pub const fn query_heads(self) -> usize { + QH + } + /// Key/value head count. + #[must_use] + pub const fn kv_heads(self) -> usize { + KVH + } + /// Coordinates per attention head. + #[must_use] + pub const fn head_dim(self) -> usize { + D + } + /// FFN geometry for the same token rows and input pitch. + #[must_use] + pub const fn ffn(self) -> FfnShape { + self.ffn + } + + /// Check the explicit const widths used by typed projection weights. + /// + /// Stable const generics cannot express products such as `QH * D` in field + /// types, so this relation is checked once at the composition boundary. + pub fn validate_projection_widths( + self, + ) -> Result<(), DecoderError> { + if QWIDTH != self.attention.query_width() || KVWIDTH != self.attention.kv_width() { + Err(DecoderError::ProjectionWidthMismatch) + } else { + Ok(()) + } + } + + /// Validate head geometry against role- and width-checked decoder weights. + pub fn validate_weights( + self, + _weights: &DecoderWeights<'_, H, I, QWIDTH, KVWIDTH, D>, + ) -> Result<(), DecoderError> { + self.validate_projection_widths::() + } +} + +/// Private supertrait restricting normalization roles. +mod sealed { + pub trait Sealed {} +} + +/// A sealed semantic role for one RMS normalization weight vector. +pub trait NormRole: sealed::Sealed {} + +/// Hidden-state normalization before attention projections. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct InputNorm; +/// Hidden-state normalization before the FFN. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct PostAttentionNorm; +/// Per-query-head normalization before RoPE. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QueryNorm; +/// Per-key-head normalization before RoPE. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct KeyNorm; +/// Hidden-state normalization after the decoder stack. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FinalNorm; + +impl sealed::Sealed for InputNorm {} +impl sealed::Sealed for PostAttentionNorm {} +impl sealed::Sealed for QueryNorm {} +impl sealed::Sealed for KeyNorm {} +impl sealed::Sealed for FinalNorm {} +impl NormRole for InputNorm {} +impl NormRole for PostAttentionNorm {} +impl NormRole for QueryNorm {} +impl NormRole for KeyNorm {} +impl NormRole for FinalNorm {} + +/// Borrowed BF16 RMS normalization weights, distinguished by role and width. +/// +/// A query norm cannot be substituted for the equally sized key norm: +/// ```compile_fail,E0308 +/// use cute_rs::decoder::{KeyNorm, NormWeights, QueryNorm}; +/// fn wrong_role(query: NormWeights<'_, QueryNorm, 128>) { +/// let _: NormWeights<'_, KeyNorm, 128> = query; +/// } +/// ``` +/// The immutable borrow lasts through the checked view's final use: +/// ```compile_fail,E0502 +/// use cuda_core::DeviceBuffer; +/// use cute_rs::{BFloat16, decoder::{InputNorm, NormWeights}}; +/// fn mutate(_: &mut DeviceBuffer) {} +/// fn alias(buffer: &mut DeviceBuffer) { +/// let weights = NormWeights::::new(buffer).unwrap(); +/// mutate(buffer); +/// let _ = weights.buffer(); +/// } +/// ``` +pub struct NormWeights<'a, Role: NormRole, const WIDTH: usize> { + buffer: &'a DeviceBuffer, + role: PhantomData, +} + +impl<'a, Role: NormRole, const WIDTH: usize> NormWeights<'a, Role, WIDTH> { + /// Check a positive logical width and BF16 allocation extent, allowing guards. + pub fn new(buffer: &'a DeviceBuffer) -> Result { + RowMajorShape::dense(1, WIDTH)?.validate_len(buffer.len())?; + Ok(Self { + buffer, + role: PhantomData, + }) + } + /// Borrow the existing allocation for a checked kernel launch. + #[must_use] + pub const fn buffer(&self) -> &'a DeviceBuffer { + self.buffer + } + /// Logical number of normalization weights, excluding trailing padding. + #[must_use] + pub const fn width(&self) -> usize { + WIDTH + } +} + +/// Borrowed BF16 bias for a projection with a specific role and output width. +pub struct BiasWeights<'a, Role: ProjectionRole, const WIDTH: usize> { + buffer: &'a DeviceBuffer, + role: PhantomData, +} + +impl<'a, Role: ProjectionRole, const WIDTH: usize> BiasWeights<'a, Role, WIDTH> { + /// Check a positive logical width and BF16 allocation extent, allowing guards. + pub fn new(buffer: &'a DeviceBuffer) -> Result { + RowMajorShape::dense(1, WIDTH)?.validate_len(buffer.len())?; + Ok(Self { + buffer, + role: PhantomData, + }) + } + /// Borrow the existing allocation for a checked kernel launch. + #[must_use] + pub const fn buffer(&self) -> &'a DeviceBuffer { + self.buffer + } + /// Logical number of bias elements, excluding trailing padding. + #[must_use] + pub const fn width(&self) -> usize { + WIDTH + } +} + +/// Role-checked projections and per-head normalization weights for attention. +/// +/// Query and key weights remain distinct even for multi-head attention where +/// their dimensions match: +/// ```compile_fail,E0308 +/// use cute_rs::ffn::{Key, ProjectionWeights, Query}; +/// fn wrong_role(query: ProjectionWeights<'_, Query, 8, 16>) { +/// let _: ProjectionWeights<'_, Key, 8, 16> = query; +/// } +/// ``` +pub struct AttentionWeights< + 'a, + const H: usize, + const QWIDTH: usize, + const KVWIDTH: usize, + const D: usize, +> { + query: ProjectionWeights<'a, Query, H, QWIDTH>, + key: ProjectionWeights<'a, Key, H, KVWIDTH>, + value: ProjectionWeights<'a, Value, H, KVWIDTH>, + output: ProjectionWeights<'a, AttentionOutput, QWIDTH, H>, + query_norm: NormWeights<'a, QueryNorm, D>, + key_norm: NormWeights<'a, KeyNorm, D>, +} + +impl<'a, const H: usize, const QWIDTH: usize, const KVWIDTH: usize, const D: usize> + AttentionWeights<'a, H, QWIDTH, KVWIDTH, D> +{ + /// Compose already checked weights; head-count consistency is checked by + /// [`DecoderShape::validate_weights`] before launch. + #[must_use] + pub const fn new( + query: ProjectionWeights<'a, Query, H, QWIDTH>, + key: ProjectionWeights<'a, Key, H, KVWIDTH>, + value: ProjectionWeights<'a, Value, H, KVWIDTH>, + output: ProjectionWeights<'a, AttentionOutput, QWIDTH, H>, + query_norm: NormWeights<'a, QueryNorm, D>, + key_norm: NormWeights<'a, KeyNorm, D>, + ) -> Self { + Self { + query, + key, + value, + output, + query_norm, + key_norm, + } + } + /// Query projection weights. + #[must_use] + pub const fn query(&self) -> &ProjectionWeights<'a, Query, H, QWIDTH> { + &self.query + } + /// Key projection weights. + #[must_use] + pub const fn key(&self) -> &ProjectionWeights<'a, Key, H, KVWIDTH> { + &self.key + } + /// Value projection weights. + #[must_use] + pub const fn value(&self) -> &ProjectionWeights<'a, Value, H, KVWIDTH> { + &self.value + } + /// Projection from query-head output width back to hidden width. + #[must_use] + pub const fn output(&self) -> &ProjectionWeights<'a, AttentionOutput, QWIDTH, H> { + &self.output + } + /// Shared per-query-head RMS weights. + #[must_use] + pub const fn query_norm(&self) -> &NormWeights<'a, QueryNorm, D> { + &self.query_norm + } + /// Shared per-key-head RMS weights. + #[must_use] + pub const fn key_norm(&self) -> &NormWeights<'a, KeyNorm, D> { + &self.key_norm + } +} + +/// Borrowed weights for one pre-norm decoder layer. +/// +/// Predictor hidden width 1024 cannot stand in for Talker hidden width 2048: +/// ```compile_fail,E0308 +/// use cute_rs::decoder::{DecoderShape, DecoderWeights}; +/// fn wrong_hidden(weights: &DecoderWeights<'_, 1024, 3072, 2048, 1024, 128>) { +/// let talker = DecoderShape::<2048, 6144, 16, 8, 128>::new(1, 2048).unwrap(); +/// talker.validate_weights(weights).unwrap(); +/// } +/// ``` +pub struct DecoderWeights< + 'a, + const H: usize, + const I: usize, + const QWIDTH: usize, + const KVWIDTH: usize, + const D: usize, +> { + attention: AttentionWeights<'a, H, QWIDTH, KVWIDTH, D>, + input_norm: NormWeights<'a, InputNorm, H>, + post_attention_norm: NormWeights<'a, PostAttentionNorm, H>, + ffn: FfnWeights<'a, H, I>, +} + +impl<'a, const H: usize, const I: usize, const QWIDTH: usize, const KVWIDTH: usize, const D: usize> + DecoderWeights<'a, H, I, QWIDTH, KVWIDTH, D> +{ + /// Compose role- and width-checked layer arguments without allocation. + #[must_use] + pub const fn new( + attention: AttentionWeights<'a, H, QWIDTH, KVWIDTH, D>, + input_norm: NormWeights<'a, InputNorm, H>, + post_attention_norm: NormWeights<'a, PostAttentionNorm, H>, + ffn: FfnWeights<'a, H, I>, + ) -> Self { + Self { + attention, + input_norm, + post_attention_norm, + ffn, + } + } + /// Attention projections and per-head normalization weights. + #[must_use] + pub const fn attention(&self) -> &AttentionWeights<'a, H, QWIDTH, KVWIDTH, D> { + &self.attention + } + /// Hidden-state normalization before attention. + #[must_use] + pub const fn input_norm(&self) -> &NormWeights<'a, InputNorm, H> { + &self.input_norm + } + /// Hidden-state normalization before the FFN. + #[must_use] + pub const fn post_attention_norm(&self) -> &NormWeights<'a, PostAttentionNorm, H> { + &self.post_attention_norm + } + /// Gate/up/down FFN projections. + #[must_use] + pub const fn ffn(&self) -> &FfnWeights<'a, H, I> { + &self.ffn + } +} + +/// Checked vocabulary-table geometry independent of the token sequence. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct EmbeddingShape { + table: RowMajorShape, +} + +impl EmbeddingShape { + /// Check a nonempty `[vocabulary, embedding_width]` table and its row pitch. + pub fn new(row_stride: usize) -> Result { + if VOCAB == 0 { + return Err(DecoderError::EmptyVocabulary); + } + Ok(Self { + table: RowMajorShape::new(VOCAB, WIDTH, row_stride)?, + }) + } + /// Logical table shape with checkpoint row orientation. + #[must_use] + pub const fn table(self) -> RowMajorShape { + self.table + } + /// Dense embedding output for a checked number of tokens. + pub fn output(self, tokens: usize) -> Result { + Ok(RowMajorShape::dense(tokens, WIDTH)?) + } +} + +/// Borrowed BF16 embedding table with checked vocabulary and embedding width. +pub struct EmbeddingWeights<'a, const VOCAB: usize, const WIDTH: usize> { + buffer: &'a DeviceBuffer, + shape: EmbeddingShape, +} + +impl<'a, const VOCAB: usize, const WIDTH: usize> EmbeddingWeights<'a, VOCAB, WIDTH> { + /// Borrow a `[vocabulary, width]` table, permitting trailing allocation guards. + pub fn new( + buffer: &'a DeviceBuffer, + row_stride: usize, + ) -> Result { + let shape = EmbeddingShape::new(row_stride)?; + shape.table().validate_len(buffer.len())?; + Ok(Self { buffer, shape }) + } + /// Existing table allocation. + #[must_use] + pub const fn buffer(&self) -> &'a DeviceBuffer { + self.buffer + } + /// Checked table geometry. + #[must_use] + pub const fn shape(&self) -> EmbeddingShape { + self.shape + } + /// Dense output extent for IDs checked against this exact vocabulary bound. + pub fn output_shape(&self, ids: TokenIds<'_, VOCAB>) -> Result { + self.shape.output(ids.len()) + } +} + +/// Immutably borrowed host token IDs checked against an exclusive vocab bound. +/// +/// These IDs are ready for an operator to copy to device memory. This type does +/// not certify an unrelated device ID buffer or a subsequent mutation to one. +/// +/// The host indices cannot change while the checked view is in use: +/// ```compile_fail,E0506 +/// use cute_rs::decoder::TokenIds; +/// let mut ids = [0, 1]; +/// let checked = TokenIds::<2>::new(&ids).unwrap(); +/// ids[0] = 99; +/// let _ = checked.as_slice(); +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TokenIds<'a, const VOCAB: usize> { + ids: &'a [u32], +} + +impl<'a, const VOCAB: usize> TokenIds<'a, VOCAB> { + /// Reject an empty vocabulary or the first token outside `[0, VOCAB)`. + pub fn new(ids: &'a [u32]) -> Result { + if VOCAB == 0 { + return Err(DecoderError::EmptyVocabulary); + } + for (index, &token) in ids.iter().enumerate() { + if token as usize >= VOCAB { + return Err(DecoderError::InvalidToken { + index, + token, + vocabulary: VOCAB, + }); + } + } + Ok(Self { ids }) + } + /// Checked host token indices. + #[must_use] + pub const fn as_slice(self) -> &'a [u32] { + self.ids + } + /// Number of token indices. + #[must_use] + pub const fn len(self) -> usize { + self.ids.len() + } + /// Whether the token sequence requires no embedding lookup. + #[must_use] + pub const fn is_empty(self) -> bool { + self.ids.is_empty() + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/mod.rs new file mode 100644 index 0000000000..1e443bef70 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/mod.rs @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Output staging operations grouped by architecture. +//! +//! The SM120 staging types and helpers are also exported here. The compiler +//! hooks have one path, `sm120::__compiler`. + +pub mod sm120; + +pub use sm120::{ + SM120_EPILOGUE_BYTES, SM120_EPILOGUE_COLS, SM120_EPILOGUE_ELEMENTS, SM120_EPILOGUE_HALF_BYTES, + SM120_EPILOGUE_HALF_COLS, SM120_EPILOGUE_HALF_ELEMENTS, SM120_EPILOGUE_ROWS, + Sm120Epilogue128x128, Sm120EpilogueHalfLayout, Sm120EpilogueWarp128x128, + sm120_epilogue_atom_origin, sm120_epilogue_physical_offset, +}; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/sm120.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/sm120.rs new file mode 100644 index 0000000000..0d11f4b380 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/epilogue/sm120.rs @@ -0,0 +1,395 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Store an SM120 128x128 result tile in shared memory. +//! +//! A thread is one GPU worker. A warp is 32 threads that run together. A CTA, +//! also called a block, is a group of warps that can use the same shared +//! memory. +//! A lane is a thread's position from 0 to 31 inside its warp. +//! +//! One tensor-core operation gives each thread four `f32` results in +//! [`FragmentC`]. Together, one warp owns a 16x8 result block. This +//! module calls that small block an *atom*. `stmatrix` is a warp-wide +//! instruction that stores packed register values into shared memory. +//! +//! ```text +//! each thread: [c0, c1, c2, c3] f32 +//! │ convert two at a time +//! ▼ +//! [c0,c1] [c2,c3] f16 +//! │ one stmatrix for the warp +//! ▼ +//! one 16x8 f16 atom in shared memory +//! ``` +//! +//! The Tensor Memory Accelerator (TMA) is hardware that copies complete +//! tiles. A TMA store using the 128-byte (B128) swizzle can move 64 `f16` +//! columns at once. The 128-column result therefore uses two adjacent 128x64 +//! halves. A swizzle is a fixed address rearrangement that reduces cases where +//! many threads request the same shared-memory bank and have to wait; each +//! half uses its own. +//! +//! ```text +//! logical 128x128 C tile +//! ┌──────── 128x64 ────────┬──────── 128x64 ────────┐ +//! │ left: 16 KiB │ right: 16 KiB │ +//! └────────────────────────┴────────────────────────┘ +//! shared allocation: 32 KiB total +//! ``` +//! +//! Kernel code uses the small methods below. At the compiler boundary they +//! stay visible as one result-tile flow: select the shared tile, select a +//! warp, store the accumulator, synchronize, and expose two TMA halves. + +use core::marker::PhantomData; + +use crate::copy::cooperative::SmemTile; +use crate::markers::{ComposedLayout, RowMajor, Swizzle}; +use crate::mma::FragmentC; + +/// Layout of one 128x64 half that a B128 TMA store can read. +/// +/// `S<3,3,3>` describes the address swizzle in `f16` elements. Because one +/// `f16` is two bytes, the same layout is TMA's byte-based `S<3,4,3>` B128 +/// mode. +pub type Sm120EpilogueHalfLayout = ComposedLayout, 0, RowMajor<128, 64>>; + +/// Rows in the full output tile. +pub const SM120_EPILOGUE_ROWS: usize = 128; +/// Columns in the full output tile. +pub const SM120_EPILOGUE_COLS: usize = 128; +/// Columns moved by one B128 TMA store. +pub const SM120_EPILOGUE_HALF_COLS: usize = 64; +/// `f16` elements in one 128x64 half. +pub const SM120_EPILOGUE_HALF_ELEMENTS: usize = SM120_EPILOGUE_ROWS * SM120_EPILOGUE_HALF_COLS; +/// Bytes in one 128x64 half. +pub const SM120_EPILOGUE_HALF_BYTES: usize = + SM120_EPILOGUE_HALF_ELEMENTS * core::mem::size_of::(); +/// `f16` elements in both halves. +pub const SM120_EPILOGUE_ELEMENTS: usize = SM120_EPILOGUE_ROWS * SM120_EPILOGUE_COLS; +/// Bytes in both halves. +pub const SM120_EPILOGUE_BYTES: usize = SM120_EPILOGUE_ELEMENTS * core::mem::size_of::(); + +/// Starting columns of the eight N8 atoms owned by one warp, before its N offset. +const N_ATOM_OFFSETS: [usize; 8] = [0, 8, 32, 40, 64, 72, 96, 104]; + +/// Top-left `(row, column)` of one warp's 16x8 result atom. +/// +/// Eight compute warps split the CTA tile. `warp_id % 4` selects one of four +/// 16-row positions. `warp_id / 4` selects the left or right 16-column group. +/// `m_band` adds either 0 or 64 rows. `n_slot` selects one of eight fixed +/// column offsets. +/// +/// ```text +/// row = 16*(warp_id % 4) + 64*m_band +/// col = 16*(warp_id / 4) + [0,8,32,40,64,72,96,104][n_slot] +/// ``` +/// +/// Expected inputs are `warp_id` in `0..8`, `m_band` in `0..2`, and +/// `n_slot` in `0..8`. +#[must_use] +pub const fn sm120_epilogue_atom_origin( + warp_id: usize, + m_band: usize, + n_slot: usize, +) -> (usize, usize) { + let warp_m = warp_id % 4; + let warp_n = warp_id / 4; + ( + 16 * warp_m + 64 * m_band, + 16 * warp_n + N_ATOM_OFFSETS[n_slot], + ) +} + +/// Convert logical `(row, column)` into an `f16` offset in shared memory. +/// +/// The address swizzle is applied separately inside each 128x64 half. This +/// helper is mainly for layout tests and debugging. Kernels normally call +/// [`Sm120EpilogueWarp128x128::store_atom`]. +/// +/// Expected inputs are `row` in `0..128` and `col` in `0..128`. +#[must_use] +pub const fn sm120_epilogue_physical_offset(row: usize, col: usize) -> usize { + let half = col / SM120_EPILOGUE_HALF_COLS; + let local_col = col % SM120_EPILOGUE_HALF_COLS; + let logical = row * SM120_EPILOGUE_HALF_COLS + local_col; + // S<3,3,3> changes address bits 3..5 using bits 6..8. It leaves bits + // 0..2 unchanged, so each 8-f16 (16-byte) stmatrix chunk stays together. + let swizzled = logical ^ ((logical & (0b111 << 6)) >> 3); + half * SM120_EPILOGUE_HALF_ELEMENTS + swizzled +} + +/// A typed view of the 32 KiB shared result buffer. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct Sm120Epilogue128x128 { + /// Base of the complete 32-KiB FP16 shared output tile. + base: *mut f16, +} + +impl Sm120Epilogue128x128 { + /// Required alignment of the shared allocation. + pub const ALIGNMENT: usize = 1024; + /// Size of the full shared allocation. + pub const BYTES: usize = SM120_EPILOGUE_BYTES; + + /// Treat an existing shared allocation as this result buffer. + /// + /// # Safety + /// + /// `base` must point to [`Self::BYTES`] writable bytes in shared memory + /// and be aligned to 1024 bytes. The allocation must remain valid until + /// every `stmatrix` write and asynchronous TMA read from it has finished. + #[must_use] + #[inline(always)] + pub const unsafe fn from_raw(base: *mut f16) -> Self { + unsafe { __compiler::epilogue_smem_overlay(base) } + } + + /// Select the part of the result tile owned by one compute warp. + /// + /// # Safety + /// + /// `warp_id` must be in `0..8`. `lane` must be the calling thread's + /// hardware lane number in `0..32`. All 32 lanes must select the same warp + /// partition, each with its own correct lane, and take the same + /// control-flow path for later stores. + #[must_use] + #[inline(always)] + pub const unsafe fn get_slice(self, warp_id: usize, lane: u32) -> Sm120EpilogueWarp128x128 { + unsafe { __compiler::epilogue_warp_slice(self, warp_id, lane) } + } + + /// Wait until the previous asynchronous store no longer reads this tile. + /// + /// # Safety + /// + /// All eight compute warps must call this together. The ninth, producer + /// warp must not participate in this counted barrier. `self` must refer + /// to the live shared epilogue allocation used by those warps. + #[inline(always)] + pub unsafe fn sync_reusable(self) { + unsafe { __compiler::epilogue_sync_reusable(self) }; + } + + /// Publish all compute-warp stores, then release the TMA issuer. + /// + /// This semantic boundary includes the generic-to-async shared proxy + /// publication required before TMA reads the tile. Each backend selects + /// its own fence and counted-barrier implementation. + /// + /// # Safety + /// + /// All eight compute warps must call this together. The producer warp + /// must not participate, and `self` must refer to their live shared + /// epilogue allocation. + #[inline(always)] + pub unsafe fn sync_ready_for_tma(self) { + unsafe { __compiler::epilogue_sync_ready_for_tma(self) }; + } + + /// View one 128x64 half as the source of a typed TMA store. + /// + /// # Safety + /// + /// `HALF` must be 0 for the left half or 1 for the right half. The caller + /// must provide the required synchronization and must not change this half + /// until the asynchronous TMA store has finished reading it. + #[must_use] + #[inline(always)] + pub unsafe fn tma_half(self) -> SmemTile { + unsafe { __compiler::epilogue_half::(self) } + } +} + +/// Stable compiler boundaries for the SM120 output-tile protocol. +/// +/// The public methods above are the readable source API. These exact helpers +/// let the importer preserve the same tile and synchronization story with or +/// without MIR inlining. Every returned value is still an ordinary Rust +/// pointer or scalar carrier. +#[doc(hidden)] +pub mod __compiler { + use super::*; + + /// Compiler boundary attaching the fixed output-tile interpretation to shared storage. + #[inline(never)] + pub const unsafe fn epilogue_smem_overlay(base: *mut f16) -> Sm120Epilogue128x128 { + Sm120Epilogue128x128 { base } + } + + /// Compiler boundary selecting one warp and lane’s output atoms. + #[inline(never)] + pub const unsafe fn epilogue_warp_slice( + tile: Sm120Epilogue128x128, + warp_id: usize, + lane: u32, + ) -> Sm120EpilogueWarp128x128 { + Sm120EpilogueWarp128x128 { + base: tile.base, + warp_id, + lane, + } + } + + /// Compiler boundary waiting for the old TMA reader before reusing the output tile. + /// + /// # Safety + /// Participants and barrier counts must satisfy Sm120Epilogue128x128::wait_reusable. + #[inline(never)] + pub unsafe fn epilogue_sync_reusable(tile: Sm120Epilogue128x128) { + let _ = tile; + unreachable!("cute-rs epilogue sync executed outside recognized device compilation") + } + + /// Compiler boundary publishing accumulator stores and releasing the TMA issuer. + /// + /// # Safety + /// Every output writer must follow the public epilogue publication and barrier protocol. + #[inline(never)] + pub unsafe fn epilogue_sync_ready_for_tma(tile: Sm120Epilogue128x128) { + let _ = tile; + unreachable!("cute-rs epilogue publication executed outside recognized device compilation") + } + + /// Compiler boundary exposing one 128-by-64 FP16 half to a matching TMA store. + #[inline(never)] + pub unsafe fn epilogue_half( + tile: Sm120Epilogue128x128, + ) -> SmemTile { + const { assert!(HALF < 2) }; + SmemTile { + base: unsafe { tile.base.add(HALF * SM120_EPILOGUE_HALF_ELEMENTS) }, + capacity: SM120_EPILOGUE_HALF_ELEMENTS, + layout: PhantomData, + } + } +} + +/// The result locations owned by one compute warp. +#[derive(Clone, Copy)] +pub struct Sm120EpilogueWarp128x128 { + /// Base of the complete shared output tile, in f16 elements. + base: *mut f16, + /// Compute-warp index selecting this warp’s output atoms. + warp_id: usize, + /// Lane index selecting this thread’s accumulator values. + lane: u32, +} + +impl Sm120EpilogueWarp128x128 { + /// Convert and store one complete 16x8 result atom. + /// + /// `M_BAND` is 0 or 1. `N_SLOT` is in `0..8`. Each lane's four values + /// land here inside the atom: + /// + /// ```text + /// c0,c1 -> row lane/4, columns 2*(lane%4) + {0,1} + /// c2,c3 -> row lane/4 + 8, columns 2*(lane%4) + {0,1} + /// ``` + /// + /// The two packed pairs are the exact input format required by the + /// non-transposed `stmatrix.m8n8.x2` instruction. + /// + /// # Safety + /// + /// All 32 lanes of the warp must call the same `M_BAND` and `N_SLOT` on + /// the same control-flow path. This view must have been created with the + /// calling hardware warp and lane. No other warp may write the same atom + /// at the same time. + #[inline(always)] + pub unsafe fn store_atom(&self, acc: FragmentC) { + const { + assert!(M_BAND < 2); + assert!(N_SLOT < 8); + } + + let _ = (self.base, self.warp_id, self.lane, acc); + unreachable!("cute-rs epilogue atom store executed outside recognized device compilation") + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + use crate::markers::ReifySmem2D; + use cuda_device::mma_frag; + + /// An epilogue half must have the shape, byte count and swizzle expected by one B128 TMA store. + #[test] + fn half_layout_is_exactly_one_b128_tma_tile() { + assert_eq!(::SHAPE, (128, 64)); + assert_eq!(::STRIDE, (64, 1)); + assert_eq!(::NUM_BITS, 3); + assert_eq!(::NUM_BASE, 3); + assert_eq!(::NUM_SHIFT, 3); + assert_eq!(SM120_EPILOGUE_HALF_BYTES, 16 * 1024); + assert_eq!(Sm120Epilogue128x128::BYTES, 32 * 1024); + } + + /// The eight compute warps must own every logical output atom exactly once. + #[test] + fn eight_warps_cover_the_logical_tile_once() { + let mut owners = [0u8; SM120_EPILOGUE_ELEMENTS]; + for warp in 0..8 { + for m_band in 0..2 { + for n_slot in 0..8 { + let (atom_row, atom_col) = sm120_epilogue_atom_origin(warp, m_band, n_slot); + for lane in 0..32 { + for j in 0..4 { + let (r, c) = mma_frag::acc_coords(lane, j); + let logical = (atom_row + r) * SM120_EPILOGUE_COLS + atom_col + c; + assert_eq!( + owners[logical], + 0, + "overlap at ({},{})", + atom_row + r, + atom_col + c + ); + owners[logical] = 1; + } + } + } + } + } + assert!(owners.iter().all(|&count| count == 1)); + } + + /// Each epilogue half must be a bijection and occupy storage disjoint from the other half. + #[test] + fn two_swizzled_halves_are_disjoint_bijections() { + let mut owners = [false; SM120_EPILOGUE_ELEMENTS]; + for row in 0..SM120_EPILOGUE_ROWS { + for col in 0..SM120_EPILOGUE_COLS { + let physical = sm120_epilogue_physical_offset(row, col); + assert!(physical < owners.len()); + assert!(!owners[physical], "physical collision at {physical}"); + owners[physical] = true; + } + } + assert!(owners.iter().all(|&claimed| claimed)); + assert_eq!(sm120_epilogue_physical_offset(0, 0), 0); + assert_eq!(sm120_epilogue_physical_offset(0, 64), 8192); + } + + /// Every lane-provided epilogue row address must satisfy stmatrix’s 16-byte alignment. + #[test] + fn every_stmatrix_row_address_is_sixteen_byte_aligned() { + for warp in 0..8 { + for m_band in 0..2 { + for n_slot in 0..8 { + let (atom_row, atom_col) = sm120_epilogue_atom_origin(warp, m_band, n_slot); + for address_lane in 0..16 { + let physical = + sm120_epilogue_physical_offset(atom_row + address_lane, atom_col); + assert_eq!((physical * core::mem::size_of::()) % 16, 0); + } + } + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/ffn.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/ffn.rs new file mode 100644 index 0000000000..db28ed2ec3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/ffn.rs @@ -0,0 +1,461 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Checked host arguments for BF16 projections and feed-forward compositions. +//! +//! These types describe existing allocations; they neither allocate nor launch +//! work. Matrix dimensions and pitches are checked once at construction. Weight +//! roles and projection dimensions remain in Rust types when composing an FFN. +//! All weights retain checkpoint orientation: `[output, input]`, with contiguous +//! input coordinates within each row. No implicit transpose or conversion occurs. +//! +//! Assigning a role asserts the meaning of the supplied weights. Shapes cannot +//! establish that an importer labeled their contents correctly. Buffer context, +//! stream ordering, kernel participation and disjoint output writes remain launch +//! obligations; these argument types do not prove them. +//! +//! ``` +//! use cuda_core::DeviceBuffer; +//! use cute_rs::BFloat16; +//! use cute_rs::ffn::{Down, FfnError, FfnWeights, Gate, ProjectionWeights, Up}; +//! +//! fn checked<'a>( +//! gate: &'a DeviceBuffer, +//! up: &'a DeviceBuffer, +//! down: &'a DeviceBuffer, +//! ) -> Result, FfnError> { +//! Ok(FfnWeights::new( +//! ProjectionWeights::::new(gate, 1024)?, +//! ProjectionWeights::::new(up, 1024)?, +//! ProjectionWeights::::new(down, 3072)?, +//! )) +//! } +//! ``` + +use core::{fmt, marker::PhantomData, mem::size_of}; +use cuda_core::DeviceBuffer; + +use crate::numeric::BFloat16; + +/// A rejected BF16 matrix shape or allocation extent. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FfnError { + /// A matrix width or projection dimension is zero. + ZeroWidth, + /// Consecutive rows overlap because their pitch is smaller than the width. + StrideTooSmall, + /// An element calculation overflowed or the BF16 byte extent exceeds `isize::MAX`. + ExtentOverflow, + /// The allocation does not cover every logical matrix coordinate. + BufferTooShort, +} + +impl fmt::Display for FfnError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::ZeroWidth => "FFN matrix widths and projection dimensions must be positive", + Self::StrideTooSmall => "matrix row stride is smaller than its column count", + Self::ExtentOverflow => { + "BF16 matrix or slice byte extent exceeds the addressable range" + } + Self::BufferTooShort => "buffer does not cover the complete matrix extent", + }) + } +} + +impl core::error::Error for FfnError {} + +/// Validated row-major BF16 coordinates, with the pitch measured in elements. +/// +/// The minimum allocation ends at the last logical element: a `3 x 5` matrix +/// with pitch `8` requires `21` elements. Extra trailing padding is permitted. +/// Empty row sets require no elements, but their width and pitch remain valid. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RowMajorShape { + /// Number of logical rows. + rows: usize, + /// Number of logical columns in every row. + cols: usize, + /// Distance in BF16 elements between consecutive row starts. + row_stride: usize, + /// Minimum buffer extent including padding between rows. + required_len: usize, +} + +impl RowMajorShape { + /// Check a row-major BF16 shape without allocating any memory. + /// + /// Rejects zero width, overlapping rows, overflowing element offsets and + /// byte extents larger than `isize::MAX`, as required by Rust slice indexing. + pub fn new(rows: usize, cols: usize, row_stride: usize) -> Result { + if cols == 0 { + return Err(FfnError::ZeroWidth); + } + if row_stride < cols { + return Err(FfnError::StrideTooSmall); + } + let required_len = if rows == 0 { + 0 + } else { + (rows - 1) + .checked_mul(row_stride) + .and_then(|last_row| last_row.checked_add(cols)) + .ok_or(FfnError::ExtentOverflow)? + }; + check_byte_extent(required_len)?; + Ok(Self { + rows, + cols, + row_stride, + required_len, + }) + } + + /// Check a row-major BF16 shape whose rows have no padding. + pub fn dense(rows: usize, cols: usize) -> Result { + Self::new(rows, cols, cols) + } + + /// Number of logical rows. + #[must_use] + pub const fn rows(self) -> usize { + self.rows + } + + /// Number of logical columns in each row. + #[must_use] + pub const fn cols(self) -> usize { + self.cols + } + + /// Distance in BF16 elements between consecutive row starts. + #[must_use] + pub const fn row_stride(self) -> usize { + self.row_stride + } + + /// Minimum number of BF16 elements covering every logical coordinate. + #[must_use] + pub const fn required_len(self) -> usize { + self.required_len + } + + /// Check an existing buffer's element count, permitting trailing padding. + /// + /// The full slice extent, including any trailing padding, must fit within + /// `isize::MAX` bytes. This check does not establish buffer context or lifetime. + pub fn validate_len(self, len: usize) -> Result<(), FfnError> { + check_byte_extent(len)?; + if len < self.required_len { + Err(FfnError::BufferTooShort) + } else { + Ok(()) + } + } +} + +/// Reject BF16 buffers whose complete byte extent cannot be represented by a Rust slice. +fn check_byte_extent(len: usize) -> Result<(), FfnError> { + match len.checked_mul(size_of::()) { + Some(bytes) if bytes <= isize::MAX as usize => Ok(()), + _ => Err(FfnError::ExtentOverflow), + } +} + +/// Private supertrait prevents unreviewed projection roles. +mod sealed { + /// Marks the supported projection weight roles. + pub trait Sealed {} +} + +/// A sealed semantic role for projection weights and their optional bias. +pub trait ProjectionRole: sealed::Sealed {} + +/// Weights producing the gated branch before SiLU. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Gate; + +/// Weights producing the branch multiplied by the activated gate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Up; + +/// Weights projecting the multiplied intermediate values back to hidden width. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Down; + +/// Weights projecting hidden states into query heads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Query; + +/// Weights projecting hidden states into key heads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Key; + +/// Weights projecting hidden states into value heads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Value; + +/// Weights projecting concatenated attention heads back to hidden width. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AttentionOutput; + +/// Weights mapping an embedding width into a decoder's hidden width. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct InputProjection; + +/// Weights mapping decoder hidden states back into the external latent width. +/// +/// The input and output adapters remain distinct even when widths happen to match: +/// +/// ```compile_fail,E0308 +/// use cute_rs::ffn::{InputProjection, OutputProjection, ProjectionWeights}; +/// fn wrong_role(weight: ProjectionWeights<'_, InputProjection, 16, 16>) { +/// let _: ProjectionWeights<'_, OutputProjection, 16, 16> = weight; +/// } +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OutputProjection; + +/// Weights expanding channels before a pointwise channel activation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ChannelExpand; + +/// Weights contracting activated channels back to the residual width. +/// +/// ```compile_fail,E0308 +/// use cute_rs::ffn::{ChannelExpand, ChannelContract, ProjectionWeights}; +/// fn wrong_role(weight: ProjectionWeights<'_, ChannelExpand, 16, 16>) { +/// let _: ProjectionWeights<'_, ChannelContract, 16, 16> = weight; +/// } +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ChannelContract; + +/// Weights for the first projection in a text input adapter. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TextProjection1; + +/// Weights for the second projection in a text input adapter. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TextProjection2; + +/// Weights projecting hidden states into vocabulary logits. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Logits; + +impl sealed::Sealed for Gate {} +impl sealed::Sealed for Up {} +impl sealed::Sealed for Down {} +impl ProjectionRole for Gate {} +impl ProjectionRole for Up {} +impl ProjectionRole for Down {} +impl sealed::Sealed for Query {} +impl sealed::Sealed for Key {} +impl sealed::Sealed for Value {} +impl sealed::Sealed for AttentionOutput {} +impl sealed::Sealed for InputProjection {} +impl sealed::Sealed for OutputProjection {} +impl sealed::Sealed for ChannelExpand {} +impl sealed::Sealed for ChannelContract {} +impl sealed::Sealed for TextProjection1 {} +impl sealed::Sealed for TextProjection2 {} +impl sealed::Sealed for Logits {} +impl ProjectionRole for Query {} +impl ProjectionRole for Key {} +impl ProjectionRole for Value {} +impl ProjectionRole for AttentionOutput {} +impl ProjectionRole for InputProjection {} +impl ProjectionRole for OutputProjection {} +impl ProjectionRole for ChannelExpand {} +impl ProjectionRole for ChannelContract {} +impl ProjectionRole for TextProjection1 {} +impl ProjectionRole for TextProjection2 {} +impl ProjectionRole for Logits {} + +/// Borrowed BF16 projection weights with a checked `[OUTPUT, INPUT]` layout. +/// +/// The role and dimensions prevent incompatible weights from being passed to a +/// checked FFN composition. The owning [`DeviceBuffer`] retains allocation and +/// context lifetime; this wrapper never assumes ownership or converts its data. +/// +/// An FP32 allocation is not BF16 storage: +/// +/// ```compile_fail,E0308 +/// use cuda_core::DeviceBuffer; +/// use cute_rs::ffn::{Gate, ProjectionWeights}; +/// fn wrong_dtype(weights: &DeviceBuffer) { +/// let _ = ProjectionWeights::::new(weights, 8); +/// } +/// ``` +/// +/// The weight borrow prevents mutable access while the checked view is used: +/// +/// ```compile_fail,E0502 +/// use cuda_core::DeviceBuffer; +/// use cute_rs::BFloat16; +/// use cute_rs::ffn::{Gate, ProjectionWeights}; +/// fn exclusive_use(_: &mut DeviceBuffer) {} +/// fn alias(weights: &mut DeviceBuffer) { +/// let gate = ProjectionWeights::::new(weights, 8).unwrap(); +/// exclusive_use(weights); +/// let _ = gate.buffer(); +/// } +/// ``` +pub struct ProjectionWeights<'a, Role: ProjectionRole, const INPUT: usize, const OUTPUT: usize> { + /// The existing, immutably borrowed allocation. + buffer: &'a DeviceBuffer, + /// Checked checkpoint orientation and row pitch. + shape: RowMajorShape, + /// Distinguishes equally shaped weights with different semantic roles. + role: PhantomData, +} + +impl<'a, Role: ProjectionRole, const INPUT: usize, const OUTPUT: usize> + ProjectionWeights<'a, Role, INPUT, OUTPUT> +{ + /// Borrow weights in checkpoint orientation, with input coordinates contiguous. + /// + /// Both projection dimensions must be positive. The caller explicitly labels + /// the role; this validates dimensions and capacity, not semantic contents. + pub fn new(buffer: &'a DeviceBuffer, row_stride: usize) -> Result { + if OUTPUT == 0 { + return Err(FfnError::ZeroWidth); + } + let shape = RowMajorShape::new(OUTPUT, INPUT, row_stride)?; + shape.validate_len(buffer.len())?; + Ok(Self { + buffer, + shape, + role: PhantomData, + }) + } + + /// Borrow the existing allocation for a checked kernel launch. + #[must_use] + pub const fn buffer(&self) -> &'a DeviceBuffer { + self.buffer + } + + /// The checked `[OUTPUT, INPUT]` shape and its physical row pitch. + #[must_use] + pub const fn shape(&self) -> RowMajorShape { + self.shape + } +} + +/// Gate/up/down weights whose roles and dimensions agree for one FFN. +/// +/// Equal dimensions do not make gate and up weights interchangeable: +/// +/// ```compile_fail,E0308 +/// use cute_rs::ffn::{Down, FfnWeights, Gate, ProjectionWeights, Up}; +/// fn swapped( +/// gate: ProjectionWeights<'_, Gate, 8, 16>, +/// up: ProjectionWeights<'_, Up, 8, 16>, +/// down: ProjectionWeights<'_, Down, 16, 8>, +/// ) { +/// let _ = FfnWeights::<8, 16>::new(up, gate, down); +/// } +/// ``` +/// +/// The down projection must reverse the hidden/intermediate dimensions: +/// +/// ```compile_fail,E0308 +/// use cute_rs::ffn::{Down, FfnWeights, Gate, ProjectionWeights, Up}; +/// fn wrong_down_orientation( +/// gate: ProjectionWeights<'_, Gate, 8, 16>, +/// up: ProjectionWeights<'_, Up, 8, 16>, +/// down: ProjectionWeights<'_, Down, 8, 16>, +/// ) { +/// let _ = FfnWeights::<8, 16>::new(gate, up, down); +/// } +/// ``` +pub struct FfnWeights<'a, const HIDDEN: usize, const INTERMEDIATE: usize> { + /// Gate projection from hidden to intermediate width. + gate: ProjectionWeights<'a, Gate, HIDDEN, INTERMEDIATE>, + /// Up projection from hidden to intermediate width. + up: ProjectionWeights<'a, Up, HIDDEN, INTERMEDIATE>, + /// Down projection from intermediate to hidden width. + down: ProjectionWeights<'a, Down, INTERMEDIATE, HIDDEN>, +} + +impl<'a, const HIDDEN: usize, const INTERMEDIATE: usize> FfnWeights<'a, HIDDEN, INTERMEDIATE> { + /// Compose previously checked weights; Rust verifies their roles and dimensions. + #[must_use] + pub const fn new( + gate: ProjectionWeights<'a, Gate, HIDDEN, INTERMEDIATE>, + up: ProjectionWeights<'a, Up, HIDDEN, INTERMEDIATE>, + down: ProjectionWeights<'a, Down, INTERMEDIATE, HIDDEN>, + ) -> Self { + Self { gate, up, down } + } + + /// Weights for the gate projection. + #[must_use] + pub const fn gate(&self) -> &ProjectionWeights<'a, Gate, HIDDEN, INTERMEDIATE> { + &self.gate + } + + /// Weights for the up projection. + #[must_use] + pub const fn up(&self) -> &ProjectionWeights<'a, Up, HIDDEN, INTERMEDIATE> { + &self.up + } + + /// Weights for the down projection. + #[must_use] + pub const fn down(&self) -> &ProjectionWeights<'a, Down, INTERMEDIATE, HIDDEN> { + &self.down + } +} + +/// Checked input and workspace extents for fixed hidden/intermediate widths. +/// +/// Input rows may be padded. Hidden-width outputs and intermediate workspace are +/// dense. Empty row sets remain valid so callers can skip all allocations/launches. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FfnShape { + /// The input matrix, including its caller-selected pitch. + input: RowMajorShape, + /// Dense hidden-width output or normalization workspace. + hidden: RowMajorShape, + /// Dense intermediate-width projection/activation workspace. + intermediate: RowMajorShape, +} + +impl FfnShape { + /// Check all input and workspace extents before allocation. + pub fn new(rows: usize, input_stride: usize) -> Result { + Ok(Self { + input: RowMajorShape::new(rows, HIDDEN, input_stride)?, + hidden: RowMajorShape::dense(rows, HIDDEN)?, + intermediate: RowMajorShape::dense(rows, INTERMEDIATE)?, + }) + } + + /// Number of logical token rows in every activation matrix. + #[must_use] + pub const fn rows(self) -> usize { + self.input.rows() + } + + /// Checked input shape, including any row padding. + #[must_use] + pub const fn input(self) -> RowMajorShape { + self.input + } + + /// Dense shape for hidden-width outputs and normalization workspace. + #[must_use] + pub const fn hidden(self) -> RowMajorShape { + self.hidden + } + + /// Dense shape for projection outputs and multiplied activations. + #[must_use] + pub const fn intermediate(self) -> RowMajorShape { + self.intermediate + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/lib.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/lib.rs new file mode 100644 index 0000000000..47f682788c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/lib.rs @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CuTe-style building blocks for GPU kernels written in Rust. +//! +//! A tensor view describes storage without loading it. Split it into tiles, +//! then select the part one thread will process: +//! +//! ``` +//! use cute_rs::Tensor; +//! let values = [0.0f32; 10]; +//! let last = Tensor::from_slice(&values).zipped_divide::<4>().slice(2); +//! assert_eq!(last.base(), 8); +//! assert!(!last.is_full()); // Two values remain; use scalar loads for this tile. +//! ``` +//! +//! See `cuda-oxide/cute-rs-preview/examples/elementwise_cute` in the repository for a +//! complete load/add/store kernel and launch. +//! Host descriptor constructors require the `host` feature; device operations +//! require cuda-oxide's device compiler. +//! +//! # Modules +//! +//! - [`tensor`] describes storage, register tiles and thread partitions. +//! - [`copy`] groups cooperative, tiled, vector and TMA transfers. +//! - [`mma`] contains common fragments and the BF16 and block-scaled families. +//! - [`attention`] groups host metadata and device-available SM100 operations. +//! - [`epilogue`] contains SM120 output staging. +//! +//! Every module has one path, its defining path. A root item export such as +//! [`TensorMap`] names the same item as its defining module, here +//! [`copy::tma::TensorMap`]. Architecture ownership and pipeline modules +//! retain their existing roles. +//! +//! The typed wrappers are inlined, so they add no device-side storage or call +//! overhead. Functions such as [`load_tile`] and [`copy_g2s`] are compiler +//! entry points: during device compilation, the compiler replaces each call +//! with the matching CuTe operation. +//! +//! Copy operations come in two sizes: +//! +//! - [`load_tile`] and [`store_tile`] move one thread's short vector. +//! - [`copy_g2s`] lets every thread in a block help move one larger tile from +//! global memory to shared memory. +//! +//! [`assume`] tells the compiler that an index has the alignment needed by +//! a later operation. +//! +//! # Small glossary +//! +//! ```text +//! thread one running lane +//! warp 32 threads that run together +//! block / CTA a group of threads that can share shared memory +//! tile a small rectangular part of a tensor +//! fragment one thread's register-held part of a tile +//! MMA a tensor-core matrix multiply-and-add operation +//! TMA hardware that copies complete multidimensional tiles +//! +//! global memory ──► shared memory ──► registers +//! whole GPU one block one thread +//! ``` + +#![no_std] +#![feature(f16)] +#![deny(missing_docs)] + +#[cfg(feature = "host")] +extern crate alloc; + +pub mod arch; +pub mod attention; +pub mod block_scaled; +pub mod copy; +#[cfg(feature = "host")] +pub mod decoder; +pub mod epilogue; +#[cfg(feature = "host")] +pub mod ffn; +pub mod markers; +pub mod mma; +pub mod numeric; +pub mod pipeline; +pub mod reduction; +pub mod scheduler; +pub mod sm100; +pub mod static_layout; +pub mod tensor; + +pub use block_scaled::{ + BlockScaledTensor, BlockScaledThreadRow, BlockScaledTile64, DenseBlockScaleKMajor, KMajor, + LoadedBlockScaledTile64, LoadedNvfp4Tile64, Mkl, MmaScalePair, MmaScaleQuad, + Mxf4BlockScaledTensor, Nkl, Nvfp4BlockScaledTensor, Nvfp4ThreadRow, Nvfp4Tile64, ScalePack4, + SharedScaleAtom, Sm1xxBlockScaleKMajor, Sm120ScaleAtom, +}; +pub use copy::cooperative::{GmemMatrix, SmemTile, copy_g2s}; +pub use copy::tiled::{GlobalCopyTensor, SharedTensor, ThrCopy, TiledCopy}; +pub use copy::tma::{TensorMap, copy_tma_2d, copy_tma_s2g_2d, size_in_bytes}; +pub use copy::vector::{ + GlobalStoreElement, GlobalStoreTile, assume, load_tile, store_global_no_allocate, store_tile, +}; +pub use cute_layout as layout; +pub use epilogue::{ + SM120_EPILOGUE_BYTES, SM120_EPILOGUE_COLS, SM120_EPILOGUE_ELEMENTS, SM120_EPILOGUE_HALF_BYTES, + SM120_EPILOGUE_HALF_COLS, SM120_EPILOGUE_HALF_ELEMENTS, SM120_EPILOGUE_ROWS, + Sm120Epilogue128x128, Sm120EpilogueHalfLayout, Sm120EpilogueWarp128x128, + sm120_epilogue_atom_origin, sm120_epilogue_physical_offset, +}; +pub use markers::{C, ColMajor, ComposedLayout, CopyG2S, Layout, LeadingDim, RowMajor, Swizzle}; +pub use mma::block_scaled::{ + Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxf4ScalePairs, Mxf4ScalePairs128, Mxf4ScaleStage, + Mxf4ScaleTile128, Mxf4TiledMma, +}; +pub use mma::{ + FragmentA, FragmentB, FragmentBPair, FragmentC, load_matrix_a, load_matrix_b, + load_matrix_b_pair, load_matrix_b_pair_k_major, +}; +pub use numeric::{ + BFloat16, E4M3x2, E4M3x4, Float4E2M1FN, Float4E2M1FNx2, Float8E4M3FN, Float8E8M0FNU, Mxf4E2M1, + UE8M0x2, UE8M0x4, e4m3_to_f32, +}; +pub use pipeline::{Consumer, PipelineState, PipelineTmaAsync, PipelineTmaStore, Producer}; +pub use scheduler::{StaticPersistentTileScheduler, WorkTileInfo}; +pub use sm100::{ + OperandA, OperandB, PipelineTmaUmma, PipelineUmmaAsync, Sm100ClusterTmaCopy, Sm100SharedTile, + Sm100TiledMma, Sm100Tmem, Sm100TmemEpilogue, +}; +pub use tensor::{Contiguous1D, RegisterTile, Tensor, TensorElement, TensorMut, Tile1D, Zipped1D}; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/markers.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/markers.rs new file mode 100644 index 0000000000..7208812d32 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/markers.rs @@ -0,0 +1,229 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Layout facts stored in Rust types. +//! +//! These marker types have no runtime data. They let the compiler see a +//! layout directly in a function's type arguments. +//! +//! Most code can use a short form such as `RowMajor<32, 8>`. The full form +//! stores the same shape and strides explicitly, with Rust tuples grouping +//! the modes: +//! +//! ```text +//! 32 rows × 8 columns, row-major +//! +//! Layout< (C<32>, C<8>), (C<8>, C<1>) > +//! └── shape ──┘ └─ stride ─┘ +//! row: 8 +//! column: 1 +//! ``` + +use core::marker::PhantomData; + +/// One integer stored in a type, usually a size or stride. +pub struct C; + +/// A layout made from a `Shape` and its `Stride`. +/// +/// A stride is the number of elements to skip when a coordinate increases by +/// one. Nested modes are Rust tuples of [`C`] leaves, which can themselves be +/// nested: `(A, (B, C))` and `((A, B), C)` describe different layouts. +/// +/// This zero-sized marker is read from a function's type arguments. The +/// runtime value with the same shape and stride is [`cute_layout::Layout`], +/// reachable as `cute_rs::layout::Layout`. +pub struct Layout( + #[doc = "Carries corresponding shape and stride descriptions."] PhantomData<(Shape, Stride)>, +); + +/// A `ROWS × COLS` matrix stored one complete row after another. +/// +/// ```text +/// [a00 a01 a02] [a10 a11 a12] +/// ``` +pub struct RowMajor; + +/// A `ROWS × COLS` matrix stored one complete column after another. +/// +/// ```text +/// [a00 a10] [a01 a11] [a02 a12] +/// ``` +pub struct ColMajor; + +/// One global-to-shared copy atom of `BYTES` bytes, issued as `cp.async` +/// (CuTeDSL `cpasync.CopyG2SOp`). +pub struct CopyG2S; + +/// The distance between rows, with a compile-time alignment promise. +/// +/// `elements` is known only while the kernel runs. The type promises that it +/// is divisible by `DIVISOR`: +/// +/// ```text +/// LeadingDim<4> { elements: 132 } +/// │ │ +/// │ └─ 132 elements between row starts +/// └──────────────── 132 % 4 == 0 +/// ``` +/// +/// It intentionally has no `new` helper. Build it with a struct literal so +/// the promise stays visible at the copy call where the compiler needs it. +#[repr(transparent)] +pub struct LeadingDim { + /// Number of elements from the start of one row to the next. + pub elements: usize, +} + +/// Internal link between [`LeadingDim`] and block-wide copy operations. +/// +/// This trait is public only because it appears in a public function bound. +/// Users do not implement it; the compiler accepts [`LeadingDim`] only. +#[doc(hidden)] +pub trait LeadingDimMarker {} + +impl LeadingDimMarker for LeadingDim {} + +/// A CuTe XOR swizzle written as `Swizzle`. +/// +/// A swizzle changes shared-memory addresses so threads are less likely to +/// request the same memory bank at once and be forced to wait: +/// +/// ```text +/// normal address ── XOR selected bits ──► swizzled address +/// ``` +/// +/// - `B`: number of bits XORed. +/// - `M`: low bits left unchanged. +/// - `S`: distance from those low bits to the second XOR field. +/// +/// At this API boundary, the bit positions count elements. The compiler +/// converts them to byte positions after it knows the element type. +pub struct Swizzle; + +/// One address rule built by applying `Outer`, adding `OFFSET`, then applying +/// `Inner`: `Inner(OFFSET + Outer(coordinate))`. +/// +/// ```text +/// coordinate ── Outer ──► element offset +/// + OFFSET +/// │ +/// ▼ +/// Inner ──► final address +/// ``` +/// +/// The names are CuTe's. `Inner` is the swizzle, the map on the outside of +/// the expression that runs last; `Outer` is the affine layout that runs +/// first. The type arguments sit in the order of CuTeDSL's +/// `make_composed_layout(inner, offset, outer)`, so +/// `ComposedLayout, 0, RowMajor<128, 64>>` reads like that +/// call. +/// +/// `OFFSET` and `Outer` count elements here. The compiler converts the whole +/// rule to bytes after it knows the element size. The runtime value with the +/// same parts is [`cute_layout::ComposedLayout`], reachable as +/// `cute_rs::layout::ComposedLayout`. +pub struct ComposedLayout( + #[doc = "Carries the inner swizzle and the outer layout."] PhantomData<(Inner, Outer)>, +); + +/// Makes a two-dimensional shared-memory layout readable by host code. +/// +/// The device compiler reads the layout from its Rust type. Host code needs +/// ordinary constants to build a TMA descriptor. This trait gives both sides +/// the same values: +/// +/// ```text +/// one layout type +/// ├──► device compiler +/// └──► host TMA descriptor +/// ``` +/// +/// The swizzle constants carry the field names of `cute_layout::Swizzle`, +/// which are CuTeDSL's `Swizzle.num_bits`, `num_base`, and `num_shift`. +pub trait ReifySmem2D { + /// Tile extents as `(rows, columns)`. + const SHAPE: (i64, i64); + /// Elements skipped when the row, then the column, increases by one. + const STRIDE: (i64, i64); + /// Fixed starting offset in elements; `0` for a plain layout. + const OFFSET: i64; + /// Swizzle `B`, the number of XOR-ed bits, in element units; `0` means no + /// swizzle. + const NUM_BITS: u32; + /// Swizzle `M`: the number of low address bits left unchanged. + const NUM_BASE: u32; + /// Swizzle `S`: the distance between the two address fields. + const NUM_SHIFT: i32; +} + +impl ReifySmem2D + for Layout<(C, C), (C, C)> +{ + const SHAPE: (i64, i64) = (R, CC); + const STRIDE: (i64, i64) = (SR, SC); + const OFFSET: i64 = 0; + const NUM_BITS: u32 = 0; + const NUM_BASE: u32 = 4; + const NUM_SHIFT: i32 = 3; +} + +impl ReifySmem2D for RowMajor { + const SHAPE: (i64, i64) = (R, CC); + const STRIDE: (i64, i64) = (CC, 1); + const OFFSET: i64 = 0; + const NUM_BITS: u32 = 0; + const NUM_BASE: u32 = 4; + const NUM_SHIFT: i32 = 3; +} + +impl ReifySmem2D for ColMajor { + const SHAPE: (i64, i64) = (R, CC); + const STRIDE: (i64, i64) = (1, R); + const OFFSET: i64 = 0; + const NUM_BITS: u32 = 0; + const NUM_BASE: u32 = 4; + const NUM_SHIFT: i32 = 3; +} + +/// Read one composed layout. Do not nest `ComposedLayout`: this implementation +/// reports only its own inner swizzle and would hide a nested one. +impl ReifySmem2D + for ComposedLayout, OFFSET, Outer> +{ + const SHAPE: (i64, i64) = Outer::SHAPE; + const STRIDE: (i64, i64) = Outer::STRIDE; + const OFFSET: i64 = OFFSET; + const NUM_BITS: u32 = B; + const NUM_BASE: u32 = M; + const NUM_SHIFT: i32 = S; +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Nested layout, swizzle and copy markers must occupy zero runtime bytes. + #[test] + fn markers_are_compile_time_only() { + /// Two-dimensional 32-by-8 marker shape used to test zero-sized composition. + type Shape = (C<32>, C<8>); + /// Row-major marker strides corresponding to the test shape. + type Stride = (C<8>, C<1>); + /// Pairs the test’s shape and strides without runtime storage. + type Affine = Layout; + /// Adds a swizzle to verify composed markers also remain zero-sized. + type Smem = ComposedLayout, 0, Affine>; + + assert_eq!(core::mem::size_of::>(), 0); + assert_eq!(core::mem::size_of::(), 0); + assert_eq!(core::mem::size_of::>(), 0); + assert_eq!( + core::mem::size_of::>(), + core::mem::size_of::() + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/bf16.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/bf16.rs new file mode 100644 index 0000000000..cd183cf814 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/bf16.rs @@ -0,0 +1,392 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Typed BF16 m16n8k16 operand partitions and FP32 accumulation. +//! +//! Static thread/value layouts determine every operand and result coordinate. +//! Load callbacks can describe ordinary matrices or virtual im2col views. +//! CTM supplies only the final instruction and packed physical loads. +use crate::markers::{C, Layout}; +use crate::numeric::BFloat16; +use crate::static_layout::StaticLayout; +use core::marker::PhantomData; + +type Threads = (C<4>, C<8>); +/// A's `(lane, scalar value)` map into a row-major 16-by-16 logical tile. +pub type ALayout = Layout<(Threads, ((C<2>, C<2>), C<2>)), ((C<2>, C<16>), ((C<1>, C<128>), C<8>))>; +/// B's `(lane, scalar value)` map into the stored 8-by-16 `(N,K)` tile. +pub type BLayout = Layout<(Threads, (C<2>, C<2>)), ((C<2>, C<16>), (C<1>, C<8>))>; +/// C's `(lane, register)` map into a row-major 16-by-8 logical tile. +pub type CLayout = Layout<(Threads, (C<2>, C<2>)), ((C<2>, C<8>), (C<1>, C<64>))>; + +/// This lane's BF16 A registers, whose layout is established by its partition. +#[derive(Clone, Copy)] +pub struct FragmentA(ctm_rs::mma::Bf16A); +/// This lane's BF16 B registers, distinct from an A operand. +#[derive(Clone, Copy)] +pub struct FragmentB(ctm_rs::mma::Bf16B); +/// This lane's FP32 accumulator in the atom's C layout. +#[derive(Clone, Copy)] +pub struct FragmentC(ctm_rs::mma::Accumulator16x8); +impl FragmentC { + /// Positive FP32 zero in every register, for starting one warp product. + pub const ZERO: Self = Self(ctm_rs::mma::Accumulator16x8::zero()); +} + +/// One lane's handle to the BF16 16-by-8-by-16 matrix operation. +#[derive(Clone, Copy)] +pub struct Bf16Mma16x8x16 { + lane: usize, +} +impl Bf16Mma16x8x16 { + /// Select a lane in the 32-thread atom. Device use must pass the actual lane. + #[inline(always)] + pub fn get_slice(lane: usize) -> Self { + assert!(lane < 32); + Self { lane } + } + /// Partition A's sixteen rows and sixteen reduction coordinates. + #[inline(always)] + pub fn partition_a(self) -> APartition { + APartition { lane: self.lane } + } + /// Partition B's eight stored output rows and sixteen reduction coordinates. + #[inline(always)] + pub fn partition_b(self) -> BPartition { + BPartition { lane: self.lane } + } + /// Partition the sixteen-by-eight result tile. + #[inline(always)] + pub fn partition_c(self) -> CPartition { + CPartition { lane: self.lane } + } + /// Accumulate one BF16 product, with no conversion or reassociation. + /// + /// # Safety + /// All 32 lanes must execute together on SM80 or newer, using their actual + /// lane handles and A/B/C fragments of the same logical warp product. + #[must_use] + #[inline(always)] + pub unsafe fn mma_atom_call(self, c: FragmentC, a: FragmentA, b: FragmentB) -> FragmentC { + FragmentC(unsafe { ctm_rs::mma::mma_bf16(c.0, a.0, b.0) }) + } +} +/// A lane's immutable partition of the A operand. +#[derive(Clone, Copy)] +pub struct APartition { + lane: usize, +} +/// A lane's immutable partition of the B operand. +#[derive(Clone, Copy)] +pub struct BPartition { + lane: usize, +} +/// A lane's partition of the accumulator coordinates. +#[derive(Clone, Copy)] +pub struct CPartition { + lane: usize, +} + +/// Borrowed storage proof for one/two-row BF16 decode with complete K/N tiles. +/// +/// Construction checks bounds, even row pitches and four-byte alignment once. +/// The borrow retains operand storage; it does not establish GPU readiness. +pub struct PackedDecode<'a> { + // Raw carriers admit null as a bit pattern, so Option has an ordinary + // discriminant supported by the device enum ABI. Only the checked borrowed + // constructor creates this private payload; the phantom retains both loans. + input: *const BFloat16, + weight: *const BFloat16, + rows: usize, + input_stride: usize, + weight_stride: usize, + borrow: PhantomData<&'a [BFloat16]>, +} +impl<'a> PackedDecode<'a> { + /// Validate the storage needed by packed decode operand loads. + #[must_use] + #[inline(always)] + pub fn new( + input: &'a [BFloat16], + weight: &'a [BFloat16], + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) -> Option { + if Self::supported( + input, + weight, + rows, + width, + outputs, + input_stride, + weight_stride, + ) { + Some(Self { + input: input.as_ptr(), + weight: weight.as_ptr(), + rows, + input_stride, + weight_stride, + borrow: PhantomData, + }) + } else { + None + } + } + + /// Borrow a validated packed view inside the successful validation branch. + /// + /// Returns whether validation succeeded. The callback runs exactly once on + /// success and never on failure; its borrowed proof cannot leave the call. + /// Construction and use stay in the same branch, so device code need not + /// merge an optional aggregate containing the operand descriptors. + #[inline(always)] + #[allow(clippy::too_many_arguments)] + pub fn with_validated( + input: &'a [BFloat16], + weight: &'a [BFloat16], + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + use_proof: impl FnOnce(&Self), + ) -> bool { + if Self::supported( + input, + weight, + rows, + width, + outputs, + input_stride, + weight_stride, + ) { + let source = Self { + input: input.as_ptr(), + weight: weight.as_ptr(), + rows, + input_stride, + weight_stride, + borrow: PhantomData, + }; + use_proof(&source); + true + } else { + false + } + } + + #[inline(always)] + fn supported( + input: &[BFloat16], + weight: &[BFloat16], + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) -> bool { + matches!(rows, 1 | 2) + && width != 0 + && width.is_multiple_of(16) + && outputs != 0 + && outputs.is_multiple_of(8) + && input_stride >= width + && weight_stride >= width + && input_stride.is_multiple_of(2) + && weight_stride.is_multiple_of(2) + && (input.as_ptr() as usize).is_multiple_of(4) + && (weight.as_ptr() as usize).is_multiple_of(4) + && rows <= input.len() / input_stride + && outputs <= weight.len() / weight_stride + } +} + +#[inline(always)] +fn coordinate(lane: usize, value: usize) -> (usize, usize) { + let offset = L::crd2idx(lane + 32 * value); + (offset / COLS, offset % COLS) +} +impl APartition { + /// Load this lane's eight BF16 elements, preserving low/high pair order. + /// The callback receives local `(M,K)` coordinates in register order. + #[inline(always)] + pub fn load_with(self, mut read: impl FnMut(usize, usize) -> BFloat16) -> FragmentA { + let mut pair = |i| { + let (r, c) = coordinate::(self.lane, i); + let low = read(r, c).to_bits(); + let (r, c) = coordinate::(self.lane, i + 1); + u32::from(low) | (u32::from(read(r, c).to_bits()) << 16) + }; + FragmentA(ctm_rs::mma::Bf16A::from_registers([ + pair(0), + pair(2), + pair(4), + pair(6), + ])) + } + + /// Load a complete decode K fragment, repeating valid row addresses and + /// zero-masking absent logical rows without divergent loads. + /// + /// # Safety + /// `k` must be a multiple of 16 selecting a complete K tile within the + /// width passed to `PackedDecode::new` or `PackedDecode::with_validated`. + /// Operand data must be ready to read. + #[inline(always)] + pub unsafe fn load_decode_packed(self, source: &PackedDecode<'_>, k: usize) -> FragmentA { + let (row, low_k) = coordinate::(self.lane, 0); + let (_, high_k) = coordinate::(self.lane, 4); + let physical_row = row & (source.rows - 1); + let mask = 0u32.wrapping_sub((row < source.rows) as u32); + let base = source.input; + let low = unsafe { + ctm_rs::memory::load_aligned_u32( + base.add(physical_row * source.input_stride + k + low_k) + .cast(), + ) + }; + let high = unsafe { + ctm_rs::memory::load_aligned_u32( + base.add(physical_row * source.input_stride + k + high_k) + .cast(), + ) + }; + FragmentA(ctm_rs::mma::Bf16A::from_registers([ + low & mask, + 0, + high & mask, + 0, + ])) + } +} +impl BPartition { + /// Load this lane's four BF16 elements from local stored `(N,K)` coordinates. + #[inline(always)] + pub fn load_with(self, mut read: impl FnMut(usize, usize) -> BFloat16) -> FragmentB { + let mut pair = |i| { + let (r, c) = coordinate::(self.lane, i); + let low = read(r, c).to_bits(); + let (r, c) = coordinate::(self.lane, i + 1); + u32::from(low) | (u32::from(read(r, c).to_bits()) << 16) + }; + FragmentB(ctm_rs::mma::Bf16B::from_registers([pair(0), pair(2)])) + } + + /// Load one complete eight-output decode weight tile with packed reads. + /// + /// # Safety + /// `n` and `k` must select complete N=8 and K=16 tiles, aligned to those + /// dimensions and within the extents passed to `PackedDecode::new` or + /// `PackedDecode::with_validated`. + /// Operand data must be ready to read. + #[inline(always)] + pub unsafe fn load_decode_packed( + self, + source: &PackedDecode<'_>, + n: usize, + k: usize, + ) -> FragmentB { + let (row, low_k) = coordinate::(self.lane, 0); + let (_, high_k) = coordinate::(self.lane, 2); + let base = source.weight; + let low = unsafe { + ctm_rs::memory::load_aligned_u32( + base.add((n + row) * source.weight_stride + k + low_k) + .cast(), + ) + }; + let high = unsafe { + ctm_rs::memory::load_aligned_u32( + base.add((n + row) * source.weight_stride + k + high_k) + .cast(), + ) + }; + FragmentB(ctm_rs::mma::Bf16B::from_registers([low, high])) + } +} +impl CPartition { + /// Visit each accumulator once in register order with its local `(M,N)`. + /// The callback applies the explicit epilogue and output predication. + #[inline(always)] + pub fn store_with(self, accumulator: FragmentC, mut write: impl FnMut(usize, usize, f32)) { + let values = accumulator.0.registers(); + let mut i = 0; + while i < 4 { + let (r, c) = coordinate::(self.lane, i); + write(r, c, values[i]); + i += 1; + } + } + + /// Visit only initialized decode result rows, in the original register order. + #[inline(always)] + pub fn store_decode_with( + self, + accumulator: FragmentC, + source: &PackedDecode<'_>, + mut write: impl FnMut(usize, usize, f32), + ) { + let (row, _) = coordinate::(self.lane, 0); + if row < source.rows { + let values = accumulator.0.registers(); + let mut i = 0; + while i < 2 { + let (r, c) = coordinate::(self.lane, i); + write(r, c, values[i]); + i += 1; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Swapping pair halves or the upper A row band changes these PTX fixtures. + #[test] + fn operand_storage_preserves_bf16_bits_and_register_order() { + let a = Bf16Mma16x8x16::get_slice(0) + .partition_a() + .load_with(|r, k| BFloat16::from_bits(0x8000 + (r * 16 + k) as u16)); + assert_eq!( + a.0.registers(), + [0x8001_8000, 0x8081_8080, 0x8009_8008, 0x8089_8088] + ); + let b = Bf16Mma16x8x16::get_slice(31) + .partition_b() + .load_with(|n, k| BFloat16::from_bits(0x8000 + (n * 16 + k) as u16)); + assert_eq!(b.0.registers(), [0x8077_8076, 0x807f_807e]); + } + + /// Result mapping must preserve every FP32 payload until the explicit epilogue. + #[test] + fn result_partition_preserves_values_including_special_bits() { + let bits = [0x8000_0000, 0x7fc1_2345, 0x0000_0001, 0xff80_0000]; + let acc = FragmentC(ctm_rs::mma::Accumulator16x8::from_registers( + bits.map(f32::from_bits), + )); + let mut result = [(0, 0, 0); 4]; + let mut index = 0; + Bf16Mma16x8x16::get_slice(31) + .partition_c() + .store_with(acc, |r, c, value| { + result[index] = (r, c, value.to_bits()); + index += 1; + }); + assert_eq!( + result, + [ + (7, 6, bits[0]), + (7, 7, bits[1]), + (15, 6, bits[2]), + (15, 7, bits[3]) + ] + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/__compiler.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/__compiler.rs new file mode 100644 index 0000000000..8f4b520bcd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/__compiler.rs @@ -0,0 +1,214 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Exact compiler boundaries for SM120 block-scaled MMA. + +use super::tiled::a_atom_128; +use super::*; +use crate::block_scaled::{Mkl, Nkl, ScalePack4, SharedScaleAtom, Sm120ScaleAtom}; +use crate::copy::tiled::SharedTensor; +use crate::epilogue::sm120::Sm120EpilogueWarp128x128; +use crate::mma::{FragmentA, FragmentC}; +use crate::numeric::Mxf4E2M1; +use core::marker::PhantomData; + +/// Compiler boundary assigning a lane to its typed shared-layout MMA view. +#[inline(never)] +pub const fn tiled_mma_slice(lane: u32) -> Mxf4TiledMma { + Mxf4TiledMma { + lane, + layout: PhantomData, + } +} + +/// Compiler boundary filling all 16 per-lane accumulator fragments with one FP32 value. +#[inline(never)] +pub const fn fragment_fill(fill: f32) -> Mxf4AccumulatorTile2x8 { + let cell = FragmentC([fill; 4]); + Mxf4AccumulatorTile2x8 { + c00: cell, + c01: cell, + c02: cell, + c03: cell, + c04: cell, + c05: cell, + c06: cell, + c07: cell, + c10: cell, + c11: cell, + c12: cell, + c13: cell, + c14: cell, + c15: cell, + c16: cell, + c17: cell, + } +} + +/// Compiler boundary storing both row bands and all eight N fragments. +/// +/// # Safety +/// The collective participation, shared storage and publication contracts of the epilogue slice apply. +#[inline(never)] +pub unsafe fn epilogue_store_fragment( + slice: Sm120EpilogueWarp128x128, + tile: Mxf4AccumulatorTile2x8, +) { + unsafe { + slice.store_atom::<0, 0>(tile.c00); + slice.store_atom::<0, 1>(tile.c01); + slice.store_atom::<0, 2>(tile.c02); + slice.store_atom::<0, 3>(tile.c03); + slice.store_atom::<0, 4>(tile.c04); + slice.store_atom::<0, 5>(tile.c05); + slice.store_atom::<0, 6>(tile.c06); + slice.store_atom::<0, 7>(tile.c07); + slice.store_atom::<1, 0>(tile.c10); + slice.store_atom::<1, 1>(tile.c11); + slice.store_atom::<1, 2>(tile.c12); + slice.store_atom::<1, 3>(tile.c13); + slice.store_atom::<1, 4>(tile.c14); + slice.store_atom::<1, 5>(tile.c15); + slice.store_atom::<1, 6>(tile.c16); + slice.store_atom::<1, 7>(tile.c17); + } +} + +/// Compiler boundary selecting the scale pairs for one K64 half. +/// +/// # Safety +/// The half index must be zero or one. +#[inline(never)] +pub const unsafe fn fragment_slice_k(scales: Mxf4ScaleTile128, half: usize) -> Mxf4ScalePairs128 { + Mxf4ScalePairs128 { + a0: unsafe { scales.a0.pair_at_unchecked(half) }, + a1: unsafe { scales.a1.pair_at_unchecked(half) }, + b0: unsafe { scales.b0.pair_at_unchecked(half) }, + b1: unsafe { scales.b1.pair_at_unchecked(half) }, + b2: unsafe { scales.b2.pair_at_unchecked(half) }, + b3: unsafe { scales.b3.pair_at_unchecked(half) }, + b4: unsafe { scales.b4.pair_at_unchecked(half) }, + b5: unsafe { scales.b5.pair_at_unchecked(half) }, + b6: unsafe { scales.b6.pair_at_unchecked(half) }, + b7: unsafe { scales.b7.pair_at_unchecked(half) }, + } +} + +/// Compiler boundary loading this lane’s two A and eight B scale packs. +/// +/// # Safety +/// Shared scale atoms must be initialized and live; lane and warp coordinates must match the consuming MMA. +#[inline(never)] +pub unsafe fn mma_load_scales( + mma: &Mxf4TiledMma, + a: &SharedScaleAtom, + b: &SharedScaleAtom, + warp_m: usize, + warp_n: usize, +) -> Mxf4ScaleTile128 { + let q = (mma.lane >> 2) as usize; + let r = (mma.lane & 3) as usize; + let a_provider = q + 8 * (r & 1); + let a_row0 = 16 * warp_m + a_provider; + let a_row1 = 64 + 16 * warp_m + a_provider; + let b_row0 = 16 * warp_n + q; + let a_base = a.carrier().base.cast_const(); + let b_base = b.carrier().base.cast_const(); + + unsafe { + Mxf4ScaleTile128 { + a0: ScalePack4::from_bits(*a_base.add(Sm120ScaleAtom::word_offset(a_row0))), + a1: ScalePack4::from_bits(*a_base.add(Sm120ScaleAtom::word_offset(a_row1))), + b0: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0))), + b1: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 8))), + b2: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 32))), + b3: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 40))), + b4: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 64))), + b5: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 72))), + b6: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 96))), + b7: ScalePack4::from_bits(*b_base.add(Sm120ScaleAtom::word_offset(b_row0 + 104))), + } + } +} + +/// Compiler boundary loading the warp’s two A fragments for one K64 half. +/// +/// # Safety +/// Shared storage, lane participation and coordinates must satisfy Mxf4TiledMma::load_a. +#[inline(never)] +pub unsafe fn mma_load_a( + mma: &Mxf4TiledMma, + source: &SharedTensor, + warp_m: usize, + k_half: usize, +) -> (FragmentA, FragmentA) { + unsafe { + ( + mma.load_a(source, (a_atom_128(warp_m, 0), k_half)), + mma.load_a(source, (a_atom_128(warp_m, 1), k_half)), + ) + } +} + +/// Compiler boundary borrowing a lazy B view without loading its eight fragments. +#[inline(never)] +pub const fn mma_partition_b<'a, SharedLayout>( + _mma: &Mxf4TiledMma, + source: &'a SharedTensor, + warp_n: usize, + k_half: usize, +) -> Mxf4BTileK64<'a, SharedLayout> { + Mxf4BTileK64 { + base: source.carrier().base, + capacity: source.carrier().capacity, + warp_n, + k_half, + source: PhantomData, + } +} + +/// Compiler boundary issuing the four B-pair steps of a two-band, eight-column fragment tile. +/// +/// # Safety +/// All lanes must supply matching live A/B fragments and scales under the public tiled-MMA collective contract. +#[inline(never)] +pub unsafe fn tiled_gemm( + accumulators: &mut Mxf4AccumulatorTile2x8, + mma: &Mxf4TiledMma, + a: (FragmentA, FragmentA), + b: Mxf4BTileK64<'_, SharedLayout>, + scales: Mxf4ScalePairs128, +) { + let (a0, a1) = a; + + { + let (b0, b1) = unsafe { b.load::<0>(mma) }; + accumulators.c00 = unsafe { mma.gemm(a0, scales.a0, b0, scales.b0, accumulators.c00) }; + accumulators.c01 = unsafe { mma.gemm(a0, scales.a0, b1, scales.b1, accumulators.c01) }; + accumulators.c10 = unsafe { mma.gemm(a1, scales.a1, b0, scales.b0, accumulators.c10) }; + accumulators.c11 = unsafe { mma.gemm(a1, scales.a1, b1, scales.b1, accumulators.c11) }; + } + { + let (b2, b3) = unsafe { b.load::<1>(mma) }; + accumulators.c02 = unsafe { mma.gemm(a0, scales.a0, b2, scales.b2, accumulators.c02) }; + accumulators.c03 = unsafe { mma.gemm(a0, scales.a0, b3, scales.b3, accumulators.c03) }; + accumulators.c12 = unsafe { mma.gemm(a1, scales.a1, b2, scales.b2, accumulators.c12) }; + accumulators.c13 = unsafe { mma.gemm(a1, scales.a1, b3, scales.b3, accumulators.c13) }; + } + { + let (b4, b5) = unsafe { b.load::<2>(mma) }; + accumulators.c04 = unsafe { mma.gemm(a0, scales.a0, b4, scales.b4, accumulators.c04) }; + accumulators.c05 = unsafe { mma.gemm(a0, scales.a0, b5, scales.b5, accumulators.c05) }; + accumulators.c14 = unsafe { mma.gemm(a1, scales.a1, b4, scales.b4, accumulators.c14) }; + accumulators.c15 = unsafe { mma.gemm(a1, scales.a1, b5, scales.b5, accumulators.c15) }; + } + { + let (b6, b7) = unsafe { b.load::<3>(mma) }; + accumulators.c06 = unsafe { mma.gemm(a0, scales.a0, b6, scales.b6, accumulators.c06) }; + accumulators.c07 = unsafe { mma.gemm(a0, scales.a0, b7, scales.b7, accumulators.c07) }; + accumulators.c16 = unsafe { mma.gemm(a1, scales.a1, b6, scales.b6, accumulators.c16) }; + accumulators.c17 = unsafe { mma.gemm(a1, scales.a1, b7, scales.b7, accumulators.c17) }; + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/epilogue.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/epilogue.rs new file mode 100644 index 0000000000..ffb51d4895 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/epilogue.rs @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Store complete SM120 accumulator fragments through the shared epilogue. + +use super::{__compiler, Mxf4AccumulatorTile2x8}; +use crate::epilogue::sm120::Sm120EpilogueWarp128x128; + +impl Sm120EpilogueWarp128x128 { + /// Convert and store one warp's complete 2 x 8 result tile. + /// + /// ```text + /// 2 x 8 FP32 cells -> FP16 -> shared memory -> TMA output + /// ``` + /// + /// Every cell has a fixed position. No runtime array index is used. + /// + /// # Safety + /// + /// All 32 lanes must call this together on the slice for their current + /// warp. The caller must follow the synchronization rules documented by + /// [`Sm120EpilogueWarp128x128::store_atom`]. + #[inline(always)] + pub unsafe fn store_tile(self, tile: Mxf4AccumulatorTile2x8) { + unsafe { __compiler::epilogue_store_fragment(self, tile) }; + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/fragments.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/fragments.rs new file mode 100644 index 0000000000..1f7e0b73e1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/fragments.rs @@ -0,0 +1,311 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed scale, accumulator, and lazy B fragments for SM120 MMA. + +use super::{__compiler, Mxf4TiledMma}; +use crate::block_scaled::{Mkl, MmaScalePair, Nkl, ScalePack4}; +use crate::copy::tiled::SharedTensor; +use crate::mma::{FragmentA, FragmentB, FragmentC}; +use crate::numeric::Mxf4E2M1; +use core::marker::PhantomData; + +/// Scales for the five inputs used by one warp in one K=64 step. +/// +/// ```text +/// one A fragment + four B fragments +/// 1 + 4 = 5 scale pairs +/// ``` +/// +/// [`Mkl`] marks the A scale. [`Nkl`] marks the B scales, so Rust rejects an +/// A/B scale swap. +#[derive(Clone, Copy)] +pub struct Mxf4ScalePairs { + /// Scale pair for A. + pub a: MmaScalePair, + /// Scale pair for B columns 0 through 7. + pub b0: MmaScalePair, + /// Scale pair for B columns 8 through 15. + pub b1: MmaScalePair, + /// Scale pair for B columns 16 through 23. + pub b2: MmaScalePair, + /// Scale pair for B columns 24 through 31. + pub b3: MmaScalePair, +} + +/// Five packed scale words for one warp and one K=128 stage. +/// +/// Each word contains four 8-bit scales. The lower two scales cover the first +/// K=64 step; the upper two cover the second K=64 step. +#[derive(Clone, Copy)] +pub struct Mxf4ScaleStage { + /// Four K32 scales for the warp’s A fragment. + pub(super) a: ScalePack4, + /// Four K32 scales for B columns 0 through 7. + pub(super) b0: ScalePack4, + /// Four K32 scales for B columns 8 through 15. + pub(super) b1: ScalePack4, + /// Four K32 scales for B columns 16 through 23. + pub(super) b2: ScalePack4, + /// Four K32 scales for B columns 24 through 31. + pub(super) b3: ScalePack4, +} + +/// Packed scales for one warp's part of a 128 x 128 output tile. +/// +/// One warp computes two M positions and eight N positions: +/// +/// ```text +/// A scales: [a0] [a1] +/// B scales: [b0] [b1] [b2] [b3] [b4] [b5] [b6] [b7] +/// ``` +/// +/// Each word contains the four scales needed by the two K=64 steps in one +/// K=128 stage. +#[derive(Clone, Copy)] +pub struct Mxf4ScaleTile128 { + /// Four K32 scales for the first 16-row A band. + pub(super) a0: ScalePack4, + /// Four K32 scales for the second A band, 64 rows after the first. + pub(super) a1: ScalePack4, + /// Four K32 scales for B’s first N8 fragment. + pub(super) b0: ScalePack4, + /// Four K32 scales for B’s second N8 fragment. + pub(super) b1: ScalePack4, + /// Four K32 scales for B’s third N8 fragment. + pub(super) b2: ScalePack4, + /// Four K32 scales for B’s fourth N8 fragment. + pub(super) b3: ScalePack4, + /// Four K32 scales for B’s fifth N8 fragment. + pub(super) b4: ScalePack4, + /// Four K32 scales for B’s sixth N8 fragment. + pub(super) b5: ScalePack4, + /// Four K32 scales for B’s seventh N8 fragment. + pub(super) b6: ScalePack4, + /// Four K32 scales for B’s eighth N8 fragment. + pub(super) b7: ScalePack4, +} + +/// Scale pairs for one warp's 16 MMA cells in one K=64 step. +/// +/// The two A pairs combine with all eight B pairs: +/// +/// ```text +/// b0 b1 b2 b3 b4 b5 b6 b7 +/// a0 x x x x x x x x +/// a1 x x x x x x x x +/// ``` +#[derive(Clone, Copy)] +pub struct Mxf4ScalePairs128 { + /// A scale for rows starting at `16 * warp_m`. + pub a0: MmaScalePair, + /// A scale for rows starting at `64 + 16 * warp_m`. + pub a1: MmaScalePair, + /// B scale for columns starting at `16 * warp_n`. + pub b0: MmaScalePair, + /// B scale for columns starting at `16 * warp_n + 8`. + pub b1: MmaScalePair, + /// B scale for columns starting at `32 + 16 * warp_n`. + pub b2: MmaScalePair, + /// B scale for columns starting at `40 + 16 * warp_n`. + pub b3: MmaScalePair, + /// B scale for columns starting at `64 + 16 * warp_n`. + pub b4: MmaScalePair, + /// B scale for columns starting at `72 + 16 * warp_n`. + pub b5: MmaScalePair, + /// B scale for columns starting at `96 + 16 * warp_n`. + pub b6: MmaScalePair, + /// B scale for columns starting at `104 + 16 * warp_n`. + pub b7: MmaScalePair, +} + +/// One warp's FP32 results for a 128 x 128 output tile. +/// +/// One MMA cell is a 16 x 8 result. Each lane holds four FP32 values from +/// every cell. +/// +/// One warp owns this 2 x 8 grid of cells: +/// +/// ```text +/// eight N positions +/// 0 1 2 3 4 5 6 7 +/// first M band [ ] [ ] [ ] [ ] [ ] [ ] [ ] [ ] +/// second M band [ ] [ ] [ ] [ ] [ ] [ ] [ ] [ ] +/// ``` +/// +/// The 16 cells have fixed names inside the type. The compiler therefore +/// keeps them in registers without a runtime array index. +#[derive(Clone, Copy)] +pub struct Mxf4AccumulatorTile2x8 { + /// Per-lane accumulators for row band 0 and N8 fragment 0. + pub(super) c00: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 1. + pub(super) c01: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 2. + pub(super) c02: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 3. + pub(super) c03: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 4. + pub(super) c04: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 5. + pub(super) c05: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 6. + pub(super) c06: FragmentC, + /// Per-lane accumulators for row band 0 and N8 fragment 7. + pub(super) c07: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 0. + pub(super) c10: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 1. + pub(super) c11: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 2. + pub(super) c12: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 3. + pub(super) c13: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 4. + pub(super) c14: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 5. + pub(super) c15: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 6. + pub(super) c16: FragmentC, + /// Per-lane accumulators for row band 1 and N8 fragment 7. + pub(super) c17: FragmentC, +} + +/// One warp's eight B fragments for one K=64 step. +/// +/// This is a view, so creating it does not load all eight fragments. During +/// [`Mxf4AccumulatorTile2x8::gemm`], it loads two at a time: +/// +/// ```text +/// load b0,b1 -> compute -> load b2,b3 -> compute -> ... +/// ``` +/// +/// That keeps fewer temporary values in registers. +pub struct Mxf4BTileK64<'a, SharedLayout> { + /// Shared packed-value base, addressed in f16 carrier elements. + pub(super) base: *mut f16, + /// Available f16 carrier elements from the base pointer. + pub(super) capacity: usize, + /// Warp’s N partition, selecting its group of eight B fragments. + pub(super) warp_n: usize, + /// K64 half of the staged K128 input: zero or one. + pub(super) k_half: usize, + /// Borrows the original shared tensor and preserves its layout and B-operand role. + pub(super) source: PhantomData<&'a SharedTensor>, +} + +impl Mxf4BTileK64<'_, SharedLayout> { + /// Load B fragments `2 * PAIR` and `2 * PAIR + 1`. + /// + /// # Safety + /// + /// All 32 lanes must call this together. `mma` must contain each calling + /// thread's real lane number. This view's `warp_n` and `k_half` must each + /// be 0 or 1, and its shared tile must still contain the completed stage. + #[inline(always)] + pub(super) unsafe fn load( + &self, + mma: &Mxf4TiledMma, + ) -> (FragmentB, FragmentB) { + const { assert!(PAIR < 4) }; + let source = unsafe { + SharedTensor::::from_raw_parts( + self.base, + self.capacity, + ) + }; + unsafe { mma.load_b_pair_128(&source, self.warp_n, PAIR, self.k_half) } + } +} + +impl Mxf4AccumulatorTile2x8 { + /// Start a new output tile with every FP32 value set to zero. + #[must_use] + #[inline(always)] + pub const fn zero() -> Self { + __compiler::fragment_fill(0.0) + } + + /// Add one K=64 input step to all 16 result cells. + /// + /// ```text + /// two A fragments x eight B fragments + /// | + /// v + /// this 2 x 8 FP32 result tile + /// ``` + /// + /// `a` contains the two A fragments already loaded for this warp. `b` + /// selects the eight matching B fragments and loads them two at a time. + /// `scales` must cover the same A, B, and K=64 positions. + /// + /// [`Mxf4TiledMma::gemm`] is the one-cell form of this operation; this + /// method runs sixteen of those cells for the warp's 2 x 8 tile. + /// + /// # Safety + /// + /// All 32 lanes must call this together. `mma` must contain each calling + /// thread's real lane number. The A fragments, B view, and scales must use + /// the same warp position and K=64 half from one completed shared-memory + /// stage. The B view's `warp_n` and `k_half` must each be 0 or 1, and that + /// stage must not be overwritten during this call. + #[inline(always)] + pub unsafe fn gemm( + &mut self, + mma: &Mxf4TiledMma, + a: (FragmentA, FragmentA), + b: Mxf4BTileK64<'_, SharedLayout>, + scales: Mxf4ScalePairs128, + ) { + unsafe { __compiler::tiled_gemm(self, mma, a, b, scales) }; + } +} + +impl Mxf4ScaleTile128 { + /// Take the scale pairs for one K=64 half of a K=128 stage. + /// + /// # Safety + /// + /// `half` must be `0` for the first half or `1` for the second. + #[must_use] + #[inline(always)] + pub const unsafe fn pairs_at_unchecked(self, half: usize) -> Mxf4ScalePairs128 { + unsafe { __compiler::fragment_slice_k(self, half) } + } +} + +impl Mxf4ScaleStage { + /// Take the scale pairs for the first or second K=64 half. + /// + /// `HALF=0` selects the first half. `HALF=1` selects the second. + #[must_use] + #[inline(always)] + pub const fn pairs(self) -> Mxf4ScalePairs { + Mxf4ScalePairs { + a: self.a.pair::(), + b0: self.b0.pair::(), + b1: self.b1.pair::(), + b2: self.b2.pair::(), + b3: self.b3.pair::(), + } + } + + /// Take one K=64 half using a runtime value. + /// + /// # Safety + /// + /// `half` must be `0` for the first half or `1` for the second. + #[must_use] + #[inline(always)] + pub const unsafe fn pairs_at_unchecked(self, half: usize) -> Mxf4ScalePairs { + Mxf4ScalePairs { + a: unsafe { self.a.pair_at_unchecked(half) }, + b0: unsafe { self.b0.pair_at_unchecked(half) }, + b1: unsafe { self.b1.pair_at_unchecked(half) }, + b2: unsafe { self.b2.pair_at_unchecked(half) }, + b3: unsafe { self.b3.pair_at_unchecked(half) }, + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/mod.rs new file mode 100644 index 0000000000..1a9202ec89 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/mod.rs @@ -0,0 +1,59 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Block-scaled MXFP4 matrix multiply for SM120 GPUs. +//! +//! This module covers the compute part of a 128 x 128 output tile: +//! +//! ```text +//! packed FP4 A + packed FP4 B + 8-bit scales +//! | +//! v +//! shared memory +//! | +//! v +//! fragments held by each GPU lane +//! | +//! v +//! m16n8k64 tensor-core instructions +//! | +//! v +//! 2 x 8 FP32 result cells +//! ``` +//! +//! Words used below: +//! +//! - **M** and **N** are the output row and column dimensions. **K** is the +//! shared dimension whose products are summed. +//! - A **warp** is 32 GPU threads running one instruction together. +//! - A **lane** is one thread inside that warp. +//! - A **fragment** is the small part of a matrix held by one lane. +//! - **MMA** means matrix multiply-accumulate: `result = A * B + result`. +//! - An **MMA cell** is one 16 x 8 result made by a tensor-core instruction. +//! - A **scale pair** contains two scales: one for each K=32 half of K=64. +//! - A **stage** is one K=128 input block held in shared memory. +//! - **TMA** is the GPU copy engine used to move whole tiles. +//! - **MXFP4** stores 4-bit E2M1 numbers with an 8-bit UE8M0 scale for each +//! group of 32 numbers. +//! +//! These Rust types add checks and names. CuTe-aware compilers retain the +//! shared-memory MMA story and choose its implementation for their backend. + +mod epilogue; +mod fragments; +mod tiled; + +pub use fragments::{ + Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxf4ScalePairs, Mxf4ScalePairs128, Mxf4ScaleStage, + Mxf4ScaleTile128, +}; +pub use tiled::Mxf4TiledMma; + +/// Stable compiler boundaries for the SM120 shared-memory MMA path. +#[doc(hidden)] +pub mod __compiler; +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tests.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tests.rs new file mode 100644 index 0000000000..8bb97ef3f2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tests.rs @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks SM120 accumulator, B-fragment and scale-provider mappings."] +use super::tiled::{a_atom_128, b_pair_atom_128}; +use super::{Mxf4AccumulatorTile2x8, Mxf4BTileK64, Mxf4TiledMma}; +use crate::block_scaled::{Mkl, Nkl, SharedScaleAtom, Sm120ScaleAtom}; +use crate::mma::FragmentC; + +/// Accumulator and lazy B views must retain the intended register and pointer carrier sizes. +#[test] +fn accumulator_and_lazy_b_keep_their_physical_carriers_explicit() { + assert_eq!( + core::mem::size_of::(), + 16 * core::mem::size_of::() + ); + assert_eq!( + core::mem::size_of::>(), + // shared base + capacity + warp-N + K-half + 4 * core::mem::size_of::() + ); +} + +/// Warp row bands and B-pair indices must match CuTe’s 128-by-128 fragment assignment. +#[test] +fn cta_128_fragment_permutation_matches_cute() { + assert_eq!( + [ + [16 * a_atom_128(0, 0), 16 * a_atom_128(0, 1)], + [16 * a_atom_128(1, 0), 16 * a_atom_128(1, 1)], + [16 * a_atom_128(2, 0), 16 * a_atom_128(2, 1)], + [16 * a_atom_128(3, 0), 16 * a_atom_128(3, 1)], + ], + [[0, 64], [16, 80], [32, 96], [48, 112]] + ); + assert_eq!( + [ + [ + 16 * b_pair_atom_128(0, 0), + 16 * b_pair_atom_128(0, 1), + 16 * b_pair_atom_128(0, 2), + 16 * b_pair_atom_128(0, 3), + ], + [ + 16 * b_pair_atom_128(1, 0), + 16 * b_pair_atom_128(1, 1), + 16 * b_pair_atom_128(1, 2), + 16 * b_pair_atom_128(1, 3), + ], + ], + [[0, 32, 64, 96], [16, 48, 80, 112]] + ); +} + +/// Each lane’s scale load must select exactly the canonical two A and eight B words. +#[test] +fn canonical_scale_atom_loads_exact_two_a_and_eight_b_packs() { + let mut a_words = + core::array::from_fn::<_, { Sm120ScaleAtom::WORDS }, _>(|word| 0xa500_0000 | word as u32); + let mut b_words = + core::array::from_fn::<_, { Sm120ScaleAtom::WORDS }, _>(|word| 0xb500_0000 | word as u32); + // Both arrays are aligned, contain one complete scale block, and stay + // alive while the views are used. + let a = unsafe { + SharedScaleAtom::::from_raw_parts(a_words.as_mut_ptr(), Sm120ScaleAtom::WORDS) + }; + // The separate B array follows the same rules. + let b = unsafe { + SharedScaleAtom::::from_raw_parts(b_words.as_mut_ptr(), Sm120ScaleAtom::WORDS) + }; + + // Lane 0 provides A rows 0/64 and B rows + // 0,8,32,40,64,72,96,104 for warp (0,0). + let lane0 = Mxf4TiledMma::<()>::get_slice(0); + // The arrays are valid and lane/warp positions are in range. + let first = unsafe { lane0.load_scale_atom_128(&a, &b, 0, 0) }; + assert_eq!( + [first.a0.bits(), first.a1.bits()], + [0xa500_0000, 0xa500_0002] + ); + assert_eq!( + [ + first.b0.bits(), + first.b1.bits(), + first.b2.bits(), + first.b3.bits(), + first.b4.bits(), + first.b5.bits(), + first.b6.bits(), + first.b7.bits(), + ], + [ + 0xb500_0000, + 0xb500_0020, + 0xb500_0001, + 0xb500_0021, + 0xb500_0002, + 0xb500_0022, + 0xb500_0003, + 0xb500_0023, + ] + ); + + // Lane 31 checks the last rows and words for warp position (3,1). + let lane31 = Mxf4TiledMma::<()>::get_slice(31); + // The arrays are valid and lane/warp positions are in range. + let last = unsafe { lane31.load_scale_atom_128(&a, &b, 3, 1) }; + assert_eq!([last.a0.bits(), last.a1.bits()], [0xa500_007d, 0xa500_007f]); + assert_eq!( + [ + last.b0.bits(), + last.b1.bits(), + last.b2.bits(), + last.b3.bits(), + last.b4.bits(), + last.b5.bits(), + last.b6.bits(), + last.b7.bits(), + ], + [ + 0xb500_005c, + 0xb500_007c, + 0xb500_005d, + 0xb500_007d, + 0xb500_005e, + 0xb500_007e, + 0xb500_005f, + 0xb500_007f, + ] + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tiled.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tiled.rs new file mode 100644 index 0000000000..78160588c3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/block_scaled/tiled.rs @@ -0,0 +1,338 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Warp MMA slices and shared-memory fragment loading. + +use super::{__compiler, Mxf4BTileK64, Mxf4ScaleStage, Mxf4ScaleTile128}; +use crate::block_scaled::{Mkl, MmaScalePair, Nkl, ScalePack4, SharedScaleAtom}; +use crate::copy::tiled::SharedTensor; +use crate::markers::ColMajor; +use crate::mma::{FragmentA, FragmentB, FragmentC, load_matrix_a}; +use crate::numeric::{Mxf4E2M1, UE8M0x4}; +use core::marker::PhantomData; + +// Map `(warp M position, first-or-second M block)` to a 16-row block number. +#[inline(always)] +/// Map a warp M partition and one of its two row bands to an A atom index. +pub(super) const fn a_atom_128(warp_m: usize, atom_m: usize) -> usize { + warp_m + 4 * atom_m +} + +// Map `(warp N position, B pair number)` to a 16-row block number in N x K B. +#[inline(always)] +/// Map a warp N partition and B-pair slot to the first N8 atom index. +pub(super) const fn b_pair_atom_128(warp_n: usize, pair_n: usize) -> usize { + warp_n + 2 * pair_n +} + +/// One lane's handle to the SM120 MXFP4 tensor-core operation. +/// +/// All 32 lanes create this handle with their own lane number. They then use +/// it together to load fragments and run `m16n8k64`: +/// +/// ```text +/// 32 lane handles -> 32 lanes load together -> one warp-wide MMA +/// ``` +#[derive(Clone, Copy)] +pub struct Mxf4TiledMma { + /// Calling thread’s lane index in the 32-thread warp. + pub(super) lane: u32, + /// Retains the shared-memory coordinate map for fragment loads. + pub(super) layout: PhantomData, +} + +impl Mxf4TiledMma { + /// Create the handle for one lane of the warp. + /// + /// The later load methods require `lane` to be the calling thread's real + /// lane number, from 0 through 31. + #[must_use] + #[inline(always)] + pub const fn get_slice(lane: u32) -> Self { + __compiler::tiled_mma_slice(lane) + } + + /// Create a view of the eight B fragments for one K=64 step. + /// + /// ```text + /// shared B tile + warp N position + K half + /// | + /// v + /// eight-fragment view + /// ``` + /// + /// The view loads two fragments at a time when it is passed to + /// [`Mxf4AccumulatorTile2x8::gemm`](super::Mxf4AccumulatorTile2x8::gemm). `warp_n` must be 0 or 1; + /// `k_half` must be 0 or 1 before that unsafe operation is called. + #[must_use] + #[inline(always)] + pub const fn get_b_tile_k64<'a>( + &self, + source: &'a SharedTensor, + warp_n: usize, + k_half: usize, + ) -> Mxf4BTileK64<'a, SharedLayout> { + __compiler::mma_partition_b(self, source, warp_n, k_half) + } + + /// Load this lane's part of a 16 x 64 FP4 A block. + /// + /// Four FP4 values use one 16-bit storage value, so the 16 x 64 FP4 block + /// occupies the same bytes as a 16 x 16 `f16` block. + /// + /// # Safety + /// + /// All 32 lanes must call this together. This handle must contain the + /// calling thread's lane number. `warp_tile` must select an initialized, + /// correctly aligned block inside `source`. + #[must_use] + #[inline(always)] + pub unsafe fn load_a( + &self, + source: &SharedTensor, + warp_tile: (usize, usize), + ) -> FragmentA { + unsafe { load_matrix_a::(source.carrier(), warp_tile, self.lane) } + } + + /// Load two neighboring B fragments. + /// + /// B is stored as `N x K`, with K next to K in memory: + /// + /// ```text + /// one 16-row B load -> fragment for 8 columns + fragment for 8 columns + /// ``` + /// + /// This lets A and B use the same `ldmatrix.x4` load. + /// + /// # Safety + /// + /// The same rules as [`Self::load_a`] apply. + #[must_use] + #[inline(always)] + pub unsafe fn load_b_pair( + &self, + source: &SharedTensor, + warp_tile: (usize, usize), + ) -> (FragmentB, FragmentB) { + let loaded = + unsafe { load_matrix_a::(source.carrier(), warp_tile, self.lane) }; + ( + FragmentB([loaded.0[0], loaded.0[2]]), + FragmentB([loaded.0[1], loaded.0[3]]), + ) + } + + /// Load the two A fragments owned by this warp. + /// + /// The four M warp positions receive these row blocks: + /// + /// ```text + /// warp_m first block second block + /// 0 0 64 + /// 1 16 80 + /// 2 32 96 + /// 3 48 112 + /// ``` + /// + /// # Safety + /// + /// The rules from [`Self::load_a`] apply. `warp_m` must be 0 through 3, + /// `k_half` must be 0 or 1, and `source` must contain the complete staged + /// 128 x 128 FP4 tile. + #[must_use] + #[inline(always)] + pub unsafe fn load_a_128( + &self, + source: &SharedTensor, + warp_m: usize, + k_half: usize, + ) -> (FragmentA, FragmentA) { + unsafe { __compiler::mma_load_a(self, source, warp_m, k_half) } + } + + /// Load one of this warp's four B-fragment pairs. + /// + /// Each pair gives two neighboring 8-column fragments: + /// + /// ```text + /// pair_n 0 1 2 3 + /// base col 0 32 64 96 + /// plus 16 * warp_n + /// ``` + /// + /// # Safety + /// + /// The rules from [`Self::load_b_pair`] apply. `warp_n` must be 0 or 1, + /// `pair_n` must be 0 through 3, and `k_half` must be 0 or 1. `source` + /// must contain the complete staged 128 x 128 FP4 tile. + #[must_use] + #[inline(always)] + pub unsafe fn load_b_pair_128( + &self, + source: &SharedTensor, + warp_n: usize, + pair_n: usize, + k_half: usize, + ) -> (FragmentB, FragmentB) { + unsafe { self.load_b_pair(source, (b_pair_atom_128(warp_n, pair_n), k_half)) } + } + + /// Load five packed scale words for a 64 x 32 warp tile. + /// + /// One word is for A. Four words are for B. Each word contains four + /// 8-bit scales, so both K=64 steps reuse these loads. + /// + /// # Safety + /// + /// `ROWS` must be positive. Both shared tensors must contain the requested + /// `word`. `warp_m` and `warp_n` must select blocks inside those tensors. + /// This handle must contain the calling thread's real lane number. + #[must_use] + #[inline(always)] + pub unsafe fn load_scale_stage( + &self, + a: &SharedTensor, Mkl>, + b: &SharedTensor, Nkl>, + word: usize, + warp_m: usize, + warp_n: usize, + ) -> Mxf4ScaleStage { + const { assert!(ROWS > 0) }; + let rows = ROWS as usize; + let q = (self.lane >> 2) as usize; + let r = (self.lane & 3) as usize; + let a_row = 16 * warp_m + q + 8 * (r & 1); + let b_row = 32 * warp_n + q; + let column = word * rows; + let a_base = a.carrier().base.cast_const(); + let b_base = b.carrier().base.cast_const(); + + unsafe { + Mxf4ScaleStage { + a: ScalePack4::from_bits(*a_base.add(column + a_row)), + b0: ScalePack4::from_bits(*b_base.add(column + b_row)), + b1: ScalePack4::from_bits(*b_base.add(column + b_row + 8)), + b2: ScalePack4::from_bits(*b_base.add(column + b_row + 16)), + b3: ScalePack4::from_bits(*b_base.add(column + b_row + 24)), + } + } + } + + /// Load the ten packed scale words for one warp's 128 x 128 tile part. + /// + /// The words match the same two-by-eight grid as the result cells: + /// + /// ```text + /// A row blocks: 16*warp_m, 64 + 16*warp_m + /// B pair starts: 16*warp_n + {0, 32, 64, 96} + /// ``` + /// + /// One word contains four K=32 scales, covering both K=64 steps. Every + /// lane reads a valid word, even when the hardware does not use that + /// lane's scale value. + /// + /// # Safety + /// + /// Both shared tensors must contain `word`. `ROWS` must be at least 128, + /// `word` must be 0 or 1, `warp_m` must be 0 through 3, and `warp_n` must + /// be 0 or 1. This handle must contain the calling thread's lane number. + #[must_use] + #[inline(always)] + pub unsafe fn load_scale_tile_128( + &self, + a: &SharedTensor, Mkl>, + b: &SharedTensor, Nkl>, + word: usize, + warp_m: usize, + warp_n: usize, + ) -> Mxf4ScaleTile128 { + const { assert!(ROWS >= 128) }; + let rows = ROWS as usize; + let q = (self.lane >> 2) as usize; + let r = (self.lane & 3) as usize; + let a_provider = q + 8 * (r & 1); + let a_row0 = 16 * warp_m + a_provider; + let a_row1 = 64 + 16 * warp_m + a_provider; + let b_row0 = 16 * warp_n + q; + let column = word * rows; + let a_base = a.carrier().base.cast_const(); + let b_base = b.carrier().base.cast_const(); + + unsafe { + Mxf4ScaleTile128 { + a0: ScalePack4::from_bits(*a_base.add(column + a_row0)), + a1: ScalePack4::from_bits(*a_base.add(column + a_row1)), + b0: ScalePack4::from_bits(*b_base.add(column + b_row0)), + b1: ScalePack4::from_bits(*b_base.add(column + b_row0 + 8)), + b2: ScalePack4::from_bits(*b_base.add(column + b_row0 + 32)), + b3: ScalePack4::from_bits(*b_base.add(column + b_row0 + 40)), + b4: ScalePack4::from_bits(*b_base.add(column + b_row0 + 64)), + b5: ScalePack4::from_bits(*b_base.add(column + b_row0 + 72)), + b6: ScalePack4::from_bits(*b_base.add(column + b_row0 + 96)), + b7: ScalePack4::from_bits(*b_base.add(column + b_row0 + 104)), + } + } + } + + /// Load the ten packed scale words from SM120's required scale layout. + /// + /// Each shared input is one 512-byte block: + /// + /// ```text + /// 512-byte A scale block -> 2 packed words + /// 512-byte B scale block -> 8 packed words + /// ``` + /// + /// TMA copies the bytes without rearranging them. [`Sm120ScaleAtom`](crate::block_scaled::Sm120ScaleAtom) + /// converts each logical row into the physical word position required by + /// the GPU. The ten words are reused by both K=64 steps. + /// + /// # Safety + /// + /// `a` and `b` must point to initialized, `u32`-aligned shared-memory + /// blocks of [`Sm120ScaleAtom::BYTES`](crate::block_scaled::Sm120ScaleAtom::BYTES) bytes. `warp_m` must be 0 through + /// 3, `warp_n` must be 0 or 1, and this handle must contain the calling + /// thread's real lane number. + #[must_use] + #[inline(always)] + pub unsafe fn load_scale_atom_128( + &self, + a: &SharedScaleAtom, + b: &SharedScaleAtom, + warp_m: usize, + warp_n: usize, + ) -> Mxf4ScaleTile128 { + unsafe { __compiler::mma_load_scales(self, a, b, warp_m, warp_n) } + } + + /// Run one 16 x 8 x 64 block-scaled tensor-core operation. + /// + /// ```text + /// result = previous result + (A * A scale) x (B * B scale) + /// ``` + /// + /// The A scale has type [`MmaScalePair`]. The B scale has type + /// [`MmaScalePair`], so Rust rejects swapping them. + /// + /// # Safety + /// + /// All 32 lanes must call this together. `a`, `b`, and `c` must use the + /// SM120 `m16n8k64` lane layout. Each scale pair must contain the first + /// K=32 scale followed by the second K=32 scale. + #[must_use] + #[inline(always)] + pub unsafe fn gemm( + &self, + a: FragmentA, + scale_a: MmaScalePair, + b: FragmentB, + scale_b: MmaScalePair, + c: FragmentC, + ) -> FragmentC { + let _ = (self, a, scale_a, b, scale_b, c); + unreachable!("cute-rs block-scaled MMA executed outside recognized CuTe device compilation") + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/mma/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/mod.rs new file mode 100644 index 0000000000..85d5b9dac3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/mma/mod.rs @@ -0,0 +1,212 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Warp-level tensor-core loads and multiply-accumulate operations. +//! +//! A thread is one GPU worker. A warp is 32 threads that run together. The +//! tensor core splits each small matrix across the registers of all 32 +//! threads. One thread's piece is called a *fragment*. A lane is that +//! thread's position from 0 to 31 inside the warp. +//! +//! One `mma.sync.m16n8k16` performs this operation: +//! +//! ```text +//! A: 16x16 f16 ──┐ +//! ├── one warp ──► C: 16x8 f32 +//! B: 16x8 f16 ──┘ C = A*B + C +//! +//! each thread holds: FragmentA [u32; 4], FragmentB [u32; 2], FragmentC [f32; 4] +//! ``` +//! +//! The compiler retains `load_matrix_*` as typed matrix-load semantics. The +//! selected backend chooses the lane mapping and target operation. The `SmemL` +//! type tells it the shared-memory layout, including any swizzle (a fixed +//! address rearrangement). A plain pointer would lose that information. +//! +//! [`bf16`] provides typed BF16 partitions and fragments; [`block_scaled`] +//! provides the SM120 MXFP4 tiled family. Their operand contracts remain +//! distinct from the common FP16 fragments defined here. +//! +//! [`FragmentC::mma_atom_call`] issues one `mma.sync` through ctm-rs +//! `mma_f16`; [`bf16::Bf16Mma16x8x16::mma_atom_call`] does the same through +//! `mma_bf16`. +//! +pub mod bf16; +pub mod block_scaled; + +use crate::copy::cooperative::SmemTile; + +/// One thread's part of a 16x16 `f16` A tile. +/// +/// Each of the four 32-bit registers holds two adjacent `f16` values in the +/// order expected by `mma.sync`. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct FragmentA( + #[doc = "This lane’s four packed A registers in instruction order."] pub [u32; 4], +); + +/// One thread's part of a 16x8 `f16` B tile: two 32-bit registers. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct FragmentB( + #[doc = "This lane’s two packed B registers in instruction order."] pub [u32; 2], +); + +/// Two adjacent 16x8 B fragments in MMA register order. +/// +/// Produced by [`load_matrix_b_pair`], [`load_matrix_b_pair_k_major`], or by +/// retiling an A-fragment load with [`FragmentA::as_b_pair`]. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct FragmentBPair( + #[doc = "Packed B registers for two neighboring N8 fragments."] pub [u32; 4], +); + +impl FragmentBPair { + /// Select the left (`N = 0`) or right (`N = 1`) fragment. + #[inline(always)] + pub fn get(self) -> FragmentB { + const { assert!(N < 2) }; + FragmentB([self.0[N * 2], self.0[N * 2 + 1]]) + } +} + +impl FragmentA { + /// Interpret a non-transposed load of two K-major B tiles. + /// + /// The shared window contains 16 B rows and 16 reduction elements. + #[inline(always)] + pub fn as_b_pair(self) -> FragmentBPair { + FragmentBPair([self.0[0], self.0[2], self.0[1], self.0[3]]) + } +} + +/// One thread's four `f32` values from a 16x8 result tile. +/// +/// `cuda_device::mma_frag` maps these four values to result coordinates. +#[repr(transparent)] +#[derive(Clone, Copy)] +pub struct FragmentC( + #[doc = "This lane’s four FP32 result registers in instruction order."] pub [f32; 4], +); + +impl FragmentC { + /// Four zero values for starting `C = A*B + C`. + pub const ZERO: FragmentC = FragmentC([0.0; 4]); + + /// Accumulate one FP16 16x16 by 16x8 warp matrix product in FP32. + /// + /// # Safety + /// All lanes must participate with matching fragments and supported hardware. + #[must_use = "use the returned accumulator; mma_atom_call does not modify self in place"] + #[inline(always)] + pub unsafe fn mma_atom_call(self, a: FragmentA, b: FragmentB) -> Self { + use ctm_rs::mma::{Accumulator16x8, F16A, F16B, mma_f16}; + + Self( + unsafe { + mma_f16( + Accumulator16x8::from_registers(self.0), + F16A::from_registers(a.0), + F16B::from_registers(b.0), + ) + } + .registers(), + ) + } +} + +/// Load one 16x16 `f16` A tile into the 32 threads of a warp. +/// +/// Each thread receives its [`FragmentA`]. `SmemL` describes the shared-memory +/// layout and any address swizzle. +/// +/// `warp_tile = (r, c)` counts 16x16 windows. For example, `(2, 3)` starts at +/// element `(32, 48)`. +/// +/// # Safety +/// +/// All 32 lanes of one warp must call this function on the same control-flow +/// path. `lane` must be the calling thread's hardware lane number. The shared +/// tile must contain the selected window. Every selected row address must be +/// aligned to 16 bytes under `SmemL`; a composed swizzle is valid when it +/// keeps at least 16-byte chunks together. The compiler checks the parts +/// known from the types. +#[inline(never)] +pub unsafe fn load_matrix_a( + src: &SmemTile, + warp_tile: (usize, usize), + lane: u32, +) -> FragmentA { + let _ = (src, warp_tile, lane); + unreachable!("cute-rs `load_matrix_a` executed outside device compilation") +} + +/// Load one 16x8 `f16` B tile into the 32 threads of a warp. +/// +/// Each thread receives its [`FragmentB`]. `SmemL` describes the shared-memory +/// layout and any address swizzle. +/// +/// `warp_tile = (r, c)` counts 16x8 windows. The window starts at element +/// `(16*r, 8*c)`. `ldmatrix.trans` transposes while loading, turning a +/// row-major tile into the column order expected by `mma.sync`. +/// +/// # Safety +/// +/// All 32 lanes of one warp must call this function on the same control-flow +/// path, and `lane` must be the calling thread's hardware lane number. The +/// shared tile must contain the selected window, with row addresses aligned +/// as required by [`load_matrix_a`]. +#[inline(never)] +pub unsafe fn load_matrix_b( + src: &SmemTile, + warp_tile: (usize, usize), + lane: u32, +) -> FragmentB { + let _ = (src, warp_tile, lane); + unreachable!("cute-rs `load_matrix_b` executed outside device compilation") +} + +/// Load a 16x16 shared-memory window as two adjacent B fragments. +/// +/// `warp_tile` counts 16x16 windows. One `ldmatrix.x4.trans` supplies both +/// 16x8 operands, retaining the source layout and swizzle in the compiler. +/// +/// # Safety +/// The lane participation, alignment, and storage requirements of +/// [`load_matrix_a`] apply to the entire 16x16 window. +#[inline(never)] +pub unsafe fn load_matrix_b_pair( + src: &SmemTile, + warp_tile: (usize, usize), + lane: u32, +) -> FragmentBPair { + let _ = (src, warp_tile, lane); + unreachable!("cute-rs `load_matrix_b_pair` executed outside device compilation") +} + +/// Load a K-major 16x16 shared-memory window as two adjacent B fragments. +/// +/// Shared rows are B columns (`N`), and shared columns are reduction elements +/// (`K`). `warp_tile = (n, k)` selects the window starting at `(16*n, 16*k)`. +/// One non-transposed `ldmatrix.x4` assigns its four 8x8 submatrices in B-fragment +/// order, so each pair of result registers can feed one `mma.sync` directly. +/// This produces the same fragments as [`load_matrix_a`] followed by +/// [`FragmentA::as_b_pair`], without requiring a register permutation after loading. +/// +/// # Safety +/// The lane participation, alignment, and storage requirements of +/// [`load_matrix_a`] apply to the entire 16x16 window. `lane` must be the calling +/// thread's hardware lane number; the compiler assigns the submatrix addresses. +#[inline(never)] +pub unsafe fn load_matrix_b_pair_k_major( + src: &SmemTile, + warp_tile: (usize, usize), + lane: u32, +) -> FragmentBPair { + let _ = (src, warp_tile, lane); + unreachable!("cute-rs `load_matrix_b_pair_k_major` executed outside device compilation") +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/numeric.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/numeric.rs new file mode 100644 index 0000000000..a47abb6a71 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/numeric.rs @@ -0,0 +1,432 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Explicit storage and conversions for reduced-precision number formats. +//! +//! A logical value and the bits that carry it are sometimes different. For +//! example, two FP4 values share one byte: +//! +//! ```text +//! one byte +//! ┌────────────┬────────────┐ +//! │ high FP4 │ low FP4 │ +//! │ bits 7..4 │ bits 3..0 │ +//! └────────────┴────────────┘ +//! ``` +//! +//! These types keep that meaning visible: +//! +//! ```text +//! logical value storage bits +//! Float4E2M1FN u8 holds 2 FP4 values +//! Float8E8M0FNU u8 holds 1 scale +//! Mxf4E2M1 f16 holds 4 FP4 values +//! UE8M0x4 u32 holds 4 scales +//! ``` +//! +//! [`Float4E2M1FNx2`] and [`UE8M0x2`] represent values already loaded into a +//! register. Their conversion methods describe the numeric result without +//! naming a backend instruction. CuTe-aware compilers retain the enclosing +//! high-level operation and choose the implementation for their backend. + +use crate::tensor::TensorElement; + +/// One BF16 value stored in sixteen bits, distinct from FP16 and integer data. +/// +/// Arithmetic requires explicit [`Self::to_f32`] widening and +/// [`Self::from_f32`] rounding. No arithmetic operators or implicit conversions +/// are supplied. Every bit pattern is valid, including signed zero and NaNs. +#[repr(transparent)] +#[derive(Clone, Copy, Debug)] +pub struct BFloat16(#[doc = "The unchanged IEEE-style BF16 storage bits."] u16); + +impl BFloat16 { + /// Preserve an existing BF16 bit pattern without conversion. + #[must_use] + #[inline(always)] + pub const fn from_bits(bits: u16) -> Self { + Self(bits) + } + + /// Return the unchanged BF16 storage bits. + #[must_use] + #[inline(always)] + pub const fn to_bits(self) -> u16 { + self.0 + } + + /// Round FP32 to BF16 using round-to-nearest, ties-to-even. + /// + /// Preserves signed zero and infinities; NaNs remain NaNs. + #[must_use] + #[inline(always)] + pub fn from_f32(value: f32) -> Self { + Self(ctm_rs::convert::f32_to_bf16_rne(value)) + } + + /// Widen to FP32 exactly, retaining all BF16 payload bits. + #[must_use] + #[inline(always)] + pub fn to_f32(self) -> f32 { + ctm_rs::convert::bf16_to_f32(self.0) + } +} + +impl TensorElement for BFloat16 {} + +/// Evaluate SiLU using stable FP32 sigmoid branches, then round once to BF16. +/// +/// Uses non-FTZ approximate base-two exponential and FTZ approximate reciprocal. +/// This is a device operation; callers retain an explicit BF16 boundary when +/// composing it with a later product (for example, a gated activation). +#[must_use] +#[inline(always)] +pub fn silu_bf16(value: BFloat16) -> BFloat16 { + let value = value.to_f32(); + let sigmoid = if value >= 0.0 { + let e = ctm_rs::float::ex2_approx_f32(-value * core::f32::consts::LOG2_E); + ctm_rs::float::rcp_approx_ftz_f32(1.0 + e) + } else { + let e = ctm_rs::float::ex2_approx_f32(value * core::f32::consts::LOG2_E); + e * ctm_rs::float::rcp_approx_ftz_f32(1.0 + e) + }; + BFloat16::from_f32(value * sigmoid) +} + +// SAFETY: repr(transparent) over u16, no padding, references or invalid patterns. +#[cfg(feature = "host")] +unsafe impl cuda_core::DeviceCopy for BFloat16 {} + +/// One logical `Float4E2M1FN` value, also called E2M1 or FP4. +/// +/// ```text +/// 4 bits: [sign][exponent exponent][fraction] +/// ``` +/// +/// Each value uses four bits, so one `u8` storage byte holds two values. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Float4E2M1FN; + +/// One unsigned `Float8E8M0FNU` scale, also called UE8M0. +/// +/// ```text +/// 8 bits: [exponent only] ──► power-of-two scale +/// ``` +/// +/// It has no sign or fraction bits. The all-ones byte represents NaN. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Float8E8M0FNU; + +/// One `Float8E4M3FN` scale, also called E4M3 (used by NVFP4 block scaling). +/// +/// ```text +/// 8 bits: [sign][4-bit exponent][3-bit mantissa] +/// ``` +/// +/// Unlike [`Float8E8M0FNU`], E4M3 carries three fraction bits so each 16-element NVFP4 +/// block scale can represent fractional steps such as `1.5` (`max_val / 6.0`). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Float8E4M3FN; + +/// Four neighboring [`Float8E8M0FNU`] scale bytes carried in one `u32`. +/// +/// ```text +/// u32: [scale 3][scale 2][scale 1][scale 0] +/// └─ lowest K group +/// ``` +/// +/// This type marks a memory view. [`UE8M0x2`] is the two-scale register value +/// used by one-thread conversions. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct UE8M0x4; + +/// Four neighboring [`Float8E4M3FN`] scale bytes carried in one `u32` (covering K=64 at SF=16). +/// +/// ```text +/// u32: [scale 3][scale 2][scale 1][scale 0] +/// K48..63 K32..47 K16..31 K0..15 +/// ``` +/// +/// This type marks a shared-memory or global-memory view of four 16-element +/// NVFP4 block scales. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct E4M3x4; + +/// Four neighboring [`Float4E2M1FN`] values packed for an MXFP4 MMA instruction. +/// +/// The four nibbles use the bits of one `f16` storage value. It is only a +/// 16-bit box; its bits do **not** represent an FP16 number. +/// +/// ```text +/// f16 bits: [FP4 3][FP4 2][FP4 1][FP4 0] +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Mxf4E2M1; + +impl TensorElement for Float4E2M1FN {} +impl TensorElement for Float8E8M0FNU {} +impl TensorElement for Float8E4M3FN {} +impl TensorElement for Mxf4E2M1 {} +impl TensorElement for UE8M0x4 {} +impl TensorElement for E4M3x4 {} + +/// Two [`Float4E2M1FN`] values packed into one byte. +/// +/// ```text +/// bits 7..4 = value 1 +/// bits 3..0 = value 0 +/// ``` +/// +/// Every bit pattern is valid, including positive and negative zero. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Float4E2M1FNx2(#[doc = "Low and high FP4 nibbles preserved in one byte."] u8); + +impl Float4E2M1FNx2 { + /// Treat `bits` as two packed E2M1 values without changing them. + #[must_use] + #[inline(always)] + pub const fn from_bits(bits: u8) -> Self { + Self(bits) + } + + /// Return the original packed byte. + #[must_use] + #[inline(always)] + pub const fn to_bits(self) -> u8 { + self.0 + } + + /// Convert both FP4 values to `f32`, returning the low nibble first. + /// + /// ```text + /// packed E2M1 byte ──► decode two logical values ──► (f32, f32) + /// ``` + /// + /// Both conversion steps are exact. + #[must_use] + #[inline(always)] + pub fn to_f32x2(self) -> (f32, f32) { + (e2m1_to_f32(self.0 & 0x0f), e2m1_to_f32(self.0 >> 4)) + } +} + +/// Two [`Float8E8M0FNU`] scales packed into one 16-bit register. +/// +/// ```text +/// bits 15..8 = scale 1 +/// bits 7..0 = scale 0 +/// ``` +/// +/// Every byte is valid. The byte `255` represents NaN. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct UE8M0x2(#[doc = "Two exponent-only scales, with the first scale in the low byte."] u16); + +impl UE8M0x2 { + /// Pack `low` as scale 0 and `high` as scale 1. + #[must_use] + #[inline(always)] + pub const fn from_bytes(low: u8, high: u8) -> Self { + Self((low as u16) | ((high as u16) << 8)) + } + + /// Treat `bits` as two packed scales without changing them. + #[must_use] + #[inline(always)] + pub const fn from_bits(bits: u16) -> Self { + Self(bits) + } + + /// Return the original packed 16 bits. + #[must_use] + #[inline(always)] + pub const fn to_bits(self) -> u16 { + self.0 + } + + /// Convert both scales to `f32`, returning the low byte first. + /// + /// ```text + /// packed UE8M0 pair ──► decode two powers of two ──► (f32, f32) + /// ``` + /// + /// UE8M0 values are powers of two, so both steps are exact. NaN stays NaN. + #[must_use] + #[inline(always)] + pub fn to_f32x2(self) -> (f32, f32) { + ( + ue8m0_to_f32(self.0 as u8), + ue8m0_to_f32((self.0 >> 8) as u8), + ) + } +} + +/// Two [`Float8E4M3FN`] scales packed into one 16-bit register. +/// +/// ```text +/// bits 15..8 = scale 1 (K group 1) +/// bits 7..0 = scale 0 (K group 0) +/// ``` +/// +/// Magnitude `0x7f` (`0x7f` and `0xff`) represents NaN; all other bit patterns +/// decode to exact finite `f32` values in `[-448.0, 448.0]`. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct E4M3x2(#[doc = "Two E4M3FN scales, with the first scale in the low byte."] u16); + +impl E4M3x2 { + /// Pack `low` as scale 0 and `high` as scale 1. + #[must_use] + #[inline(always)] + pub const fn from_bytes(low: u8, high: u8) -> Self { + Self((low as u16) | ((high as u16) << 8)) + } + + /// Treat `bits` as two packed E4M3 scales without changing them. + #[must_use] + #[inline(always)] + pub const fn from_bits(bits: u16) -> Self { + Self(bits) + } + + /// Return the original packed 16 bits. + #[must_use] + #[inline(always)] + pub const fn to_bits(self) -> u16 { + self.0 + } + + /// Convert both E4M3 scales to `f32`, returning the low byte first. + /// + /// ```text + /// packed E4M3 pair ──► decode two fractional scales ──► (f32, f32) + /// ``` + #[must_use] + #[inline(always)] + pub fn to_f32x2(self) -> (f32, f32) { + (e4m3_to_f32(self.0 as u8), e4m3_to_f32((self.0 >> 8) as u8)) + } +} + +#[inline(always)] +/// Decode one `Float8E4M3FN` byte to `f32`; `0x7f` and `0xff` become quiet NaN. +pub const fn e4m3_to_f32(byte: u8) -> f32 { + let sign = (byte >> 7) as u32; + let exp = ((byte >> 3) & 0x0f) as u32; + let mantissa = (byte & 0x07) as u32; + if exp == 0x0f && mantissa == 0x07 { + f32::from_bits(0x7fff_ffff) + } else if exp == 0 { + if mantissa == 0 { + f32::from_bits(sign << 31) + } else { + let mag = (mantissa as f32) * (1.0 / 512.0); + if sign == 0 { mag } else { -mag } + } + } else { + f32::from_bits((sign << 31) | ((exp + 120) << 23) | (mantissa << 20)) + } +} + +#[inline(always)] +/// Decode the low FP4 E2M1 nibble, including its sign bit, to f32. +fn e2m1_to_f32(nibble: u8) -> f32 { + let magnitude = match nibble & 0x07 { + 0 => 0.0, + 1 => 0.5, + 2 => 1.0, + 3 => 1.5, + 4 => 2.0, + 5 => 3.0, + 6 => 4.0, + _ => 6.0, + }; + if nibble & 0x08 == 0 { + magnitude + } else { + -magnitude + } +} + +#[inline(always)] +/// Decode one unsigned exponent-only scale; 0xff becomes NaN. +fn ue8m0_to_f32(byte: u8) -> f32 { + match byte { + // UE8M0 code zero denotes 2^-127, exactly halfway through the first + // f32 subnormal exponent bin. + 0 => f32::from_bits(0x0040_0000), + // The all-ones code is NaN. Keep a stable quiet payload without + // requiring a backend conversion intrinsic. + 0xff => f32::from_bits(0x7fff_ffff), + exponent => f32::from_bits((exponent as u32) << 23), + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Packing and unpacking FP4 pairs must preserve all eight storage bits. + #[test] + fn e2m1_pair_preserves_all_bits() { + for bits in 0u8..=u8::MAX { + assert_eq!(Float4E2M1FNx2::from_bits(bits).to_bits(), bits); + } + } + + /// The low UE8M0 byte must decode as the first scale and the high byte as the second. + #[test] + fn scale_pair_uses_low_byte_first() { + let pair = UE8M0x2::from_bytes(0x7f, 0x80); + assert_eq!(pair.to_bits(), 0x807f); + assert_eq!(UE8M0x2::from_bits(0x807f), pair); + } + + /// Every signed FP4 code must decode to the exact expected f32 bit pattern. + #[test] + fn e2m1_conversion_covers_every_code_bit_exactly() { + /// Independent positive FP4 reference values for all three-bit magnitude codes. + const MAGNITUDES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + for (code, expected) in MAGNITUDES.into_iter().enumerate() { + assert_eq!(e2m1_to_f32(code as u8).to_bits(), expected.to_bits()); + assert_eq!( + e2m1_to_f32(code as u8 | 0x08).to_bits(), + (-expected).to_bits() + ); + } + } + + /// Exponent-only decoding must handle the lowest exponent, normal codes and the NaN sentinel. + #[test] + fn ue8m0_conversion_handles_floor_normal_and_nan_codes() { + assert_eq!(ue8m0_to_f32(0).to_bits(), 0x0040_0000); + assert_eq!(ue8m0_to_f32(1).to_bits(), 0x0080_0000); + assert_eq!(ue8m0_to_f32(126), 0.5); + assert_eq!(ue8m0_to_f32(127), 1.0); + assert_eq!(ue8m0_to_f32(128), 2.0); + assert_eq!(ue8m0_to_f32(254).to_bits(), 0x7f00_0000); + assert!(ue8m0_to_f32(255).is_nan()); + } + + /// E4M3FN decoding must handle signed zero, subnormals, fractional normals, max 448.0, and NaN. + #[test] + fn e4m3_conversion_handles_zero_subnormal_fractional_max_and_nan() { + assert_eq!(e4m3_to_f32(0x00).to_bits(), 0x0000_0000); + assert_eq!(e4m3_to_f32(0x80).to_bits(), 0x8000_0000); + assert_eq!(e4m3_to_f32(0x01), 1.0 / 512.0); + assert_eq!(e4m3_to_f32(0x81), -1.0 / 512.0); + assert_eq!(e4m3_to_f32(0x38), 1.0); + assert_eq!(e4m3_to_f32(0x3c), 1.5); + assert_eq!(e4m3_to_f32(0xbc), -1.5); + assert_eq!(e4m3_to_f32(0x7e), 448.0); + assert_eq!(e4m3_to_f32(0xfe), -448.0); + assert!(e4m3_to_f32(0x7f).is_nan()); + assert!(e4m3_to_f32(0xff).is_nan()); + let pair = E4M3x2::from_bytes(0x38, 0x3c); + assert_eq!(pair.to_f32x2(), (1.0, 1.5)); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/pipeline.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/pipeline.rs new file mode 100644 index 0000000000..d21c0b1262 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/pipeline.rs @@ -0,0 +1,611 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Move Tensor Memory Accelerator (TMA) tiles through shared buffers safely. +//! +//! TMA is hardware that copies a complete tile between global memory and +//! shared memory. +//! +//! A thread is one GPU worker. A warp is 32 threads that run together. A CTA, +//! also called a block, is a group of warps that share memory. +//! +//! A *stage* is one reusable shared-memory buffer. Producer threads load the +//! next tile. Consumer warps use a full tile for computation. Two barriers +//! per stage act like ready signals: +//! +//! ```text +//! producer consumer warps +//! wait until stage is EMPTY wait until stage is FULL +//! mark expected TMA byte count read and compute +//! start TMA load lane 0 of each warp marks EMPTY +//! move to next stage move to next stage +//! ``` +//! +//! Stages form a ring: `0 -> 1 -> ... -> 0`. A one-bit phase distinguishes +//! the new use of a slot from its previous use. New buffers are empty, so the +//! producer starts at phase 1 and the consumer at phase 0. The phase flips +//! only when the stage index wraps to slot 0, as in the CUTLASS pipeline of +//! the same name. + +use core::marker::PhantomData; +use core::mem::{align_of, size_of}; + +use cuda_device::barrier::Barrier; + +/// Marks a cursor used by the thread that starts TMA loads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Producer {} + +/// Marks a cursor used by warps that read loaded tiles. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Consumer {} + +/// A cursor pointing to one slot in the circular stage buffer. +/// +/// `Role` prevents using a consumer cursor in a producer operation, or the +/// reverse. The cursor stores only a slot number and a phase bit. The kernel +/// keeps its K-loop counter separately. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct PipelineState { + /// Current reusable stage index in 0..STAGES. + slot: u32, + /// Parity awaited on this role’s current barrier slot. + phase: u32, + /// Distinguishes producer and consumer initial phases and allowed operations. + role: PhantomData, +} + +impl PipelineState { + /// Start writing at slot 0. Phase 1 sees a new buffer as empty. + #[must_use] + #[inline(always)] + pub fn new() -> Self { + __compiler::pipeline_state_new_producer::() + } +} + +impl Default for PipelineState { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +impl PipelineState { + /// Start reading the first tile produced in slot 0. + #[must_use] + #[inline(always)] + pub fn new() -> Self { + __compiler::pipeline_state_new_consumer::() + } +} + +impl Default for PipelineState { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +impl PipelineState { + /// Current stage index in `0..STAGES`, CuTeDSL's `PipelineState.index`. + /// + /// It is a `usize` so it can scale a shared-buffer offset directly; the + /// SM100 attention ring's `RingCursor::index` is a `u32`. + #[must_use] + #[inline(always)] + pub fn index(&self) -> usize { + __compiler::pipeline_state_slot(self) + } + + /// Phase bit expected from the current slot's barrier. + #[must_use] + #[inline(always)] + pub const fn phase(&self) -> u32 { + self.phase + } + + /// Move to the next slot and flip the phase when the ring wraps. + #[inline(always)] + pub fn advance(&mut self) { + __compiler::pipeline_state_advance(self); + } +} + +/// A CTA-local TMA load pipeline with one full and one empty barrier per stage. +/// +/// This is CuTeDSL's `pipeline.PipelineTmaAsync`. +/// +/// - `STAGES` is the number of shared tile buffers. +/// - `CONSUMER_WARPS` is the number of warps that read each tile. Lane 0 of +/// each warp releases the stage through [`Self::consumer_release`]. +/// - `TX_BYTES` is the total number of bytes loaded into one stage. +/// +/// This value stores only two pointers. The caller owns the barriers and tile +/// buffers in shared memory. +#[derive(Clone, Copy)] +pub struct PipelineTmaAsync { + /// First of STAGES full barriers, completed by TMA transaction bytes. + full: *mut Barrier, + /// First of STAGES empty barriers, released by consumer warps. + empty: *mut Barrier, +} + +/// Track TMA stores from shared memory to global memory. +/// +/// Hardware tracks these stores in copy groups, so this type stores no data +/// and uses no memory barriers: +/// +/// ```text +/// issue one or more shared -> global copies +/// │ +/// ▼ +/// producer_commit() put them in one group +/// │ +/// ▼ +/// producer_acquire() wait until the oldest shared buffer is reusable +/// │ +/// ▼ +/// producer_tail() wait until no store still reads shared memory +/// ``` +/// +/// `STAGES` is the number of shared result buffers. After +/// [`Self::producer_acquire`], at most `STAGES - 1` committed groups still +/// read shared memory, as in CuTeDSL's `pipeline.PipelineTmaStore`. The SM120 +/// device compiler accepts `STAGES == 1` today; other stage counts are +/// rejected at device compile time. The SM100 form with the same name and +/// methods, over a `u32` stage count, is [`crate::sm100::PipelineTmaStore`]. +#[derive(Clone, Copy, Debug, Default)] +pub struct PipelineTmaStore; + +impl PipelineTmaStore { + /// Most store groups that may still read shared memory after acquire. + pub const MAX_PENDING: u32 = { + assert!(STAGES > 0); + assert!(STAGES <= u32::MAX as usize); + (STAGES - 1) as u32 + }; + + /// Create a store pipeline. The value occupies no bytes. + #[must_use] + #[inline(always)] + pub const fn new() -> Self { + let _ = Self::MAX_PENDING; + Self + } + + /// Put this thread's TMA stores since the last commit into one group. + #[inline(always)] + pub fn producer_commit(&self) { + __compiler::tma_store_producer_commit::(); + } + + /// Wait until the oldest shared result buffer can be changed or reused. + /// + /// This waits for TMA to finish reading shared memory. The later global + /// write may still be in progress; that does not prevent buffer reuse. + #[inline(always)] + pub fn producer_acquire(&self) { + __compiler::tma_store_producer_acquire::(); + } + + /// Wait until every TMA store has finished reading its shared source. + #[inline(always)] + pub fn producer_tail(&self) { + __compiler::tma_store_producer_tail::(); + } +} + +impl + PipelineTmaAsync +{ + /// Required alignment of each shared-memory barrier array. + pub const STORAGE_ALIGN: usize = align_of::(); + + /// Bytes used by adjacent `FULL[STAGES]` and `EMPTY[STAGES]` arrays. + /// + /// Using this constant also checks the type parameters. Each empty + /// barrier expects one arrival from lane 0 of every consumer warp. + /// + /// ```compile_fail + /// use cute_rs::PipelineTmaAsync; + /// const _: usize = PipelineTmaAsync::<0, 1, 16>::STORAGE_BYTES; + /// ``` + pub const STORAGE_BYTES: usize = { + assert!(STAGES > 0); + assert!(STAGES <= u32::MAX as usize); + assert!(CONSUMER_WARPS > 0); + assert!(CONSUMER_WARPS <= 32); + assert!(TX_BYTES > 0); + 2 * STAGES * size_of::() + }; + + /// Use one allocation containing the full barriers, then empty barriers. + /// + /// ```text + /// base -> [FULL x STAGES][EMPTY x STAGES] + /// ``` + /// + /// # Safety + /// + /// `base` must point to at least [`Self::STORAGE_BYTES`] writable bytes in + /// shared memory and be aligned to [`Self::STORAGE_ALIGN`]. The first + /// `STAGES` entries must be the full barriers and the next `STAGES` the + /// empty barriers. This pipeline must have exclusive use of that storage + /// until all asynchronous work and [`Self::producer_tail`] have finished. + #[must_use] + #[inline(always)] + pub unsafe fn from_raw_base(base: *mut Barrier) -> Self { + unsafe { __compiler::tma_load_pipeline_from_raw_base(base) } + } + + /// Use separately allocated full and empty barrier arrays. + /// + /// # Safety + /// + /// `full` and `empty` must point to different shared-memory arrays of + /// `STAGES` barriers, each aligned to [`Self::STORAGE_ALIGN`]. This + /// pipeline must have exclusive use of both arrays until all asynchronous + /// work and [`Self::producer_tail`] have finished. + #[must_use] + #[inline(always)] + pub unsafe fn from_raw_parts(full: *mut Barrier, empty: *mut Barrier) -> Self { + let _ = Self::STORAGE_BYTES; + Self { full, empty } + } + + /// Initialize all full and empty barriers. + /// + /// A full barrier expects one producer arrival. An empty barrier expects + /// one arrival from lane 0 of each consumer warp. Every thread in the CTA + /// must call this method on the same control-flow path. Only `init_thread` + /// writes the barriers; the method publishes those writes and synchronizes + /// the full CTA before returning. + /// + /// # Safety + /// + /// The pointer rules of the constructor must still hold. `init_thread` + /// must be a valid thread number in this CTA. Neither barrier array may + /// already be active or accessed at the same time by another pipeline. + #[inline(always)] + pub unsafe fn init(self, init_thread: u32) { + unsafe { __compiler::tma_load_pipeline_init(self, init_thread) }; + } + + /// Wait until the current write slot is empty. + /// + /// The first slot starts ready, so the first load uses the same operation + /// as every later load. + /// + /// # Safety + /// + /// `state` must be this pipeline's current producer cursor, owned by one + /// producer leader. An acquire for the same slot must not overlap another + /// acquire or be repeated without advancing the cursor. + #[inline(always)] + pub unsafe fn producer_acquire(self, state: &PipelineState) { + unsafe { __compiler::pipeline_producer_acquire(self, state) }; + } + + /// Tell the current full barrier to expect `TX_BYTES` from TMA. + /// + /// The producer leader calls this once after [`Self::producer_acquire`] + /// and before starting the stage's TMA loads. Pass the returned barrier + /// pointer to those loads. + /// + /// # Safety + /// + /// `state` must point to the slot just acquired from this pipeline. The + /// TMA loads attached to the returned barrier must complete exactly + /// `TX_BYTES` bytes in total. + #[must_use] + #[inline(always)] + pub unsafe fn producer_expect_tx( + self, + state: &PipelineState, + ) -> *mut Barrier { + unsafe { __compiler::pipeline_producer_expect_tx(self, state) } + } + + /// Wait until TMA has filled the consumer's current slot. + /// + /// # Safety + /// + /// `state` must be this pipeline's current consumer cursor. Every consumer + /// warp must complete this wait before reading the matching shared buffer. + #[inline(always)] + pub unsafe fn consumer_wait(self, state: &PipelineState) { + unsafe { __compiler::pipeline_consumer_wait(self, state) }; + } + + /// Mark the current read slot empty after all consumers finish. + /// + /// Every lane in a consumer warp may call this method once, but only + /// hardware lane 0 updates the barrier. Exactly `CONSUMER_WARPS` warps + /// must each contribute one lane-0 update after their last read. + /// + /// # Safety + /// + /// `state` must point to the slot previously passed to + /// [`Self::consumer_wait`]. The calling warp must have finished all reads + /// from that shared buffer. Exactly `CONSUMER_WARPS` different warps must + /// participate, with no warp contributing more than one lane-0 update. + #[inline(always)] + pub unsafe fn consumer_release(self, state: &PipelineState) { + unsafe { __compiler::pipeline_consumer_release(self, state) }; + } + + /// Wait until consumers have released every stage that may still be busy. + /// + /// Starting at the producer's current slot, this waits through one full + /// turn of the ring. Slots never used by a short startup pass immediately; + /// used slots wait for the matching phase. This is CUTLASS's + /// producer-tail rule. + /// + /// # Safety + /// + /// `state` must be the producer cursor immediately after its final issued + /// stage. Call this exactly once before the producer exits or any pipeline + /// shared storage is reused. + #[inline(always)] + pub unsafe fn producer_tail(self, state: PipelineState) { + unsafe { __compiler::pipeline_producer_tail(self, state) }; + } +} + +/// Stable compiler boundaries for the TMA load-pipeline protocol. +/// +/// Public methods keep the normal Rust API. Device compilation recognizes +/// these exact calls; protocol state math remains ordinary testable Rust. +#[doc(hidden)] +pub mod __compiler { + use super::*; + + /// Compiler boundary waiting until the oldest shared store buffer can be reused. + #[inline(never)] + pub fn tma_store_producer_acquire() { + let _ = PipelineTmaStore::::MAX_PENDING; + unreachable!("cute-rs TMA-store acquire executed outside recognized device compilation") + } + + /// Compiler boundary committing the issuer’s pending TMA stores as one group. + #[inline(never)] + pub fn tma_store_producer_commit() { + let _ = PipelineTmaStore::::MAX_PENDING; + unreachable!("cute-rs TMA-store commit executed outside recognized device compilation") + } + + /// Compiler boundary waiting until no committed group still reads shared memory. + #[inline(never)] + pub fn tma_store_producer_tail() { + let _ = PipelineTmaStore::::MAX_PENDING; + unreachable!("cute-rs TMA-store tail executed outside recognized device compilation") + } + + /// Construct slot zero with empty parity one, validating a nonempty ring. + #[inline(never)] + pub fn pipeline_state_new_producer() -> PipelineState { + const { assert!(STAGES > 0) }; + const { assert!(STAGES <= u32::MAX as usize) }; + PipelineState { + slot: 0, + phase: 1, + role: PhantomData, + } + } + + /// Construct slot zero with full parity zero, validating a nonempty ring. + #[inline(never)] + pub fn pipeline_state_new_consumer() -> PipelineState { + const { assert!(STAGES > 0) }; + const { assert!(STAGES <= u32::MAX as usize) }; + PipelineState { + slot: 0, + phase: 0, + role: PhantomData, + } + } + + /// Compiler boundary exposing the current reusable stage index. + #[inline(never)] + pub fn pipeline_state_slot( + state: &PipelineState, + ) -> usize { + state.slot as usize + } + + /// Advance one slot and toggle parity only when the ring wraps. + #[inline(never)] + pub fn pipeline_state_advance( + state: &mut PipelineState, + ) { + const { assert!(STAGES > 0) }; + const { assert!(STAGES <= u32::MAX as usize) }; + state.slot = state.slot.wrapping_add(1); + if state.slot as usize == STAGES { + state.slot = 0; + state.phase ^= 1; + } + } + + /// Compiler boundary splitting adjacent full and empty barrier arrays. + /// + /// # Safety + /// The base must satisfy the public PipelineTmaAsync storage and lifetime contract. + #[inline(never)] + pub unsafe fn tma_load_pipeline_from_raw_base< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + base: *mut Barrier, + ) -> PipelineTmaAsync { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + PipelineTmaAsync { + full: base, + empty: unsafe { base.add(STAGES) }, + } + } + + /// Compiler boundary initializing full and empty arrival counts for every stage. + /// + /// # Safety + /// One initializer must publish fresh barrier storage before participants use it. + #[inline(never)] + pub unsafe fn tma_load_pipeline_init< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + init_thread: u32, + ) { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline.full, pipeline.empty, init_thread); + unreachable!("cute-rs TMA-load init executed outside recognized device compilation") + } + + /// Compiler boundary waiting for consumers to release the selected stage. + /// + /// # Safety + /// The cursor must name the producer’s current live phase with matching consumer releases. + #[inline(never)] + pub unsafe fn pipeline_producer_acquire< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + state: &PipelineState, + ) { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer acquire executed outside recognized device compilation") + } + + /// Compiler boundary attaching the stage’s fixed TMA byte count. + /// + /// # Safety + /// The producer must acquire first and issue exactly TX_BYTES to the returned completion barrier. + #[inline(never)] + pub unsafe fn pipeline_producer_expect_tx< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + state: &PipelineState, + ) -> *mut Barrier { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer expect-tx executed outside recognized device compilation") + } + + /// Compiler boundary waiting until the selected stage’s TMA copies complete. + /// + /// # Safety + /// The consumer cursor must match the producer’s current stage and phase. + #[inline(never)] + pub unsafe fn pipeline_consumer_wait< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + state: &PipelineState, + ) { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs consumer wait executed outside recognized device compilation") + } + + /// Compiler boundary signaling that a consumer warp no longer reads the stage. + /// + /// # Safety + /// The elected releasing lane must represent one completed counted consumer warp. + #[inline(never)] + pub unsafe fn pipeline_consumer_release< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + state: &PipelineState, + ) { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs consumer release executed outside recognized device compilation") + } + + /// Compiler boundary draining all stage ownership before pipeline storage is retired. + /// + /// # Safety + /// All producers must have finished issuing and all expected consumers must eventually release. + #[inline(never)] + pub unsafe fn pipeline_producer_tail< + const STAGES: usize, + const CONSUMER_WARPS: u32, + const TX_BYTES: u32, + >( + pipeline: PipelineTmaAsync, + state: PipelineState, + ) { + let _ = PipelineTmaAsync::::STORAGE_BYTES; + let _ = (pipeline, state); + unreachable!("cute-rs producer tail executed outside recognized device compilation") + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// New producer and consumer cursors must start at slot zero with opposite parities. + #[test] + fn producer_and_consumer_use_opposite_initial_phases() { + let producer = PipelineState::::new(); + let consumer = PipelineState::::new(); + assert_eq!((producer.index(), producer.phase()), (0, 1)); + assert_eq!((consumer.index(), consumer.phase()), (0, 0)); + } + + /// Advancing must toggle parity only when the stage index wraps. + #[test] + fn slot_wrap_flips_phase_once_per_ring() { + let mut state = PipelineState::::new(); + let expected = [(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1), (0, 0)]; + for (index, phase) in expected { + assert_eq!((state.index(), state.phase()), (index, phase)); + state.advance(); + } + } + + /// Pipeline state must consist only of slot and parity register values. + #[test] + fn state_is_exactly_two_device_register_values() { + assert_eq!(size_of::>(), 8); + assert_eq!(size_of::>(), 8); + } + + /// Static pipeline storage must fit aligned full and empty arrays for every stage. + #[test] + fn static_contract_sizes_two_aligned_barrier_rings() { + /// Three-stage fixture with eight consumer warps and 32 KiB of TMA bytes per stage. + type Pipeline = PipelineTmaAsync<3, 8, 32_768>; + /// Required storage for both three-stage barrier arrays. + const BYTES: usize = Pipeline::STORAGE_BYTES; + assert_eq!(Pipeline::STORAGE_ALIGN, 8); + assert_eq!(BYTES, 6 * size_of::()); + } + + /// Store pipelines must use no runtime storage and retain at most STAGES-1 pending groups. + #[test] + fn tma_store_pipeline_is_zero_storage_and_uses_stage_minus_one() { + assert_eq!(size_of::>(), 0); + assert_eq!(PipelineTmaStore::<1>::MAX_PENDING, 0); + assert_eq!(PipelineTmaStore::<4>::MAX_PENDING, 3); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/reduction.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/reduction.rs new file mode 100644 index 0000000000..ff0c7c8c7c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/reduction.rs @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Register reduction states with explicit tie and invalid-input semantics. + +/// A maximum with its original coordinate, choosing the lowest index on ties. +/// +/// Nonfinite observations make validity sticky-false. The numeric candidate may +/// still change, so consumers must check [`Self::is_valid`] before using it. +/// An empty state has index `u32::MAX`; callers must supply at least one value +/// when they require a valid in-bounds index. +#[derive(Clone, Copy, Debug)] +pub struct ArgMax { + value: f32, + index: u32, + valid: bool, +} + +impl ArgMax { + /// Identity state for a finite-value argmax reduction. + #[must_use] + #[inline(always)] + pub const fn empty() -> Self { + Self { + value: f32::NEG_INFINITY, + index: u32::MAX, + valid: true, + } + } + + /// Include a value and its original coordinate without reassociation. + #[must_use] + #[inline(always)] + pub fn observe(mut self, value: f32, index: u32) -> Self { + self.valid = self.valid && value.to_bits() & 0x7f80_0000 != 0x7f80_0000; + if value > self.value || (value == self.value && index < self.index) { + self.value = value; + self.index = index; + } + self + } + + /// Combine two partitions, retaining original indices and invalid status. + #[must_use] + #[inline(always)] + pub fn merge(mut self, other: Self) -> Self { + self.valid = self.valid && other.valid; + if other.value > self.value || (other.value == self.value && other.index < self.index) { + self.value = other.value; + self.index = other.index; + } + self + } + + /// Selected numeric candidate; meaningful as a finite maximum when valid. + #[must_use] + #[inline(always)] + pub const fn value(self) -> f32 { + self.value + } + + /// Original coordinate of the selected candidate, or `u32::MAX` if empty. + #[must_use] + #[inline(always)] + pub const fn index(self) -> u32 { + self.index + } + + /// Whether every observed input was finite (also true for an empty state). + #[must_use] + #[inline(always)] + pub const fn is_valid(self) -> bool { + self.valid + } + + /// Reduce one state per lane in XOR order 16, 8, 4, 2, 1. + /// + /// # Safety + /// All 32 lanes must execute convergently, each passing its actual lane ID + /// in `0..32`. No lane may exit before the final full-warp validity vote. + #[must_use] + #[inline(always)] + pub unsafe fn reduce_warp(mut self, lane: u32) -> Self { + let mut offset = 16; + while offset != 0 { + let value = ctm_rs::warp::shuffle_xor_f32(self.value, offset); + // SAFETY: the full-warp caller contract includes every source lane. + let index = unsafe { ctm_rs::warp::shuffle_sync(u32::MAX, self.index, lane ^ offset) }; + // Validity is reduced separately, exactly once, after candidates. + self = self.merge(Self { + value, + index, + valid: true, + }); + offset /= 2; + } + // SAFETY: all lanes reach this vote under the caller contract. + self.valid = unsafe { ctm_rs::warp::all(self.valid) }; + self + } +} + +/// Select the first maximum of a BF16 tensor using one complete warp. +/// +/// Values are partitioned cyclically over 32 lanes; each lane visits its values +/// in ascending coordinate order before the fixed XOR reduction. All lanes +/// receive the same result. Any nonfinite input invalidates that result. +/// +/// # Safety +/// All 32 lanes must execute convergently with the same input view and their +/// actual lane IDs. The tensor length must fit in `u32`; consumers must reject +/// an empty or invalid result before using its index as an address. +#[must_use] +#[inline(always)] +pub unsafe fn argmax_bf16( + input: crate::tensor::partition::StridedTensor<'_, crate::numeric::BFloat16, L>, + lane: u32, +) -> ArgMax { + let values = input.local_partition::<32>(lane as usize); + let mut index = lane; + let local = values.fold(ArgMax::empty(), |state, value| { + let state = state.observe(value.to_f32(), index); + // The terminal increment is not used as an address and may wrap for a + // u32-sized input. Each observed coordinate itself remains in bounds. + index = index.wrapping_add(32); + state + }); + // SAFETY: the caller provides the full-warp participation/lane contract. + unsafe { local.reduce_warp(lane) } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/scheduler.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/scheduler.rs new file mode 100644 index 0000000000..090b42db90 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/scheduler.rs @@ -0,0 +1,449 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Reuse each CUDA block (CTA) for several output tiles. +//! +//! A thread is one GPU worker. A warp is 32 threads that run together. A +//! block, also called a CTA, is a group of warps. A persistent kernel launches +//! only as many blocks as the GPU can keep active. Each block finishes one +//! output tile, then takes another. +//! +//! First, number every `(M tile, N tile, batch)` position. M changes first: +//! +//! ```text +//! for M_TILES=3, N_TILES=2: +//! +//! linear 0 -> (m=0, n=0, batch=0) +//! linear 1 -> (m=1, n=0, batch=0) +//! linear 2 -> (m=2, n=0, batch=0) +//! linear 3 -> (m=0, n=1, batch=0) +//! ``` +//! +//! Then block `i` visits `i`, `i + block_count`, and so on: +//! +//! ```text +//! 3 blocks, 8 tiles +//! block 0: 0 -> 3 -> 6 +//! block 1: 1 -> 4 -> 7 +//! block 2: 2 -> 5 +//! ``` +//! +//! Every tile is visited once, without a shared counter. +//! [`StaticPersistentTileScheduler::new_1d`] uses only the X launch dimension, +//! so launch `(active_blocks, 1, 1)`. + +/// Coordinates of one output tile selected by the scheduler. +/// +/// This is CuTeDSL's `WorkTileInfo`. +/// +/// This value describes a location; it does not enforce exclusive access. +/// Warps and threads inside the block must still write different elements of +/// the selected tile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WorkTileInfo { + /// Flattened tile coordinate with M changing fastest. + linear: usize, + /// Tile coordinate along the output row dimension. + m_tile: usize, + /// Tile coordinate along the output column dimension. + n_tile: usize, + /// Batch coordinate of this output tile. + batch: usize, +} + +impl + WorkTileInfo +{ + /// Return the `(M tile, N tile, batch)` coordinates together. + /// + /// This is CuTeDSL's `WorkTileInfo.tile_idx`, the `(m, n, l)` coordinate + /// of the tile. Kernels normally unpack it once at the start of a tile + /// iteration. The flattened tile number is [`Self::linear_idx`]. + #[must_use] + #[inline(always)] + pub fn tile_idx(self) -> (usize, usize, usize) { + __compiler::work_tile_coordinates(self) + } + + /// Tile number when M changes first, then N, then batch. + /// + /// This and the three accessors below read the fields on the host. In + /// device code the selected tile is a compiler handle, so read its + /// coordinates through [`Self::tile_idx`]. + #[must_use] + #[inline(always)] + pub const fn linear_idx(self) -> usize { + self.linear + } + + /// Tile position along M. + #[must_use] + #[inline(always)] + pub const fn m_tile(self) -> usize { + self.m_tile + } + + /// Tile position along N. + #[must_use] + #[inline(always)] + pub const fn n_tile(self) -> usize { + self.n_tile + } + + /// Batch position, also called logical L. + #[must_use] + #[inline(always)] + pub const fn batch(self) -> usize { + self.batch + } +} + +/// Assign a fixed output-tile grid across persistent CUDA blocks. +/// +/// Every block starts at `blockIdx.x`. Each call to +/// [`Self::advance_to_next_work`] adds +/// `gridDim.x`, so blocks visit different tile numbers without a shared +/// counter: +/// +/// ```text +/// start = blockIdx.x +/// next = current + gridDim.x +/// stop = current >= TOTAL_TILES +/// ``` +/// +/// The caller owns accumulator and pipeline state. A multi-stage pipeline's +/// slot and phase bit must continue across output tiles; that bit marks each +/// new reuse of a buffer. Only the K-loop's tile counter starts over for each +/// output tile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct StaticPersistentTileScheduler< + const M_TILES: usize, + const N_TILES: usize, + const BATCHES: usize, +> { + /// Next flattened output tile assigned to this persistent block. + current_linear: usize, + /// Number of persistent blocks; advance adds this stride to the tile index. + stride: usize, +} + +impl + StaticPersistentTileScheduler +{ + /// Total output tiles: `M_TILES * N_TILES * BATCHES`. + /// + /// Using this constant rejects zero sizes and multiplication overflow at + /// compile time. + pub const TOTAL_TILES: usize = { + assert!(M_TILES > 0, "M_TILES must be positive"); + assert!(N_TILES > 0, "N_TILES must be positive"); + assert!(BATCHES > 0, "BATCHES must be positive"); + let mn = match M_TILES.checked_mul(N_TILES) { + Some(value) => value, + None => panic!("M_TILES * N_TILES overflows usize"), + }; + match mn.checked_mul(BATCHES) { + Some(value) => value, + None => panic!("tile count overflows usize"), + } + }; + + /// Create a scheduler from this block's X index and the X grid size. + /// + /// CUDA guarantees `blockIdx.x < gridDim.x` and `gridDim.x > 0`, so this + /// path needs no runtime error check. + #[must_use] + #[inline(always)] + pub fn new_1d() -> Self { + __compiler::scheduler_new_1d::() + } + + /// Create one block's schedule from explicit numbers. + /// + /// This form is useful in CPU tests and emulators. It returns `None` when + /// `worker_count` is zero or `worker` is outside `0..worker_count`. + #[must_use] + #[inline(always)] + pub const fn for_worker(worker: usize, worker_count: usize) -> Option { + let _ = Self::TOTAL_TILES; + if worker_count == 0 || worker >= worker_count { + None + } else { + Some(Self { + current_linear: worker, + stride: worker_count, + }) + } + } + + /// True while this block still has an output tile to process. + /// + /// CuTeDSL reads this as `is_valid_tile` on the tile that + /// `get_current_work` returns. Here the test is on the scheduler, so the + /// bounds check stays outside the coordinate arithmetic in the kernel + /// loop. + #[must_use] + #[inline(always)] + pub fn has_work(&self) -> bool { + __compiler::scheduler_has_work(self) + } + + /// Return the current tile after [`Self::has_work`] returned true. + /// + /// This is CuTeDSL's `get_current_work`. Call it only while + /// [`Self::has_work`] is true; past the end of this block's schedule the + /// coordinates are meaningless. Host tests drive the same two calls. + #[must_use] + #[inline(always)] + pub fn get_current_work(&self) -> WorkTileInfo { + __compiler::scheduler_current_tile(self) + } + + /// Move to this block's next output tile. + /// + /// This is CuTeDSL's `advance_to_next_work`. The step is always one grid + /// of blocks, so there is no `advance_count` argument. + #[inline(always)] + pub fn advance_to_next_work(&mut self) { + __compiler::scheduler_advance(self); + } + + /// Choose the number of blocks for a one-dimensional persistent launch. + /// + /// This is CuTeDSL's `get_grid_shape`. CuTeDSL computes it from a + /// `PersistentTileSchedulerParams` object and returns an `(x, y, z)` + /// triple; this scheduler has no params object and its grid is one + /// dimensional, so the result is the X extent alone, for a launch of + /// `(blocks, 1, 1)`. + /// + /// The result is the smaller of the output tile count and the number of + /// blocks the GPU can keep active: + /// + /// ```text + /// active capacity = SM count * active blocks per SM + /// launch blocks = min(total tiles, active capacity) + /// ``` + /// + /// An SM is one streaming multiprocessor on the GPU. Get + /// `active_ctas_per_sm` from CUDA's occupancy query, which reports how + /// many blocks fit on each SM for this compiled kernel. Register and + /// dynamic shared-memory use change that number. This returns `None` for + /// zero capacity or an X grid larger than CUDA's `u32` launch limit. + #[must_use] + #[inline(always)] + pub const fn get_grid_shape(sm_count: u32, active_ctas_per_sm: u32) -> Option { + let total = Self::TOTAL_TILES as u64; + let resident = (sm_count as u64) * (active_ctas_per_sm as u64); + if resident == 0 { + return None; + } + let workers = if total < resident { total } else { resident }; + if workers == 0 || workers > u32::MAX as u64 { + None + } else { + Some(workers as u32) + } + } +} + +/// Stable compiler boundaries for the five scheduler operations. +/// +/// The public methods above stay pleasant to use. Device compilation +/// recognizes these exact functions; pure coordinate/state operations remain +/// ordinary testable Rust. +#[doc(hidden)] +pub mod __compiler { + use super::{StaticPersistentTileScheduler, WorkTileInfo}; + + /// Device-only boundary reading the block index and grid stride for a persistent schedule. + #[inline(never)] + pub fn scheduler_new_1d() + -> StaticPersistentTileScheduler { + let _ = StaticPersistentTileScheduler::::TOTAL_TILES; + unreachable!("cute-rs scheduler creation executed outside recognized device compilation") + } + + /// Test whether the next flattened tile is still within the fixed problem. + #[inline(never)] + pub fn scheduler_has_work( + scheduler: &StaticPersistentTileScheduler, + ) -> bool { + scheduler.current_linear + < StaticPersistentTileScheduler::::TOTAL_TILES + } + + /// Decode the current tile with M changing fastest, then N and batch. + #[inline(never)] + pub fn scheduler_current_tile< + const M_TILES: usize, + const N_TILES: usize, + const BATCHES: usize, + >( + scheduler: &StaticPersistentTileScheduler, + ) -> WorkTileInfo { + let linear = scheduler.current_linear; + let n_batch = linear / M_TILES; + let m_tile = linear - n_batch * M_TILES; + let batch = n_batch / N_TILES; + let n_tile = n_batch - batch * N_TILES; + WorkTileInfo { + linear, + m_tile, + n_tile, + batch, + } + } + + /// Expose the three logical coordinates without recomputing them. + #[inline(never)] + pub fn work_tile_coordinates< + const M_TILES: usize, + const N_TILES: usize, + const BATCHES: usize, + >( + tile: WorkTileInfo, + ) -> (usize, usize, usize) { + (tile.m_tile, tile.n_tile, tile.batch) + } + + /// Advance by one grid stride, saturating completed schedules instead of wrapping. + #[inline(never)] + pub fn scheduler_advance( + scheduler: &mut StaticPersistentTileScheduler, + ) { + scheduler.current_linear = scheduler.current_linear.saturating_add(scheduler.stride); + } +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + use core::mem::size_of; + + /// Flat tile decoding must advance M first, then N, then batch across boundaries. + #[test] + fn m_fastest_mapping_crosses_n_and_batch_boundaries() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<3, 2, 2>; + let expected = [ + (0, 0, 0, 0), + (1, 1, 0, 0), + (2, 2, 0, 0), + (3, 0, 1, 0), + (5, 2, 1, 0), + (6, 0, 0, 1), + (11, 2, 1, 1), + ]; + + for (linear, m_tile, n_tile, batch) in expected { + let scheduler = Scheduler::for_worker(linear, Scheduler::TOTAL_TILES).unwrap(); + assert!(scheduler.has_work()); + let tile = scheduler.get_current_work(); + assert_eq!( + ( + tile.linear_idx(), + tile.m_tile(), + tile.n_tile(), + tile.batch() + ), + (linear, m_tile, n_tile, batch) + ); + assert_eq!(tile.tile_idx(), (m_tile, n_tile, batch)); + } + } + + /// A 170-block persistent grid must visit all 256 output tiles exactly once. + #[test] + fn grid_stride_170_covers_the_full_16_by_16_problem_once() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + /// Persistent block count intentionally smaller than the 256-tile problem. + const WORKERS: usize = 170; + let mut seen = [false; Scheduler::TOTAL_TILES]; + + for worker in 0..WORKERS { + let mut scheduler = Scheduler::for_worker(worker, WORKERS).unwrap(); + while scheduler.has_work() { + let tile = scheduler.get_current_work(); + let linear = tile.linear_idx(); + assert!(!seen[linear], "tile {linear} was scheduled twice"); + seen[linear] = true; + assert_eq!( + linear, + tile.batch() * 16 * 16 + tile.n_tile() * 16 + tile.m_tile() + ); + scheduler.advance_to_next_work(); + } + } + + assert!(seen.into_iter().all(|visited| visited)); + } + + /// An incomplete final scheduling wave must give each worker the expected number of tiles. + #[test] + fn uneven_final_wave_has_the_expected_worker_lengths() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + + for worker in 0..170 { + let mut scheduler = Scheduler::for_worker(worker, 170).unwrap(); + let mut count = 0; + while scheduler.has_work() { + count += 1; + scheduler.advance_to_next_work(); + } + assert_eq!(count, if worker < 86 { 2 } else { 1 }); + } + } + + /// Empty grids and out-of-range worker indices must be rejected before scheduling. + #[test] + fn worker_constructor_rejects_invalid_physical_grids() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<1, 1, 1>; + assert!(Scheduler::for_worker(0, 0).is_none()); + assert!(Scheduler::for_worker(1, 1).is_none()); + assert!(Scheduler::for_worker(2, 1).is_none()); + } + + /// Grid-shape sizing must cap at total work and reject arithmetic or launch-width overflow. + #[test] + fn grid_shape_is_checked_and_capped() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + assert_eq!(Scheduler::get_grid_shape(170, 1), Some(170)); + assert_eq!(Scheduler::get_grid_shape(170, 2), Some(256)); + assert_eq!(Scheduler::get_grid_shape(512, 8), Some(256)); + assert_eq!(Scheduler::get_grid_shape(0, 1), None); + assert_eq!(Scheduler::get_grid_shape(170, 0), None); + + /// Problem whose tile count exceeds the supported physical launch width. + type TooWide = StaticPersistentTileScheduler<{ u32::MAX as usize }, 2, 1>; + assert_eq!(TooWide::get_grid_shape(u32::MAX, u32::MAX), None); + } + + /// Advancing an exhausted scheduler must not wrap back to valid work. + #[test] + fn completed_scheduler_stays_invalid() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<1, 1, 1>; + let mut scheduler = Scheduler::for_worker(0, 1).unwrap(); + assert!(scheduler.has_work()); + scheduler.advance_to_next_work(); + assert!(!scheduler.has_work()); + scheduler.advance_to_next_work(); + assert!(!scheduler.has_work()); + } + + /// A persistent scheduler must store only its next linear index and grid stride. + #[test] + fn scheduler_state_is_only_current_and_stride() { + /// Fixed output grid used to check persistent tile ownership and state size. + type Scheduler = StaticPersistentTileScheduler<16, 16, 1>; + assert_eq!(size_of::(), 2 * size_of::()); + assert_eq!(Scheduler::TOTAL_TILES, 256); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/__compiler.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/__compiler.rs new file mode 100644 index 0000000000..39e113532f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/__compiler.rs @@ -0,0 +1,139 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Exact compiler markers for paired two-CTA GEMM. + +use crate::copy::tma::TensorMap; +use core::marker::PhantomData; +use cuda_device::barrier::Barrier; +/// Device-only marker allocating COLUMNS of group-two tensor memory. +/// +/// # Safety +/// The allocating warps and shared address slot must satisfy +/// [`Sm100Tmem::alloc`](super::Sm100Tmem::alloc). +#[inline(never)] +pub unsafe fn sm100_tmem_alloc(token: *mut u32) { + let _ = token; + unreachable!("SM100 CuTe marker requires device compilation") +} +/// Device-only marker releasing a completed group-two tensor-memory allocation. +/// +/// # Safety +/// Both allocating warps must rendezvous after every asynchronous user has +/// finished, as [`Sm100Tmem::dealloc`](super::Sm100Tmem::dealloc) requires. +#[inline(never)] +pub unsafe fn sm100_tmem_dealloc(tmem: u32) { + let _ = tmem; + unreachable!("SM100 CuTe marker requires device compilation") +} +/// Device-only marker accumulating the paired N partitions with shared A-collector reuse. +/// +/// # Safety +/// The elected leader must supply complete live shared tiles and the matching 512-column accumulator. +#[inline(never)] +pub unsafe fn sm100_tiled_mma( + tmem: u32, + a: *const f16, + b0: *const f16, + b1: *const f16, + accumulate: bool, +) { + let _ = (tmem, a, b0, b1, accumulate, PhantomData::<(AL, BL0, BL1)>); + unreachable!("SM100 CuTe marker requires device compilation") +} +/// Device-only marker loading, biasing and converting one FP32 TMEM slice to shared FP16. +/// +/// # Safety +/// All epilogue lanes must supply their valid row-group address and reusable output storage. +#[inline(never)] +pub unsafe fn sm100_tmem_epilogue(tmem: u32, dst: *mut f16, tid: u32, bias: f32) { + let _ = (tmem, dst, tid, bias, PhantomData::); + unreachable!("SM100 CuTe marker requires device compilation") +} +// Every lane of the producer warp participates. CuTe selects the issuing +// lane and performs any required warp synchronization. +/// Device-only marker loading one CTA’s tile and routing completion to the leader. +/// +/// # Safety +/// All producer-warp lanes participate with the matching descriptor, stage and cluster byte expectation. +#[inline(never)] +pub unsafe fn sm100_cluster_tma_load( + desc: *const TensorMap, + dst: *mut f16, + row: u32, + col: u32, + bar: *mut Barrier, + rank: u32, +) { + let _ = (desc, dst, row, col, bar, rank); + unreachable!("SM100 CuTe marker requires device compilation") +} + +/// Device-only marker storing a published shared tile to global memory. +/// +/// # Safety +/// The elected issuer must retain matching source and descriptor storage through read completion. +#[inline(never)] +pub unsafe fn sm100_tma_store( + desc: *const TensorMap, + src: *const f16, + row: u32, + col: u32, +) { + let _ = (desc, src, row, col); + unreachable!("SM100 CuTe marker requires device compilation") +} +/// Define an exact compiler hook for one store-group ownership operation. +macro_rules! store_marker { + ($(#[$doc:meta])* $name:ident) => { + $(#[$doc])* + #[inline(never)] + pub unsafe fn $name() { + unreachable!("SM100 CuTe marker requires device compilation") + } + }; +} +store_marker!( + #[doc = "Device-only hook committing the issuer's pending TMA stores to one group."] + sm100_store_commit +); +store_marker!( + #[doc = "Device-only hook waiting until at most STAGES-1 groups still read shared storage."] + sm100_store_acquire +); +store_marker!( + #[doc = "Device-only hook waiting until every TMA store group stops reading shared storage."] + sm100_store_tail +); +/// Define an exact compiler hook for one full/empty input-stage operation. +macro_rules! pipeline_marker { + ($(#[$doc:meta])* $name:ident ($($arg:ident:$ty:ty),*)) => { + $(#[$doc])* + #[inline(never)] + pub unsafe fn $name($($arg:$ty),*) { + let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") + } + }; +} +pipeline_marker!(#[doc = "Initialize full/empty barriers; one lane per CTA publishes initialization before use."] sm100_pipeline_init(full: *mut Barrier, empty: *mut Barrier)); +pipeline_marker!(#[doc = "Wait for the producer's reusable stage; stage and phase must match the live ring position."] sm100_pipeline_acquire(empty: *mut Barrier, stage: u32, phase: u32)); +pipeline_marker!(#[doc = "Attach TX_BYTES for both CTAs; exactly one elected leader lane calls after acquisition."] sm100_pipeline_expect(full: *mut Barrier, stage: u32)); +pipeline_marker!(#[doc = "Wait for TMA completion and order later MMA; called by the leader's consuming warp."] sm100_pipeline_wait(full: *mut Barrier, stage: u32, phase: u32)); +pipeline_marker!(#[doc = "Release both CTAs' stage after prior MMA completes; called by the elected leader issuer."] sm100_pipeline_release(empty: *mut Barrier, stage: u32)); +/// Define an exact compiler hook for one accumulator ownership handoff. +macro_rules! accumulator_marker { + ($(#[$doc:meta])* $name:ident ($($arg:ident:$ty:ty),*)) => { + $(#[$doc])* + #[inline(never)] + pub unsafe fn $name($($arg:$ty),*) { + let _ = ($($arg),*); unreachable!("SM100 CuTe marker requires device compilation") + } + }; +} +accumulator_marker!(#[doc = "Initialize accumulator handoff barriers from one lane per CTA before cluster publication."] sm100_accumulator_init(empty: *mut Barrier, full: *mut Barrier)); +accumulator_marker!(#[doc = "Wait for all epilogue consumers to release TMEM; the leader's MMA warp uses empty parity."] sm100_accumulator_acquire(empty: *mut Barrier, phase: u32)); +accumulator_marker!(#[doc = "Publish completed accumulation to both CTAs; one elected leader lane commits after final MMA."] sm100_accumulator_commit(full: *mut Barrier)); +accumulator_marker!(#[doc = "Wait for an accumulator tile and order TMEM loads; every consumer warp uses full parity."] sm100_accumulator_wait(full: *mut Barrier, phase: u32)); +accumulator_marker!(#[doc = "Fence completed epilogue reads and elect one arrival per warp; every consumer lane participates."] sm100_accumulator_release(empty: *mut Barrier)); diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mma.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mma.rs new file mode 100644 index 0000000000..7266af1feb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mma.rs @@ -0,0 +1,98 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Paired N-partition FP16 MMA with shared A-collector reuse. + +use super::{__compiler, OperandA, OperandB, Sm100SharedTile}; +use crate::markers::ReifySmem2D; +use core::marker::PhantomData; + +/// Two-CTA FP16 tiled MMA with two N partitions sharing the A collector. +/// +/// The A collector holds an input so two MMA instructions can reuse it: +/// +/// ```text +/// each K=16 step: A x B0 -> first N0 columns (save A) +/// A x B1 -> next N1 columns (reuse, then release A) +/// ``` +/// +/// Each step uses a fresh A input. The complete output is +/// `256 x (N0 + N1)` with FP32 accumulation. +/// `K` must be 64; each N partition is a multiple of 16 in `16..=256`. +/// Per-CTA layouts cover A `(128, 64)`, B0 `(N0 / 2, 64)` and B1 +/// `(N1 / 2, 64)`, in FP16 elements, with the supported B128 swizzle: +/// +/// ``` +/// use cute_rs::{ComposedLayout, RowMajor, Sm100TiledMma, Swizzle}; +/// type Tile = ComposedLayout, 0, RowMajor>; +/// let mma = Sm100TiledMma::, Tile<96>, Tile<80>, 192, 160>::new(); +/// ``` +pub struct Sm100TiledMma( + #[doc = "Carries the three shared operand maps; no device storage is allocated."] + PhantomData<(AL, B0L, B1L)>, +); +impl< + AL: ReifySmem2D, + B0L: ReifySmem2D, + B1L: ReifySmem2D, + const N0: u32, + const N1: u32, + const K: u32, +> Sm100TiledMma +{ + /// Reject unsupported N partitions, K extents and per-CTA operand shapes at instantiation. + const VALID: () = { + assert!(N0 >= 16 && N0 <= 256 && N0.is_multiple_of(16)); + assert!(N1 >= 16 && N1 <= 256 && N1.is_multiple_of(16)); + assert!(N0 + N1 <= 512 && K == 64); + assert!(AL::SHAPE.0 == 128 && AL::SHAPE.1 == K as i64); + assert!(B0L::SHAPE.0 == (N0 / 2) as i64 && B0L::SHAPE.1 == K as i64); + assert!(B1L::SHAPE.0 == (N1 / 2) as i64 && B1L::SHAPE.1 == K as i64); + }; + + /// Check the tile dimensions and create the MMA configuration. + #[inline(always)] + pub const fn new() -> Self { + let () = Self::VALID; + Self(PhantomData) + } + + /// Accumulate one staged K tile into tensor memory. + /// + /// # Safety + /// + /// One elected lane in the leader CTA issues this operation after the + /// input pipeline wait. Both CTAs' tiles must be full and remain unchanged + /// until the asynchronous consumer release completes. `address` names a + /// 512-column allocation; `accumulate=false` starts a new output tile. + #[inline(always)] + pub unsafe fn gemm( + &self, + address: u32, + a: &Sm100SharedTile, + b0: &Sm100SharedTile, + b1: &Sm100SharedTile, + accumulate: bool, + ) { + unsafe { + __compiler::sm100_tiled_mma::( + address, a.base, b0.base, b1.base, accumulate, + ) + } + } +} +impl< + AL: ReifySmem2D, + B0L: ReifySmem2D, + B1L: ReifySmem2D, + const N0: u32, + const N1: u32, + const K: u32, +> Default for Sm100TiledMma +{ + fn default() -> Self { + Self::new() + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mod.rs new file mode 100644 index 0000000000..18a8c4632e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/mod.rs @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 GEMM building blocks for a pair of thread blocks (CTAs). +//! +//! Rust types describe shared-memory layouts and the two output partitions. +//! The compiler turns these calls into CuTe operations, then uses CUTLASS to +//! produce GPU code. Call them only inside device code. +//! +//! ```text +//! global A/B -> TMA -> shared tiles -> MMA -> tensor memory -> output +//! ``` + +/// Operand roles shared with the ctm-rs tcgen05 descriptors: `OperandA` marks +/// the left input and `OperandB` the right input of a matrix product, for +/// shared tiles here and for instruction operands in ctm-rs alike. +pub use ctm_rs::tcgen05::descriptors::{OperandA, OperandB}; + +mod mma; +mod pipeline; +mod shared; +mod tma; +mod tmem; + +pub use mma::Sm100TiledMma; +pub use pipeline::{PipelineTmaStore, PipelineTmaUmma, PipelineUmmaAsync}; +pub use shared::Sm100SharedTile; +pub use tma::Sm100ClusterTmaCopy; +pub use tmem::{Sm100Tmem, Sm100TmemEpilogue}; + +/// Calls recognized by the compiler; never execute these on the host. +#[doc(hidden)] +#[allow(missing_docs)] +pub mod __compiler; +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/pipeline.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/pipeline.rs new file mode 100644 index 0000000000..cd44b21cd5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/pipeline.rs @@ -0,0 +1,194 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Input, accumulator, and output ownership pipelines for paired GEMM. + +use super::__compiler; +use cuda_device::barrier::Barrier; + +/// Two-CTA TMA-to-MMA stage ring, with one completion arrival per stage. +/// +/// `TX_BYTES` is the total number of bytes copied by both CTAs per stage. +/// One lane in the leader CTA sets this expected byte count. MMA completion +/// releases the stage in both CTAs so the next TMA copy can reuse it. +pub struct PipelineTmaUmma { + /// Local array of full barriers tracking cluster-wide input-copy bytes. + full: *mut Barrier, + /// Local array of empty barriers released by group-two MMA completion. + empty: *mut Barrier, +} +impl PipelineTmaUmma { + /// Require at least one stage and a nonzero cluster transaction byte count. + const VALID: () = assert!(STAGES > 0 && TX_BYTES > 0); + /// Attach separate full and empty barrier arrays. + /// + /// # Safety + /// + /// Each pointer addresses `STAGES` aligned shared barriers exclusively + /// used by this pipeline until all consumers and producers finish. + #[inline(always)] + pub const unsafe fn from_raw(full: *mut Barrier, empty: *mut Barrier) -> Self { + let () = Self::VALID; + Self { full, empty } + } + /// Initialize both arrays from one elected lane per CTA. + /// + /// # Safety + /// + /// Publish barrier initialization and synchronize the cluster before use. + #[inline(always)] + pub unsafe fn init(&self) { + unsafe { __compiler::sm100_pipeline_init::(self.full, self.empty) } + } + /// Wait for a reusable stage, with producer phase initially one. + /// + /// # Safety + /// + /// `stage < STAGES`; `phase` must be the producer's current ring phase. + #[inline(always)] + pub unsafe fn producer_acquire(&self, stage: u32, phase: u32) { + unsafe { __compiler::sm100_pipeline_acquire::(self.empty, stage, phase) } + } + /// Attach the complete cluster's expected copy bytes to this full barrier. + /// + /// # Safety + /// + /// Exactly one elected lane of rank zero calls this after acquire. + #[inline(always)] + pub unsafe fn producer_expect_tx(&self, stage: u32) { + unsafe { __compiler::sm100_pipeline_expect::(self.full, stage) } + } + /// Return the stage's local completion barrier for typed TMA copies. + /// + /// # Safety + /// + /// `stage < STAGES`. Copies using this barrier must belong to the acquired + /// stage and phase covered by the leader's one cluster-wide expectation. + #[inline(always)] + pub unsafe fn producer_get_barrier(&self, stage: u32) -> *mut Barrier { + unsafe { self.full.add(stage as usize) } + } + /// Wait for input copies and order subsequent tensor-core operations. + /// + /// # Safety + /// + /// Rank zero's MMA warp calls with its current consumer stage and phase. + #[inline(always)] + pub unsafe fn consumer_wait(&self, stage: u32, phase: u32) { + unsafe { __compiler::sm100_pipeline_wait::(self.full, stage, phase) } + } + /// Signal both CTAs' empty barriers after pending MMA reads finish. + /// + /// # Safety + /// + /// One elected lane in the leader CTA calls after issuing all stage MMAs. + #[inline(always)] + pub unsafe fn consumer_release(&self, stage: u32) { + unsafe { __compiler::sm100_pipeline_release::(self.empty, stage) } + } +} + +/// One accumulator slot shared by the MMA producer and epilogue warps. +pub struct PipelineUmmaAsync { + /// Local barrier released when all epilogue warps finish reading the accumulator. + empty: *mut Barrier, + /// Local barrier signaled when the leader’s MMA work completes. + full: *mut Barrier, +} +impl PipelineUmmaAsync { + /// Require between one and 32 epilogue consumer warps per CTA. + const VALID: () = assert!(CONSUMER_WARPS > 0 && CONSUMER_WARPS <= 32); + /// Attach the shared empty and full barriers. + /// + /// # Safety + /// + /// Both pointers address distinct aligned barriers exclusively owned by + /// this pipeline. Each CTA has `CONSUMER_WARPS` epilogue consumer warps. + #[inline(always)] + pub const unsafe fn from_raw(empty: *mut Barrier, full: *mut Barrier) -> Self { + let () = Self::VALID; + Self { empty, full } + } + /// Initialize from one elected lane in each CTA, before cluster publication. + /// + /// # Safety + /// + /// No pipeline operation may run until initialization is visible cluster-wide. + #[inline(always)] + pub unsafe fn init(&self) { + unsafe { __compiler::sm100_accumulator_init::(self.empty, self.full) } + } + /// Wait for all epilogue warps to release the previous accumulator tile. + /// + /// # Safety + /// + /// Only the leader CTA's MMA warp acquires; its initial phase is one. + #[inline(always)] + pub unsafe fn producer_acquire(&self, phase: u32) { + unsafe { __compiler::sm100_accumulator_acquire::(self.empty, phase) } + } + /// Publish the tile to both CTAs after all pending MMA operations finish. + /// + /// # Safety + /// + /// One elected lane in the leader CTA calls after the final K stage. + #[inline(always)] + pub unsafe fn producer_commit(&self) { + unsafe { __compiler::sm100_accumulator_commit::(self.full) } + } + /// Wait for a full accumulator tile and order tensor-memory loads. + /// + /// # Safety + /// + /// Consumer warps pass their current phase, initially zero. + #[inline(always)] + pub unsafe fn consumer_wait(&self, phase: u32) { + unsafe { __compiler::sm100_accumulator_wait::(self.full, phase) } + } + /// Release this warp's portion of the accumulator back to the leader. + /// + /// # Safety + /// + /// Every lane in every epilogue warp calls after completing its tensor- + /// memory reads. The operation fences all lanes and internally elects + /// one lane to arrive at the leader CTA barrier. + #[inline(always)] + pub unsafe fn consumer_release(&self) { + unsafe { __compiler::sm100_accumulator_release::(self.empty) } + } +} + +/// A ring of shared epilogue tiles tracked by asynchronous TMA store groups. +/// +/// This is the SM100 form of CuTeDSL's `pipeline.PipelineTmaStore`, with a +/// `u32` stage count and the sm100 compiler hooks. The SM120 form with the +/// same methods over a `usize` stage count is [`crate::PipelineTmaStore`]. +#[derive(Clone, Copy, Default)] +pub struct PipelineTmaStore; +impl PipelineTmaStore { + /// Reject an empty ring of shared output buffers. + const VALID: () = assert!(STAGES > 0); + /// Define a ring containing at least one shared output buffer. + #[inline(always)] + pub const fn new() -> Self { + let () = Self::VALID; + Self + } + /// Commit the issuer's pending TMA stores to one copy group. + #[inline(always)] + pub fn producer_commit(&self) { + unsafe { __compiler::sm100_store_commit::() } + } + /// Wait until at most `STAGES-1` store groups still read shared memory. + #[inline(always)] + pub fn producer_acquire(&self) { + unsafe { __compiler::sm100_store_acquire::() } + } + /// Wait until every store group has stopped reading shared memory. + #[inline(always)] + pub fn producer_tail(&self) { + unsafe { __compiler::sm100_store_tail::() } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/shared.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/shared.rs new file mode 100644 index 0000000000..0caa688efe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/shared.rs @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared-memory tiles for paired two-CTA GEMM. + +use crate::markers::ReifySmem2D; +use core::marker::PhantomData; + +/// A non-owning FP16 shared tile with a layout and an operand role. +pub struct Sm100SharedTile { + /// Non-owning shared FP16 allocation base. + pub(super) base: *mut f16, + /// Records the shared layout and A/B operand role. + marker: PhantomData<(Layout, Role)>, +} + +impl Sm100SharedTile { + /// Interpret existing shared storage using `Layout`. + /// + /// # Safety + /// + /// `base` must remain live and contain every element addressed by `Layout`. + /// TMA input bases must be aligned to their full swizzle period; the + /// supported B128 MMA layouts require 1024-byte alignment. + #[inline(always)] + pub const unsafe fn from_raw(base: *mut f16) -> Self { + Self { + base, + marker: PhantomData, + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tests.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tests.rs new file mode 100644 index 0000000000..c9a4cc43ec --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tests.rs @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks the zero-storage paired SM100 MMA configuration."] +use super::*; +use crate::{ComposedLayout, RowMajor, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. +type A = ComposedLayout, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. +type B0 = ComposedLayout, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. +type B1 = ComposedLayout, 0, RowMajor<80, 64>>; +/// The paired SM100 MMA configuration must occupy zero runtime bytes. +#[test] +fn typed_split_plan_has_no_runtime_storage() { + let mma = Sm100TiledMma::::new(); + assert_eq!(core::mem::size_of_val(&mma), 0); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tma.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tma.rs new file mode 100644 index 0000000000..18c41ab22a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tma.rs @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Cluster copies into typed shared tiles for paired GEMM. + +use super::{__compiler, Sm100SharedTile}; +use crate::copy::tma::TensorMap; +use crate::markers::ReifySmem2D; +use core::marker::PhantomData; +use cuda_device::barrier::Barrier; + +/// A cluster TMA copy whose descriptor and destination have the same layout. +pub struct Sm100ClusterTmaCopy( + #[doc = "Carries the destination layout required by the descriptor."] PhantomData, +); +impl Sm100ClusterTmaCopy { + /// Copy an FP16 tile, routing completion to the leader CTA's barrier. + /// + /// Coordinates count elements. Each CTA supplies its rank and receives + /// only its own tile; completion bytes from both CTAs accumulate at rank 0. + /// + /// # Safety + /// + /// All 32 lanes of one producer warp in each CTA call this operation + /// together after acquiring the stage, with identical arguments within + /// that warp. The leader must attach exactly one expectation covering both + /// CTAs' copies to the same stage and phase. The native CuTe copy elects + /// one lane to issue the copy and may synchronize the warp. + /// Calling it from an already elected lane can deadlock when + /// pipeline stages are reused. + /// + /// The descriptor must encode `Layout`, destination storage must be valid + /// and aligned, and `barrier` names the local full barrier for this stage. + #[inline(always)] + pub unsafe fn copy( + desc: *const TensorMap, + dst: &mut Sm100SharedTile, + row: u32, + column: u32, + barrier: *mut Barrier, + rank: u32, + ) { + unsafe { + __compiler::sm100_cluster_tma_load::(desc, dst.base, row, column, barrier, rank) + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tmem.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tmem.rs new file mode 100644 index 0000000000..dce7492183 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/sm100/tmem.rs @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Two-CTA tensor-memory allocation and accumulator epilogues. + +use super::__compiler; +use crate::copy::tma::TensorMap; +use crate::markers::ReifySmem2D; +use core::marker::PhantomData; + +/// Allocate one tensor-memory region collectively in a two-CTA cluster. +/// +/// Tensor memory has 128 rows per CTA. `COLUMNS` is a power of two from 32 +/// through 512. Allocation and deallocation are collective warp operations. +pub struct Sm100Tmem; +impl Sm100Tmem { + /// Restrict allocation columns to supported powers of two from 32 through 512. + const VALID: () = assert!(COLUMNS >= 32 && COLUMNS <= 512 && COLUMNS.is_power_of_two()); + + /// Allocate columns and relinquish the allocation permit. + /// + /// # Safety + /// + /// Every lane of the allocating warp in both CTAs must participate. + /// `token` points to a live aligned shared `u32`. Synchronize readers + /// before loading the address written there. + #[inline(always)] + pub unsafe fn alloc(token: *mut u32) { + let () = Self::VALID; + unsafe { __compiler::sm100_tmem_alloc::(token) } + } + + /// Return a previously allocated region after both CTAs finish using it. + /// + /// # Safety + /// + /// All asynchronous MMA and loads must be complete and the two CTAs must + /// rendezvous before both allocating warps call this operation. + #[inline(always)] + pub unsafe fn dealloc(address: u32) { + let () = Self::VALID; + unsafe { __compiler::sm100_tmem_dealloc::(address) } + } +} + +/// Convert one 128x32 FP32 accumulator slice to an FP16 tile in shared memory. +pub struct Sm100TmemEpilogue( + #[doc = "Carries the shared output layout required by readback and TMA."] PhantomData, +); +impl Sm100TmemEpilogue { + /// Load a lane's 32 accumulators, add its row bias, and convert to FP16. + /// + /// # Safety + /// + /// All lanes of the four epilogue warps participate after accumulator wait. + /// `address` already selects the calling warp's 32 TMEM rows and current + /// 32-column slice. `tid` is the CTA thread index in `0..128`. `dst` has + /// the supported 128x32 layout and is reusable according to the store + /// pipeline. Publish shared writes before issuing the TMA store. + #[inline(always)] + pub unsafe fn store(address: u32, dst: *mut f16, tid: u32, bias: f32) { + unsafe { __compiler::sm100_tmem_epilogue::(address, dst, tid, bias) } + } + + /// Select one warp's 32 TMEM rows and one 32-column accumulator slice. + /// + /// `warp` is in `0..4` and the chosen columns must belong to the tile. + #[inline(always)] + pub const fn slice_address(base: u32, warp: u32, subtile: u32) -> u32 { + (((base >> 16) + warp * 32) << 16) | ((base & 0xffff) + subtile * 32) + } + + /// Start the FP16 tile's asynchronous shared-to-global TMA store. + /// + /// # Safety + /// + /// One elected store issuer calls after all writers publish their shared + /// writes and synchronize. The descriptor encodes `Layout`; coordinates + /// count elements. The shared source must remain live and unchanged until + /// this copy's committed group finishes reading it. A store-pipeline + /// acquire releases the oldest buffer in the ring and can leave this + /// newest group pending; it does not immediately release this source. + #[inline(always)] + pub unsafe fn store_async( + desc: *const TensorMap, + src: *const f16, + row: u32, + column: u32, + ) { + unsafe { __compiler::sm100_tma_store::(desc, src, row, column) } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/static_layout.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/static_layout.rs new file mode 100644 index 0000000000..1bb3abf616 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/static_layout.rs @@ -0,0 +1,173 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Allocation-free evaluation of static CuTe shape/stride markers. +//! +//! The leftmost shape mode changes fastest, as in `cute_layout::Layout`. +//! This bounded, nonnegative-stride subset is suitable for device address maps; +//! the complete algebra and validation remain in `cute-layout`. +use crate::markers::{C, ColMajor, ComposedLayout, Layout, ReifySmem2D, RowMajor, Swizzle}; + +/// A static layout mapping a flat logical coordinate to an element offset. +pub trait StaticLayout { + /// Logical element count. Invalid shapes are rejected at instantiation. + const SIZE: usize; + /// Upper bound on every in-domain physical element offset. + /// + /// Affine layouts have an exact maximum; swizzled layouts may conservatively + /// include unused offsets. The bound fits both `usize` and positive `i64`. + const MAX_OFFSET: usize; + /// Map a flat coordinate in the layout's domain to its element offset + /// without accessing memory (CuTeDSL `cute.crd2idx`). Callers that use + /// the offset for memory access must bound the coordinate. + fn crd2idx(flat: usize) -> usize; +} + +/// Convert a nonnegative layout parameter without truncating it on narrow hosts. +const fn parameter(value: i64) -> usize { + assert!(value >= 0 && value as u128 <= usize::MAX as u128); + value as usize +} + +/// Bound offsets by the signed representation consumed by the shared algebra. +const fn checked_offset_add(a: usize, b: usize) -> usize { + let result = match a.checked_add(b) { + Some(value) => value, + None => panic!("static layout physical offset overflow"), + }; + assert!( + result as u128 <= i64::MAX as u128, + "static layout offset exceeds i64" + ); + result +} + +const fn checked_offset_mul(a: usize, b: usize) -> usize { + let result = match a.checked_mul(b) { + Some(value) => value, + None => panic!("static layout physical offset overflow"), + }; + assert!( + result as u128 <= i64::MAX as u128, + "static layout offset exceeds i64" + ); + result +} + +/// Swizzling changes only the target bit field. Fill all bits up through that +/// field to bound every offset below `maximum`, including holes in the layout. +const fn swizzled_bound(maximum: usize, bits: u32, base: u32, shift: i32) -> usize { + let _ = cute_layout::Swizzle::new(bits, base, shift); + if bits == 0 { + return maximum; + } + let target_end = base + bits + if shift < 0 { (-shift) as u32 } else { 0 }; + let lower_mask = (1_u64 << target_end) - 1; + let result = maximum as u64 | lower_mask; + assert!(result as u128 <= usize::MAX as u128 && result <= i64::MAX as u64); + result as usize +} + +impl StaticLayout for Layout, C> { + const SIZE: usize = { + let _ = Self::MAX_OFFSET; + parameter(N) + }; + const MAX_OFFSET: usize = { + assert!(N > 0 && S >= 0); + checked_offset_mul(parameter(N - 1), parameter(S)) + }; + #[inline(always)] + fn crd2idx(flat: usize) -> usize { + let _ = Self::SIZE; + flat * S as usize + } +} +impl StaticLayout for Layout<(A, B), (SA, SB)> +where + Layout: StaticLayout, + Layout: StaticLayout, +{ + const SIZE: usize = { + let _ = Self::MAX_OFFSET; + match Layout::::SIZE.checked_mul(Layout::::SIZE) { + Some(n) => n, + None => panic!("static layout extent overflow"), + } + }; + const MAX_OFFSET: usize = + checked_offset_add(Layout::::MAX_OFFSET, Layout::::MAX_OFFSET); + #[inline(always)] + fn crd2idx(flat: usize) -> usize { + let _ = Self::SIZE; + Layout::::crd2idx(flat % Layout::::SIZE) + + Layout::::crd2idx(flat / Layout::::SIZE) + } +} +impl StaticLayout for RowMajor { + const SIZE: usize = , C), (C, C<1>)>>::SIZE; + const MAX_OFFSET: usize = , C), (C, C<1>)>>::MAX_OFFSET; + #[inline(always)] + fn crd2idx(flat: usize) -> usize { + , C), (C, C<1>)>>::crd2idx(flat) + } +} +impl StaticLayout for ColMajor { + const SIZE: usize = , C), (C<1>, C)>>::SIZE; + const MAX_OFFSET: usize = , C), (C<1>, C)>>::MAX_OFFSET; + #[inline(always)] + fn crd2idx(flat: usize) -> usize { + , C), (C<1>, C)>>::crd2idx(flat) + } +} +impl StaticLayout + for ComposedLayout, O, Outer> +{ + const SIZE: usize = { + let _ = Self::MAX_OFFSET; + Outer::SIZE + }; + const MAX_OFFSET: usize = + swizzled_bound(checked_offset_add(Outer::MAX_OFFSET, parameter(O)), B, M, S); + #[inline(always)] + fn crd2idx(flat: usize) -> usize { + let _ = Self::SIZE; + cute_layout::Swizzle::new(B, M, S).apply((Outer::crd2idx(flat) + O as usize) as i64) + as usize + } +} + +/// Map a `(row, column)` coordinate through a typed two-dimensional +/// shared-memory layout, including its swizzle (CuTeDSL `cute.crd2idx` on a +/// composed layout). +/// +/// Offsets and strides count elements. The static offset is added before the +/// swizzle, matching `cute_layout::ComposedLayout`. This computes an offset +/// only; callers that dereference it must establish coordinate bounds and +/// storage lifetime. +#[inline(always)] +pub fn crd2idx_2d(row: usize, column: usize) -> usize { + const { + let _ = max_offset_2d::(); + } + let plain = row * L::STRIDE.0 as usize + column * L::STRIDE.1 as usize + L::OFFSET as usize; + cute_layout::Swizzle::new(L::NUM_BITS, L::NUM_BASE, L::NUM_SHIFT).apply(plain as i64) as usize +} + +/// Bound the physical offsets of a reified two-dimensional layout at compile time. +/// +/// Rejects invalid dimensions and offsets that cannot be represented by both +/// `usize` and the signed CuTe algebra. Swizzles can conservatively include +/// holes in this bound. `crd2idx_2d` instantiates this check before device code. +pub const fn max_offset_2d() -> usize { + assert!(L::SHAPE.0 > 0 && L::SHAPE.1 > 0); + assert!(L::STRIDE.0 >= 0 && L::STRIDE.1 >= 0 && L::OFFSET >= 0); + let row_max = checked_offset_mul(parameter(L::SHAPE.0 - 1), parameter(L::STRIDE.0)); + let column_max = checked_offset_mul(parameter(L::SHAPE.1 - 1), parameter(L::STRIDE.1)); + let maximum = checked_offset_add( + checked_offset_add(row_max, column_max), + parameter(L::OFFSET), + ); + swizzled_bound(maximum, L::NUM_BITS, L::NUM_BASE, L::NUM_SHIFT) +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/mod.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/mod.rs new file mode 100644 index 0000000000..7a21f5ab3a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/mod.rs @@ -0,0 +1,924 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed views of existing tensor memory. +//! +//! A [`Tensor`] combines three things: +//! +//! ```text +//! storage: where the bits are +//! layout: how coordinates find those bits +//! element: what the bits mean +//! │ +//! ▼ +//! Tensor +//! ``` +//! +//! A tensor view does not allocate, copy, or load memory. Operations such as +//! [`Tensor::zipped_divide`] and [`Tensor::slice`] only change how the same +//! storage is viewed. [`Tensor::load`] is the first operation that reads it. +//! +//! [`partition`] describes strided views, subranges and cyclic thread +//! partitions while retaining their element and mutable-ownership contracts. + +pub mod partition; + +use core::marker::PhantomData; +use core::ops::{Add, Div, Mul, Neg, Sub}; + +use cuda_device::DisjointSlice; + +use crate::copy::vector::{assume, load_tile, store_tile}; + +/// Says which Rust storage type may hold one or more logical elements. +/// +/// For example, `Float4E2M1FN: TensorElement` means an `u8` may carry two packed +/// FP4 values. A missing implementation rejects an invalid pair at compile +/// time, such as E2M1 data stored in `f32`. +pub trait TensorElement {} + +// Plain f16 and f32 values use the same type for meaning and storage. +impl TensorElement for f16 {} +impl TensorElement for f32 {} +impl TensorElement for u8 {} +impl TensorElement for u16 {} +impl TensorElement for u32 {} +impl TensorElement for u64 {} +impl TensorElement for usize {} +impl TensorElement for i8 {} +impl TensorElement for i16 {} +impl TensorElement for i32 {} +impl TensorElement for i64 {} +impl TensorElement for isize {} + +/// One line of adjacent elements. +/// +/// ```text +/// index: 0 1 2 3 +/// memory: [a] [b] [c] [d] +/// ``` +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Contiguous1D { + /// Number of logical elements in the contiguous row. + len: usize, +} + +impl Contiguous1D { + /// Create a layout containing `len` adjacent elements. + #[must_use] + #[inline(always)] + pub const fn new(len: usize) -> Self { + Self { len } + } + + /// Return the number of elements. + #[must_use] + #[inline(always)] + pub const fn len(self) -> usize { + self.len + } + + /// Return `true` when the length is zero. + #[must_use] + #[inline(always)] + pub const fn is_empty(self) -> bool { + self.len == 0 + } +} + +/// A one-dimensional tensor split into fixed-size tiles. +/// +/// For `TILE = 4`: +/// +/// ```text +/// original: [0 1 2 3 | 4 5 6 7 | 8 9] +/// tile: 0 1 2 +/// inside: 0 1 2 3 0 1 2 3 0 1 +/// ``` +/// +/// The final tile may be shorter than `TILE`. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Zipped1D { + /// Full tensor length, used to retain a possibly partial final tile. + len: usize, +} + +impl Zipped1D { + /// Return the number of tiles, including a short final tile. + #[must_use] + #[inline(always)] + pub const fn tile_count(self) -> usize { + const { assert!(TILE > 0) }; + self.len.div_ceil(TILE) + } +} + +/// One selected tile from a [`Zipped1D`] view. +/// +/// It stores the tile's first index and the full tensor length. It does not +/// contain or load the tile values. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Tile1D { + /// Absolute first element index of the selected tile. + base: usize, + /// Full tensor length used to clip the selected tile. + len: usize, +} + +impl Tile1D { + /// Return the tile's first index in the original tensor. + #[must_use] + #[inline(always)] + pub const fn base(self) -> usize { + self.base + } + + /// Return `true` when all `TILE` positions are in bounds. + #[must_use] + #[inline(always)] + pub const fn is_full(self) -> bool { + self.valid_len() == TILE + } + + /// Return the number of positions that are in bounds. + #[must_use] + #[inline(always)] + pub const fn valid_len(self) -> usize { + if self.base >= self.len { + 0 + } else { + let remaining = self.len - self.base; + if remaining < TILE { remaining } else { TILE } + } + } +} + +/// The `N` values from one tile, held in thread-local registers. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RegisterTile( + #[doc = "Per-thread register values in increasing tile-coordinate order."] [T; N], +); + +impl RegisterTile { + /// Construct a register tile in increasing tile-coordinate order. + #[must_use] + #[inline(always)] + pub const fn from_array(values: [T; N]) -> Self { + Self(values) + } + + /// Construct registers by visiting each coordinate exactly once in order. + /// + /// Empty tiles never call the producer. Nonempty tiles initialize all + /// storage with the first produced value before replacing the remaining + /// coordinates. No default value or uninitialized storage is required. + /// The explicit fixed-trip loop is fully unrolled during device compilation. + #[must_use] + #[inline(always)] + pub fn from_fn(mut f: impl FnMut(usize) -> T) -> Self { + if N == 0 { + return Self(core::array::from_fn(f)); + } + let first = f(0); + let mut values = [first; N]; + let mut i = 1; + while i < N { + ctm_rs::thread::__unroll_config::<0>(); + values[i] = f(i); + i += 1; + } + Self(values) + } + + /// Transform each register value, preserving the tile's element count. + /// + /// Calls `f` exactly once per element in increasing order, and never for an + /// empty tile. A different output element type requires an explicit + /// conversion in `f`; this method performs no implicit numeric conversion. + #[must_use] + #[inline(always)] + pub fn map(self, mut f: impl FnMut(T) -> U) -> RegisterTile { + RegisterTile::from_fn(|i| f(self.0[i])) + } + + /// Transform corresponding values of two equally sized register tiles. + /// + /// The shared `N` rejects different element counts at compile time. Calls `f` + /// exactly once per element in increasing order, and never for empty tiles. + /// Input and output types can differ, with conversions expressed in `f`. + #[must_use] + #[inline(always)] + pub fn zip_map( + self, + rhs: RegisterTile, + mut f: impl FnMut(T, U) -> V, + ) -> RegisterTile { + RegisterTile::from_fn(|i| f(self.0[i], rhs.0[i])) + } + + /// Reduce this thread's registers in increasing element order. + /// + /// The initial value chooses the accumulation type explicitly. Calls `f` + /// exactly once per element; an empty tile returns `initial`. This does + /// not communicate between threads or reassociate floating-point sums. + #[must_use] + #[inline(always)] + pub fn fold(self, initial: U, mut f: impl FnMut(U, T) -> U) -> U { + let mut accumulator = initial; + let mut i = 0; + while i < N { + ctm_rs::thread::__unroll_config::<0>(); + accumulator = f(accumulator, self.0[i]); + i += 1; + } + accumulator + } + + /// Gather register coordinates into another fixed-size register tile. + /// + /// Coordinates are checked against `N`. Repeated coordinates implement + /// broadcasts; reordered coordinates implement register permutations. + #[must_use] + #[inline(always)] + pub fn gather(self, coordinates: [usize; M]) -> RegisterTile { + RegisterTile::from_fn(|i| self.0[coordinates[i]]) + } + + /// Borrow the values as a fixed-size array. + #[must_use] + #[inline(always)] + pub const fn as_array(&self) -> &[T; N] { + &self.0 + } + + /// Consume the tile and return its fixed-size array. + #[must_use] + #[inline(always)] + pub fn into_array(self) -> [T; N] { + self.0 + } +} + +/// Elementwise sum of two equally sized register tiles (CuTeDSL `TensorSSA + TensorSSA`). +impl, const N: usize> Add for RegisterTile { + type Output = Self; + + #[inline(always)] + fn add(mut self, rhs: Self) -> Self::Output { + let mut i = 0; + while i < N { + self.0[i] = self.0[i] + rhs.0[i]; + i += 1; + } + self + } +} + +/// Elementwise difference of two equally sized register tiles (CuTeDSL `TensorSSA - TensorSSA`). +impl, const N: usize> Sub for RegisterTile { + type Output = Self; + + #[inline(always)] + fn sub(mut self, rhs: Self) -> Self::Output { + let mut i = 0; + while i < N { + self.0[i] = self.0[i] - rhs.0[i]; + i += 1; + } + self + } +} + +/// Elementwise product of two equally sized register tiles (CuTeDSL `TensorSSA * TensorSSA`). +impl, const N: usize> Mul for RegisterTile { + type Output = Self; + + #[inline(always)] + fn mul(mut self, rhs: Self) -> Self::Output { + let mut i = 0; + while i < N { + self.0[i] = self.0[i] * rhs.0[i]; + i += 1; + } + self + } +} + +/// Elementwise quotient of two equally sized register tiles (CuTeDSL `TensorSSA / TensorSSA` +/// for floating-point lanes). +/// +/// Each lane divides with the element type's own `/`: integer lanes truncate +/// toward zero and panic on a zero divisor like the scalar operator does. +/// CuTeDSL differs on integer vectors: its `/` converts both operands to f32 +/// and divides in floating point, and its `//` floors. Parity holds for +/// floating-point lanes only. +impl, const N: usize> Div for RegisterTile { + type Output = Self; + + #[inline(always)] + fn div(mut self, rhs: Self) -> Self::Output { + let mut i = 0; + while i < N { + self.0[i] = self.0[i] / rhs.0[i]; + i += 1; + } + self + } +} + +/// Elementwise negation of a register tile (CuTeDSL `-TensorSSA`). +/// +/// Each lane is negated with the element type's own `-`, so a `+0.0` lane +/// becomes `-0.0`; CuTeDSL computes `0 - x`, which leaves that lane `+0.0`. +impl, const N: usize> Neg for RegisterTile { + type Output = Self; + + #[inline(always)] + fn neg(mut self) -> Self::Output { + let mut i = 0; + while i < N { + self.0[i] = -self.0[i]; + i += 1; + } + self + } +} + +/// A read-only view of storage with a typed element format and layout. +/// +/// The view owns no memory and performs no allocation: +/// +/// ```text +/// Tensor<'a, Element, Layout, Storage> +/// │ │ │ +/// │ │ └─ Rust type holding the bits +/// │ └────────── coordinate-to-offset rule +/// └─────────────────── meaning of the bits +/// ``` +/// +/// For ordinary values, `Storage` defaults to `Element`, so +/// `Tensor` is enough. Packed data names both types; for example, +/// `Tensor` means two logical FP4 values share each byte. +/// [`TensorElement`] checks that the pair is allowed. +/// +/// Copying this value copies only the view. The underlying storage is still +/// borrowed. This view is used inside Rust kernel code; it is not passed +/// directly as a raw kernel argument. +#[derive(Clone, Copy, Debug)] +pub struct Tensor<'a, Element, Layout, Storage = Element> +where + Element: TensorElement, + Storage: Copy, +{ + /// Borrowed storage carrying the logical elements. + pub(crate) storage: &'a [Storage], + /// Coordinate map or tile position attached to the storage. + pub(crate) layout: Layout, + /// Logical element meaning, which may differ from the storage type. + element: PhantomData, +} + +impl<'a, Element: TensorElement, Layout, Storage: Copy> + Tensor<'a, Element, Layout, Storage> +{ + /// View `storage` through `layout` without copying or loading it. + #[must_use] + #[inline(always)] + pub const fn from_storage(storage: &'a [Storage], layout: Layout) -> Self { + Self { + storage, + layout, + element: PhantomData, + } + } + + /// Return the slice that holds the physical bits. + #[must_use] + #[inline(always)] + pub const fn storage(&self) -> &'a [Storage] { + self.storage + } + + /// Return the layout value stored by this view. + #[must_use] + #[inline(always)] + pub const fn layout(&self) -> Layout + where + Layout: Copy, + { + self.layout + } +} + +impl<'a, T: Copy + TensorElement> Tensor<'a, T, Contiguous1D, T> { + /// View a slice as one contiguous row. + #[must_use] + #[inline(always)] + pub const fn from_slice(storage: &'a [T]) -> Self { + make_tensor_read(storage) + } + + /// Split the row into `(position_in_tile, tile_number)` coordinates. + /// + /// ```text + /// [0 1 2 3 | 4 5 6 7] ── zipped_divide::<4>() ──► 2 tiles of 4 + /// ``` + /// + /// This only changes the view. It does not read memory. + #[must_use] + #[inline(always)] + pub const fn zipped_divide(self) -> Tensor<'a, T, Zipped1D, T> { + zipped_divide_read::(self) + } + + /// Select tile number `tile` of the row split into `TILE`-element tiles. + /// + /// Equivalent to `zipped_divide::().slice(tile)`: CuTe's + /// `local_tile` for a one-dimensional tiler and a single tile + /// coordinate. This only changes the view. It does not read memory. + #[must_use] + #[inline(always)] + pub const fn local_tile( + self, + tile: usize, + ) -> Tensor<'a, T, Tile1D, T> { + self.zipped_divide::().slice(tile) + } +} + +impl<'a, T: Copy + TensorElement, const TILE: usize> Tensor<'a, T, Zipped1D, T> { + /// Select tile number `tile`. + /// + /// This only computes the tile's starting index. Call [`Tensor::load`] to + /// read its values. + #[must_use] + #[inline(always)] + pub const fn slice(self, tile: usize) -> Tensor<'a, T, Tile1D, T> { + slice_read::(self, tile) + } +} + +impl<'a, T: Copy + TensorElement, const TILE: usize> Tensor<'a, T, Tile1D, T> { + /// Return `true` when every position in this tile is in bounds. + /// + /// This is the tensor-level form of [`Tile1D::is_full`]. Keeping the + /// question on the view makes the complete CuTe flow visible to the + /// compiler: + /// + /// ```text + /// Tensor -> zipped_divide -> slice -> is_full + /// ``` + #[must_use] + #[inline(always)] + pub const fn is_full(self) -> bool { + tensor_tile_is_full::(self) + } + + /// Return this tile's first index in the original tensor. + /// + /// This is useful for a short final tile, where the kernel falls back to + /// scalar loads and stores. + #[must_use] + #[inline(always)] + pub const fn base(self) -> usize { + tensor_tile_base::(self) + } + + /// Load the full tile into one [`RegisterTile`] using one vector load. + /// + /// # Safety + /// + /// - [`Tile1D::is_full`] must be `true`. + /// - The tile's first address must be aligned to + /// `TILE * size_of::()` bytes. + /// - That total byte size must be a supported power of two from 4 through + /// 16. + #[must_use] + #[inline(always)] + pub unsafe fn load(self) -> RegisterTile { + unsafe { tensor_load_tile::(self) } + } +} + +/// A writable view of storage with a typed element format and layout. +/// +/// This is separate from [`Tensor`] so a read-only slice can never gain a +/// store method. `TensorMut` is not `Copy` or `Clone`: there should be one +/// owner of the writable view. +/// +/// GPU threads may still calculate overlapping indices, so store methods are +/// unsafe and state the non-overlap rules. This view is used inside Rust +/// kernel code; it is not passed directly as a raw kernel argument. +#[derive(Debug)] +pub struct TensorMut<'a, Element, Layout, Storage = Element> +where + Element: TensorElement, + Storage: Copy, +{ + /// Non-owning mutable storage base. + storage: *mut Storage, + /// Storage element count available from the base. + storage_len: usize, + /// Coordinate map or selected tile position. + layout: Layout, + /// Logical element meaning preserved across view transformations. + element: PhantomData, + /// Retains the exclusive storage borrow without manufacturing a mutable slice. + borrow: PhantomData<&'a mut [Storage]>, +} + +impl<'a, Element: TensorElement, Layout, Storage: Copy> + TensorMut<'a, Element, Layout, Storage> +{ + /// Build a writable tensor view from a raw pointer, length, and layout. + /// + /// # Safety + /// + /// - `storage` must remain valid and correctly aligned for `storage_len` + /// elements throughout `'a`. + /// - Every offset produced by `layout` must be less than `storage_len`. + /// - No other live reference may read or write the same memory in a way + /// that breaks Rust's exclusive-access rules. + #[must_use] + #[inline(always)] + pub const unsafe fn from_raw_parts( + storage: *mut Storage, + storage_len: usize, + layout: Layout, + ) -> Self { + Self { + storage, + storage_len, + layout, + element: PhantomData, + borrow: PhantomData, + } + } + + /// Return the layout value stored by this view. + #[must_use] + #[inline(always)] + pub const fn layout(&self) -> Layout + where + Layout: Copy, + { + self.layout + } +} + +impl<'a, T: Copy + TensorElement> TensorMut<'a, T, Contiguous1D, T> { + /// View a kernel output slice as one contiguous row. + #[must_use] + #[inline(always)] + pub fn from_disjoint_slice(storage: &'a mut DisjointSlice<'_, T>) -> Self { + let len = storage.len(); + unsafe { make_tensor_write(storage.as_mut_ptr(), len) } + } + + /// Write one value at its index in the original contiguous tensor. + /// + /// # Safety + /// + /// - `index` must be less than the tensor length. + /// - No concurrently running thread may access the same element in a + /// conflicting way. + #[inline(always)] + pub unsafe fn store_at(&mut self, index: usize, value: T) { + unsafe { *self.storage.add(index) = value }; + } + + /// Split the row into `(position_in_tile, tile_number)` coordinates. + /// + /// This only changes the view. It does not write memory. + #[must_use] + #[inline(always)] + pub fn zipped_divide(self) -> TensorMut<'a, T, Zipped1D, T> { + zipped_divide_write::(self) + } + + /// Select writable tile number `tile` of the row split into `TILE`-element tiles. + /// + /// Equivalent to `zipped_divide::().slice(tile)`: CuTe's + /// `local_tile` for a one-dimensional tiler and a single tile + /// coordinate. This only changes the view. It does not write memory. + #[must_use] + #[inline(always)] + pub fn local_tile(self, tile: usize) -> TensorMut<'a, T, Tile1D, T> { + self.zipped_divide::().slice(tile) + } +} + +impl<'a, T: Copy + TensorElement, const TILE: usize> TensorMut<'a, T, Zipped1D, T> { + /// Select writable tile number `tile`. + /// + /// This only computes the tile's starting index. Call [`TensorMut::store`] + /// to write it. + #[must_use] + #[inline(always)] + pub fn slice(self, tile: usize) -> TensorMut<'a, T, Tile1D, T> { + slice_write::(self, tile) + } +} + +/// Compiler boundary for a read-only contiguous tensor view. +/// +/// This stays public only so device compilation can identify the exact +/// definition across crate boundaries. Use [`Tensor::from_slice`] in kernels. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const fn make_tensor_read<'a, T: Copy + TensorElement>( + storage: &'a [T], +) -> Tensor<'a, T, Contiguous1D, T> { + Tensor::from_storage(storage, Contiguous1D::new(storage.len())) +} + +/// Compiler boundary for a writable contiguous tensor view. +/// +/// Use [`TensorMut::from_disjoint_slice`] instead. The safety rules are the +/// same as [`TensorMut::from_raw_parts`]. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const unsafe fn make_tensor_write<'a, T: Copy + TensorElement>( + storage: *mut T, + storage_len: usize, +) -> TensorMut<'a, T, Contiguous1D, T> { + unsafe { TensorMut::from_raw_parts(storage, storage_len, Contiguous1D::new(storage_len)) } +} + +/// Compiler boundary for splitting a read-only row into fixed-size tiles. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const fn zipped_divide_read<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: Tensor<'a, T, Contiguous1D, T>, +) -> Tensor<'a, T, Zipped1D, T> { + const { assert!(TILE > 0) }; + Tensor::from_storage( + tensor.storage, + Zipped1D { + len: tensor.layout.len, + }, + ) +} + +/// Compiler boundary for splitting a writable row into fixed-size tiles. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub fn zipped_divide_write<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: TensorMut<'a, T, Contiguous1D, T>, +) -> TensorMut<'a, T, Zipped1D, T> { + const { assert!(TILE > 0) }; + TensorMut { + storage: tensor.storage, + storage_len: tensor.storage_len, + layout: Zipped1D { + len: tensor.layout.len, + }, + element: PhantomData, + borrow: PhantomData, + } +} + +/// Compiler boundary for selecting one read-only tile. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const fn slice_read<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: Tensor<'a, T, Zipped1D, T>, + tile: usize, +) -> Tensor<'a, T, Tile1D, T> { + Tensor::from_storage( + tensor.storage, + Tile1D { + base: tile.saturating_mul(TILE), + len: tensor.layout.len, + }, + ) +} + +/// Compiler boundary for selecting one writable tile. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub fn slice_write<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: TensorMut<'a, T, Zipped1D, T>, + tile: usize, +) -> TensorMut<'a, T, Tile1D, T> { + TensorMut { + storage: tensor.storage, + storage_len: tensor.storage_len, + layout: Tile1D { + base: tile.saturating_mul(TILE), + len: tensor.layout.len, + }, + element: PhantomData, + borrow: PhantomData, + } +} + +impl<'a, T: Copy + TensorElement, const TILE: usize> TensorMut<'a, T, Tile1D, T> { + /// Write one value at its absolute index in the original tensor. + /// + /// Use this for a short final tile; it is the same scalar store verb as + /// `store_at` on an undivided row. Use [`TensorMut::store`] for a full + /// tile so all values are written by one vector store. + /// + /// # Safety + /// + /// - `index` must be between this tile's [`Tile1D::base`] and the end of + /// its valid range. + /// - No concurrently running thread may access the same element in a + /// conflicting way. + #[inline(always)] + pub unsafe fn store_at(&mut self, index: usize, value: T) { + unsafe { tensor_store_element_abs::(self, index, value) }; + } + + /// Write a complete [`RegisterTile`] using one vector store. + /// + /// # Safety + /// + /// - [`Tile1D::is_full`] must be `true`. + /// - The tile's first address must meet the alignment required by + /// [`crate::store_tile`]. + /// - `TILE * size_of::()` must be a supported power of two from 4 + /// through 16. + /// - No concurrently running thread may access the same elements in a + /// conflicting way. + #[inline(always)] + pub unsafe fn store(self, values: RegisterTile) { + unsafe { tensor_store_tile::(self, values) }; + } +} + +/// Compiler boundary for asking whether a selected tile is complete. +/// +/// This stays public only so device compilation can identify the exact +/// definition across crate boundaries. Use [`Tensor::is_full`] in kernels. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const fn tensor_tile_is_full<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: Tensor<'a, T, Tile1D, T>, +) -> bool { + tensor.layout.is_full() +} + +/// Compiler boundary for reading a selected tile's absolute base index. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub const fn tensor_tile_base<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: Tensor<'a, T, Tile1D, T>, +) -> usize { + tensor.layout.base() +} + +/// Compiler boundary for loading one complete selected tile. +/// +/// # Safety +/// +/// The rules are the same as [`Tensor::load`]. +#[doc(hidden)] +#[must_use] +#[inline(never)] +pub unsafe fn tensor_load_tile<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: Tensor<'a, T, Tile1D, T>, +) -> RegisterTile { + let base = unsafe { assume::(tensor.layout.base) }; + RegisterTile(unsafe { load_tile::(tensor.storage, base) }) +} + +/// Compiler boundary for storing one complete selected tile. +/// +/// # Safety +/// +/// The rules are the same as [`TensorMut::store`]. +#[doc(hidden)] +#[inline(never)] +pub unsafe fn tensor_store_tile<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: TensorMut<'a, T, Tile1D, T>, + values: RegisterTile, +) { + let base = unsafe { assume::(tensor.layout.base) }; + unsafe { store_tile::(tensor.storage, base, values.as_array()) }; +} + +/// Compiler boundary for one scalar store in a short final tile. +/// +/// # Safety +/// +/// The rules are the same as [`TensorMut::store_at`] on a selected tile. +#[doc(hidden)] +#[inline(never)] +pub unsafe fn tensor_store_element_abs<'a, T: Copy + TensorElement, const TILE: usize>( + tensor: &mut TensorMut<'a, T, Tile1D, T>, + index: usize, + value: T, +) { + unsafe { *tensor.storage.add(index) = value }; +} + +/// Regression tests for this module’s layout, representation and rejection contracts. +#[cfg(test)] +mod tests { + use super::*; + + /// Distinct logical encoding used to test a tensor backed by u8 storage. + #[derive(Clone, Copy)] + struct TestElement; + + impl TensorElement for TestElement {} + + /// A logical-element view must borrow the original slice without allocating or copying it. + #[test] + fn tensor_view_borrows_storage_without_allocating() { + let storage = [0u32; 8]; + let tensor = Tensor::::from_storage(&storage, (2usize, 4usize)); + + assert_eq!(tensor.storage().as_ptr(), storage.as_ptr()); + assert_eq!(tensor.storage().len(), storage.len()); + assert_eq!(tensor.layout(), (2, 4)); + } + + /// Tiling and slicing must preserve the borrowed storage and expose a partial final tile. + #[test] + fn zipped_divide_and_slice_are_views_over_the_same_storage() { + let storage = [0.0f32; 10]; + let tensor: Tensor<'_, f32, Contiguous1D> = Tensor::from_slice(&storage); + let tiled = tensor.zipped_divide::<4>(); + + assert_eq!(tiled.layout().tile_count(), 3); + + let middle = tiled.slice(1); + assert_eq!(middle.storage().as_ptr(), storage.as_ptr()); + assert_eq!(middle.layout().base(), 4); + assert!(middle.layout().is_full()); + + let tail = tiled.slice(2); + assert_eq!(tail.layout().base(), 8); + assert!(!tail.layout().is_full()); + } + + /// The fused tile selection must agree with the two-step chain on a full and a ragged tile. + #[test] + fn local_tile_equals_zipped_divide_then_slice() { + let storage = [0.0f32; 10]; + let tensor: Tensor<'_, f32, Contiguous1D> = Tensor::from_slice(&storage); + let tiled = tensor.zipped_divide::<4>(); + + // Tile 1 is full; tile 2 is the two-element tail. + for tile in [1, 2] { + let direct = tensor.local_tile::<4>(tile); + let chained = tiled.slice(tile); + assert_eq!(direct.storage().as_ptr(), chained.storage().as_ptr()); + assert_eq!(direct.base(), chained.base()); + assert_eq!(direct.is_full(), chained.is_full()); + } + assert!(tensor.local_tile::<4>(1).is_full()); + assert!(!tensor.local_tile::<4>(2).is_full()); + + let mut output = [0.0f32; 10]; + for tile in [1, 2] { + // Each writable view is built, queried and dropped before the + // next one, so no two views over `output` are alive together. + let direct = unsafe { make_tensor_write(output.as_mut_ptr(), output.len()) } + .local_tile::<4>(tile) + .layout(); + let chained = unsafe { make_tensor_write(output.as_mut_ptr(), output.len()) } + .zipped_divide::<4>() + .slice(tile) + .layout(); + assert_eq!(direct.base(), chained.base()); + assert_eq!(direct.is_full(), chained.is_full()); + assert_eq!(direct.is_full(), tile == 1); + } + } + + /// Register-tile addition must preserve element order and produce per-element sums. + #[test] + fn register_tiles_add_elementwise() { + let a = RegisterTile([1.0f32, 2.0, 3.0, 4.0]); + let b = RegisterTile([10.0f32, 20.0, 30.0, 40.0]); + + assert_eq!((a + b).into_array(), [11.0, 22.0, 33.0, 44.0]); + } + + /// Mutable tensor markers and lifetime tracking must add no bytes beyond pointer, length and layout. + #[test] + fn mutable_tensor_is_not_larger_than_its_runtime_fields() { + assert_eq!( + core::mem::size_of::>(), + 3 * core::mem::size_of::() + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/partition.rs b/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/partition.rs new file mode 100644 index 0000000000..14ff2f0cd9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/src/tensor/partition.rs @@ -0,0 +1,491 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Borrowed tensor storage with composable, typed cyclic coordinate layouts. +//! +//! `local_partition::

` retains `Cyclic` in the result type. Subranges +//! retain `Subrange`; nested operations compose logical coordinates +//! before the original storage layout maps them to physical offsets. No +//! transformed pointer or erased multiplied stride defines partition semantics. +//! +//! Raw strided views may have uninitialized holes between logical elements. +//! Unlike a contiguous tensor slice, their storage proof covers only the +//! coordinates described by the layout; no reference to those holes is formed. +use crate::tensor::{RegisterTile, TensorElement}; +use core::marker::PhantomData; +use ctm_rs::memory::DisjointSlice; + +mod sealed { + pub trait Sealed {} +} + +/// A library-defined logical coordinate map bound to a borrowed tensor. +/// +/// The sealed implementations preserve the originating storage proof through +/// each transformation. An arbitrary external map cannot make safe reads reach +/// an unvalidated address. +pub trait IndexLayout: sealed::Sealed + Copy { + /// Number of logical elements selected by the layout. + fn len(&self) -> usize; + /// Whether the layout selects no elements. + #[inline(always)] + fn is_empty(&self) -> bool { + self.len() == 0 + } + /// Map an in-domain logical coordinate to an originating storage offset + /// (CuTeDSL `cute.crd2idx`). This computes an integer only; it does not + /// read or write storage. + fn crd2idx(&self, index: usize) -> usize; +} + +/// One-dimensional shape and positive physical element stride. +#[derive(Clone, Copy)] +pub struct Strided1D { + len: usize, + stride: usize, +} +impl sealed::Sealed for Strided1D {} +impl IndexLayout for Strided1D { + #[inline(always)] + fn len(&self) -> usize { + self.len + } + #[inline(always)] + fn crd2idx(&self, index: usize) -> usize { + index * self.stride + } +} + +/// A contiguous logical interval of a parent coordinate map. +#[derive(Clone, Copy)] +pub struct Subrange { + parent: Parent, + start: usize, + len: usize, +} +impl sealed::Sealed for Subrange {} +impl IndexLayout for Subrange { + #[inline(always)] + fn len(&self) -> usize { + self.len + } + #[inline(always)] + fn crd2idx(&self, index: usize) -> usize { + self.parent.crd2idx(self.start + index) + } +} + +/// One of `PARTS` cyclic partitions of a parent coordinate map. +/// +/// Logical `i` maps to parent coordinate `part + i * PARTS`. The partition +/// count remains in the type even after further slicing or partitioning. +#[derive(Clone, Copy)] +pub struct Cyclic { + parent: Parent, + part: usize, + len: usize, +} +impl sealed::Sealed for Cyclic {} +impl IndexLayout for Cyclic { + #[inline(always)] + fn len(&self) -> usize { + self.len + } + #[inline(always)] + fn crd2idx(&self, index: usize) -> usize { + self.parent.crd2idx(self.part + index * PARTS) + } +} + +#[inline(always)] +fn partition( + layout: Layout, + part: usize, +) -> Cyclic { + const { + assert!(PARTS > 0); + } + assert!(part < PARTS); + let len = if part >= layout.len() { + 0 + } else { + 1 + (layout.len() - 1 - part) / PARTS + }; + Cyclic { + parent: layout, + part, + len, + } +} + +/// Read-only borrowed storage and its typed logical coordinate map. +/// +/// `Element: TensorElement` preserves the same logical/storage identity +/// contract as `Tensor`. Scalar access returns the physical storage value; +/// numerical conversion remains explicit. +pub struct StridedTensor<'a, Element, Layout = Strided1D, Storage = Element> +where + Element: TensorElement, + Layout: IndexLayout, + Storage: Copy, +{ + base: *const Storage, + layout: Layout, + borrow: PhantomData<(&'a [Storage], Element)>, +} +impl, Layout: IndexLayout, Storage: Copy> Copy + for StridedTensor<'_, Element, Layout, Storage> +{ +} +impl, Layout: IndexLayout, Storage: Copy> Clone + for StridedTensor<'_, Element, Layout, Storage> +{ + fn clone(&self) -> Self { + *self + } +} + +impl<'a, T: Copy + TensorElement> StridedTensor<'a, T> { + /// Borrow an ordinary slice with its unchanged scalar element identity. + #[inline(always)] + pub fn from_slice(storage: &'a [T]) -> Self { + unsafe { Self::from_raw_parts(storage.as_ptr(), storage.len(), 1) } + } + /// Borrow the logical elements of a strided storage range. + /// + /// # Safety + /// `stride` is positive. Every `base.add(i * stride)`, `i < len`, must be + /// live, aligned and initialized for `'a`, with representable offsets and + /// no conflicting mutation. Holes need not be initialized. Empty views + /// may use a dangling aligned pointer and never perform pointer arithmetic. + #[inline(always)] + pub const unsafe fn from_raw_parts(base: *const T, len: usize, stride: usize) -> Self { + Self { + base, + layout: Strided1D { len, stride }, + borrow: PhantomData, + } + } +} +impl<'a, Element: TensorElement, Storage: Copy> + StridedTensor<'a, Element, Strided1D, Storage> +{ + /// Borrow contiguous encoded storage without changing its logical identity. + #[inline(always)] + pub fn from_storage(storage: &'a [Storage]) -> Self { + Self { + base: storage.as_ptr(), + layout: Strided1D { + len: storage.len(), + stride: 1, + }, + borrow: PhantomData, + } + } +} +impl<'a, Element: TensorElement, Layout: IndexLayout, Storage: Copy> + StridedTensor<'a, Element, Layout, Storage> +{ + /// Number of logical storage elements selected by this view. + #[inline(always)] + pub fn len(&self) -> usize { + self.layout.len() + } + /// Whether this view contains no logical elements. + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.layout.is_empty() + } + /// Retain the coordinate map for composition or inspection. + #[inline(always)] + pub const fn layout(&self) -> &Layout { + &self.layout + } + /// Select a checked interval of logical coordinates, retaining its parent map. + #[inline(always)] + pub fn subview( + self, + start: usize, + len: usize, + ) -> StridedTensor<'a, Element, Subrange, Storage> { + assert!(start <= self.len() && len <= self.len() - start); + unsafe { self.subview_unchecked(start, len) } + } + /// Select a logical interval after an enclosing extent check. + /// # Safety + /// `start <= self.len()` and `len <= self.len() - start` must hold. + #[inline(always)] + pub unsafe fn subview_unchecked( + self, + start: usize, + len: usize, + ) -> StridedTensor<'a, Element, Subrange, Storage> { + StridedTensor { + base: self.base, + layout: Subrange { + parent: self.layout, + start, + len, + }, + borrow: PhantomData, + } + } + /// Select partition `index` of a `PARTS`-way cyclic split, preserving the + /// count in the returned type. + /// + /// Logical coordinate `i` of the result is parent coordinate + /// `index + i * PARTS`. This is CuTe's `local_partition` with a `PARTS:1` + /// thread layout: the thread count is the const generic and the thread + /// index is the argument. A ragged tail leaves the later partitions + /// shorter instead of reading past the parent. + #[inline(always)] + pub fn local_partition( + self, + index: usize, + ) -> StridedTensor<'a, Element, Cyclic, Storage> { + StridedTensor { + base: self.base, + layout: partition::(self.layout, index), + borrow: PhantomData, + } + } + /// Read a checked logical coordinate through its composed layout. + #[inline(always)] + pub fn get(&self, index: usize) -> Storage { + assert!(index < self.len()); + unsafe { self.get_unchecked(index) } + } + /// Read a logical coordinate after an enclosing bounds check. + /// # Safety + /// `index < self.len()` must hold. + #[inline(always)] + pub unsafe fn get_unchecked(&self, index: usize) -> Storage { + unsafe { self.base.add(self.layout.crd2idx(index)).read() } + } + /// Reduce logical coordinates in increasing order without reassociation. + #[inline(always)] + pub fn fold(self, initial: U, mut f: impl FnMut(U, Storage) -> U) -> U { + let mut result = initial; + let mut index = 0; + while index < self.len() { + result = f(result, unsafe { self.get_unchecked(index) }); + index += 1; + } + result + } + /// Gather the first `N` logical values into an ordered register tile. + #[inline(always)] + pub fn load(self) -> RegisterTile { + assert!(N <= self.len()); + unsafe { self.load_unchecked() } + } + /// Gather an ordered register tile after an enclosing extent check. + /// # Safety + /// `N <= self.len()` must hold. Vector alignment is not required. + #[inline(always)] + pub unsafe fn load_unchecked(self) -> RegisterTile { + RegisterTile::from_fn(|i| unsafe { self.get_unchecked(i) }) + } + /// Transform equally sized logical views, even when their layouts differ. + /// # Safety + /// Destination coordinates must have one writer and no conflicting access. + /// Destination writes cannot change input values read by another iteration + /// or thread. Both views must have valid storage throughout the operation. + #[inline(always)] + pub unsafe fn map_to, L: IndexLayout, S: Copy>( + self, + output: &mut StridedTensorMut<'_, E, L, S>, + mut f: impl FnMut(Storage) -> S, + ) { + assert!(self.len() == output.len()); + let mut i = 0; + while i < self.len() { + unsafe { + output.store_at(i, f(self.get_unchecked(i))); + } + i += 1; + } + } + /// Transform corresponding logical coordinates, preserving evaluation order. + /// # Safety + /// The destination has `map_to`'s non-overlap requirements for both inputs. + /// All views must have equal logical lengths. + #[inline(always)] + pub unsafe fn zip_map_to< + E: TensorElement, + L: IndexLayout, + S: Copy, + OE: TensorElement, + OL: IndexLayout, + OS: Copy, + >( + self, + other: StridedTensor<'_, E, L, S>, + output: &mut StridedTensorMut<'_, OE, OL, OS>, + mut f: impl FnMut(Storage, S) -> OS, + ) { + assert!(self.len() == other.len() && self.len() == output.len()); + let mut i = 0; + while i < self.len() { + unsafe { + output.store_at(i, f(self.get_unchecked(i), other.get_unchecked(i))); + } + i += 1; + } + } +} + +/// A uniquely owned writable view and its typed logical coordinate map. +/// +/// This type is neither Copy nor Clone. GPU stores still require unsafe because +/// Rust cannot prove that different hardware threads chose disjoint coordinates. +pub struct StridedTensorMut<'a, Element, Layout = Strided1D, Storage = Element> +where + Element: TensorElement, + Layout: IndexLayout, + Storage: Copy, +{ + base: *mut Storage, + layout: Layout, + borrow: PhantomData<(&'a mut [Storage], Element)>, +} +impl<'a, T: Copy + TensorElement> StridedTensorMut<'a, T> { + /// Borrow a mutable slice without changing its scalar element identity. + #[inline(always)] + pub fn from_mut_slice(storage: &'a mut [T]) -> Self { + unsafe { Self::from_raw_parts(storage.as_mut_ptr(), storage.len(), 1) } + } + /// Borrow writable device storage while retaining its original alias rules. + #[inline(always)] + pub fn from_disjoint_slice(storage: &'a mut DisjointSlice<'_, T>) -> Self { + unsafe { Self::from_raw_parts(storage.as_mut_ptr(), storage.len(), 1) } + } + /// Borrow writable strided storage with an externally established extent. + /// # Safety + /// `stride` is positive. Each `base.add(i*stride)`, `i < len`, must be live, + /// aligned and writable for `'a`, with representable pointer offsets. + /// Stores separately require disjoint active access; reads additionally + /// require initialization. Holes need not contain initialized values. + #[inline(always)] + pub const unsafe fn from_raw_parts(base: *mut T, len: usize, stride: usize) -> Self { + Self { + base, + layout: Strided1D { len, stride }, + borrow: PhantomData, + } + } +} +impl<'a, Element: TensorElement, Layout: IndexLayout, Storage: Copy> + StridedTensorMut<'a, Element, Layout, Storage> +{ + /// Number of selected logical elements. + #[inline(always)] + pub fn len(&self) -> usize { + self.layout.len() + } + /// Whether this writable view selects no elements. + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.layout.is_empty() + } + /// Inspect the retained coordinate map without granting storage access. + #[inline(always)] + pub const fn layout(&self) -> &Layout { + &self.layout + } + /// Select a checked logical interval, consuming the original writable view. + #[inline(always)] + pub fn subview( + self, + start: usize, + len: usize, + ) -> StridedTensorMut<'a, Element, Subrange, Storage> { + assert!(start <= self.len() && len <= self.len() - start); + unsafe { self.subview_unchecked(start, len) } + } + /// Reborrow the same typed view; the parent cannot be used during this borrow. + #[inline(always)] + pub fn reborrow(&mut self) -> StridedTensorMut<'_, Element, Layout, Storage> { + StridedTensorMut { + base: self.base, + layout: self.layout, + borrow: PhantomData, + } + } + /// Select a logical interval after an enclosing extent check. + /// # Safety + /// `start <= self.len()` and `len <= self.len() - start` must hold. + #[inline(always)] + pub unsafe fn subview_unchecked( + self, + start: usize, + len: usize, + ) -> StridedTensorMut<'a, Element, Subrange, Storage> { + StridedTensorMut { + base: self.base, + layout: Subrange { + parent: self.layout, + start, + len, + }, + borrow: PhantomData, + } + } + /// Select partition `index` of a `PARTS`-way cyclic split, consuming the + /// original writable view. + /// + /// Logical coordinate `i` of the result is parent coordinate + /// `index + i * PARTS`, as CuTe's `local_partition` maps it with a + /// `PARTS:1` thread layout. A ragged tail leaves the later partitions + /// shorter instead of reaching past the parent. + #[inline(always)] + pub fn local_partition( + self, + index: usize, + ) -> StridedTensorMut<'a, Element, Cyclic, Storage> { + StridedTensorMut { + base: self.base, + layout: partition::(self.layout, index), + borrow: PhantomData, + } + } + /// Store one logical coordinate using the composed coordinate map. + /// # Safety + /// `index < self.len()` and no conflicting access to this coordinate may + /// occur, including access from another GPU thread. + #[inline(always)] + pub unsafe fn store_at(&mut self, index: usize, value: Storage) { + unsafe { + self.base.add(self.layout.crd2idx(index)).write(value); + } + } + /// Read, transform and replace one initialized logical coordinate. + /// # Safety + /// `index < self.len()`, the element is initialized, and no conflicting + /// active read or write may occur. This is not an atomic operation. + #[inline(always)] + pub unsafe fn update_at(&mut self, index: usize, f: impl FnOnce(Storage) -> Storage) { + let pointer = unsafe { self.base.add(self.layout.crd2idx(index)) }; + unsafe { + pointer.write(f(pointer.read())); + } + } + /// Scatter an ordered register tile through this view's coordinate map. + /// # Safety + /// No other active access may conflict with the selected destinations. + /// This method checks that `N` logical coordinates are present. + #[inline(always)] + pub unsafe fn store(&mut self, values: RegisterTile) { + assert!(N <= self.len()); + let values = values.into_array(); + let mut i = 0; + while i < N { + ctm_rs::thread::__unroll_config::<0>(); + unsafe { + self.store_at(i, values[i]); + } + i += 1; + } + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/arch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/arch.rs new file mode 100644 index 0000000000..ea9afc1b87 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/arch.rs @@ -0,0 +1,25 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![doc = "Checks the unsafe u32 contract of the warp-uniform compiler hint."] +/// Compile checks must reject unchecked calls and non-u32 warp-uniform arguments. +#[test] +fn warp_uniform_hint_requires_an_explicit_unsafe_u32_contract() { + let _: unsafe fn(u32) -> u32 = cute_rs::arch::make_warp_uniform; + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/warp_uniform_unsafe.rs"); + tests.compile_fail("tests/compile_fail/warp_uniform_width.rs"); +} + +/// The sync-restricted fence is an unsafe forwarder whose order is a type parameter. +#[test] +fn fence_sync_restrict_takes_its_order_as_a_type() { + let _: unsafe fn() = cute_rs::arch::fence_sync_restrict::; +} + +/// The cluster coordinate triples are plain readers of the grid-level cluster registers. +#[test] +fn cluster_triples_are_plain_coordinate_readers() { + let _: fn() -> (u32, u32, u32) = cute_rs::arch::cluster_idx; + let _: fn() -> (u32, u32, u32) = cute_rs::arch::cluster_dim; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/argmax.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/argmax.rs new file mode 100644 index 0000000000..23ec248e57 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/argmax.rs @@ -0,0 +1,44 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Stable first-index selection and sticky invalid-input reduction contracts. + +use cute_rs::reduction::ArgMax; + +#[test] +fn first_index_ties_are_independent_of_partition_order() { + let a = ArgMax::empty().observe(3.0, 70).observe(3.0, 2); + let b = ArgMax::empty().observe(-8.0, 0).observe(3.0, 33); + for selected in [a.merge(b), b.merge(a)] { + assert_eq!(selected.index(), 2); + assert_eq!(selected.value(), 3.0); + assert!(selected.is_valid()); + } +} + +#[test] +fn signed_zero_ties_choose_the_first_index() { + let selected = ArgMax::empty().observe(-0.0, 8).observe(0.0, 3); + assert_eq!(selected.index(), 3); + assert_eq!(selected.value().to_bits(), 0); +} + +#[test] +fn nonfinite_input_is_sticky_across_observations_and_merges() { + for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let invalid = ArgMax::empty().observe(bad, 0).observe(5.0, 1); + let valid = ArgMax::empty().observe(10.0, 2); + assert!(!invalid.is_valid()); + assert!(!invalid.merge(valid).is_valid()); + assert!(!valid.merge(invalid).is_valid()); + } +} + +#[test] +fn empty_is_an_identity_for_finite_partitions() { + let value = ArgMax::empty().observe(-5.0, 9); + assert_eq!(ArgMax::empty().merge(value).index(), 9); + assert_eq!(value.merge(ArgMax::empty()).index(), 9); + let _: unsafe fn(ArgMax, u32) -> ArgMax = ArgMax::reduce_warp; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/attention.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/attention.rs new file mode 100644 index 0000000000..771659be28 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/attention.rs @@ -0,0 +1,264 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Attention and cache extents are checked without allocating device memory. +#![cfg(feature = "host")] + +use cute_rs::attention::{AttentionError, AttentionShape, AttentionWindow, CacheShape}; + +#[test] +fn gqa_dimensions_preserve_independent_query_and_kv_widths() { + let shape = AttentionShape::<16, 8, 128>::new().unwrap(); + assert_eq!( + (shape.query_heads(), shape.kv_heads(), shape.head_dim()), + (16, 8, 128) + ); + assert_eq!( + (shape.query_width(), shape.kv_width(), shape.groups()), + (2048, 1024, 2) + ); + assert_eq!(shape.query_len(3), Ok(6144)); + assert_eq!(shape.kv_len(3), Ok(3072)); + let small = AttentionShape::<6, 2, 5>::new().unwrap(); + assert_eq!( + (small.query_width(), small.kv_width(), small.groups()), + (30, 10, 3) + ); +} + +#[test] +fn attention_rejects_zero_dimensions_and_incomplete_query_groups() { + assert_eq!( + AttentionShape::<0, 1, 8>::new(), + Err(AttentionError::ZeroDimension) + ); + assert_eq!( + AttentionShape::<1, 0, 8>::new(), + Err(AttentionError::ZeroDimension) + ); + assert_eq!( + AttentionShape::<1, 1, 0>::new(), + Err(AttentionError::ZeroDimension) + ); + assert_eq!( + AttentionShape::<3, 2, 8>::new(), + Err(AttentionError::InvalidHeadGrouping) + ); + assert_eq!( + AttentionShape::<1, 2, 8>::new(), + Err(AttentionError::InvalidHeadGrouping) + ); +} + +#[test] +fn attention_checks_width_multiplication_and_bf16_addressability() { + assert_eq!( + AttentionShape::<{ usize::MAX }, 1, 2>::new(), + Err(AttentionError::ExtentOverflow) + ); + assert_eq!( + AttentionShape::<1, 1, { isize::MAX as usize / 2 + 1 }>::new(), + Err(AttentionError::ExtentOverflow) + ); + let widest = AttentionShape::<1, 1, { isize::MAX as usize / 2 }>::new().unwrap(); + assert_eq!(widest.query_len(1), Ok(isize::MAX as usize / 2)); + assert_eq!(widest.query_len(2), Err(AttentionError::ExtentOverflow)); +} + +#[test] +fn token_lengths_accept_empty_inputs_and_check_both_tensor_extents() { + let shape = AttentionShape::<2, 1, 1>::new().unwrap(); + assert_eq!(shape.query_len(0), Ok(0)); + assert_eq!(shape.kv_len(0), Ok(0)); + assert_eq!(shape.query_len(65), Ok(130)); + assert_eq!(shape.kv_len(65), Ok(65)); + assert_eq!( + shape.query_len(usize::MAX), + Err(AttentionError::ExtentOverflow) + ); + assert_eq!( + shape.kv_len(isize::MAX as usize / 2), + Ok(isize::MAX as usize / 2) + ); + assert_eq!( + shape.kv_len(isize::MAX as usize / 2 + 1), + Err(AttentionError::ExtentOverflow) + ); +} + +#[test] +fn cache_ranges_count_tokens_and_accept_non_tile_aligned_appends() { + let shape = CacheShape::<3, 5>::new(65).unwrap(); + assert_eq!( + (shape.capacity(), shape.width(), shape.allocation_len()), + (65, 15, 975) + ); + assert_eq!(shape.append_range(0, 17), Ok(0..17)); + assert_eq!(shape.append_range(17, 48), Ok(17..65)); + assert_eq!(shape.append_range(65, 0), Ok(65..65)); + // A range is metadata; requesting it does not mutate a cache cursor. + assert_eq!(shape.append_range(0, 17), Ok(0..17)); +} + +#[test] +fn zero_capacity_keeps_head_dimensions_valid_and_permits_empty_append() { + let empty = CacheShape::<8, 128>::new(0).unwrap(); + assert_eq!( + (empty.capacity(), empty.width(), empty.allocation_len()), + (0, 1024, 0) + ); + assert_eq!(empty.append_range(0, 0), Ok(0..0)); + assert_eq!( + empty.append_range(0, 1), + Err(AttentionError::CapacityExceeded) + ); + assert_eq!( + CacheShape::<0, 128>::new(0), + Err(AttentionError::ZeroDimension) + ); + assert_eq!( + CacheShape::<8, 0>::new(0), + Err(AttentionError::ZeroDimension) + ); + assert_eq!( + CacheShape::<1, { isize::MAX as usize / 2 + 1 }>::new(0), + Err(AttentionError::ExtentOverflow) + ); +} + +#[test] +fn append_rejects_invalid_cursors_capacity_exhaustion_and_overflow() { + let shape = CacheShape::<2, 3>::new(65).unwrap(); + assert_eq!( + shape.append_range(66, 0), + Err(AttentionError::InvalidCursor) + ); + assert_eq!( + shape.append_range(64, 2), + Err(AttentionError::CapacityExceeded) + ); + assert_eq!( + shape.append_range(65, 1), + Err(AttentionError::CapacityExceeded) + ); + assert_eq!( + shape.append_range(1, usize::MAX), + Err(AttentionError::ExtentOverflow) + ); +} + +#[test] +fn cache_allocation_checks_full_bf16_extent_before_allocation() { + let max_elements = isize::MAX as usize / 2; + let largest = CacheShape::<1, 1>::new(max_elements).unwrap(); + assert_eq!(largest.allocation_len(), max_elements); + assert_eq!( + largest.append_range(max_elements - 1, 1), + Ok(max_elements - 1..max_elements) + ); + assert_eq!( + CacheShape::<1, 1>::new(max_elements + 1), + Err(AttentionError::ExtentOverflow) + ); + assert_eq!( + CacheShape::<1, 2>::new(usize::MAX), + Err(AttentionError::ExtentOverflow) + ); + assert_eq!( + CacheShape::<{ usize::MAX }, 2>::new(0), + Err(AttentionError::ExtentOverflow) + ); +} + +#[test] +fn causal_window_tracks_absolute_query_positions_in_a_cached_prefix() { + let window = AttentionWindow::causal(3, 19, 16).unwrap(); + assert_eq!(window.query_tokens(), 3); + assert_eq!(window.key_tokens(), 19); + assert_eq!(window.causal_start(), Some(16)); + assert_eq!(window.visible_keys(0), Ok(17)); + assert_eq!(window.visible_keys(1), Ok(18)); + assert_eq!(window.visible_keys(2), Ok(19)); + + // Extra key rows remain masked; a window need not end at the cache cursor. + let with_future_keys = AttentionWindow::causal(3, 25, 16).unwrap(); + assert_eq!(with_future_keys.visible_keys(2), Ok(19)); + let decode = AttentionWindow::causal(1, 65, 64).unwrap(); + assert_eq!(decode.visible_keys(0), Ok(65)); +} + +#[test] +fn noncausal_window_exposes_all_keys_even_with_more_queries_than_keys() { + let window = AttentionWindow::noncausal(5, 3).unwrap(); + assert_eq!(window.query_tokens(), 5); + assert_eq!(window.key_tokens(), 3); + assert_eq!(window.causal_start(), None); + assert_eq!(window.visible_keys(0), Ok(3)); + assert_eq!(window.visible_keys(4), Ok(3)); +} + +#[test] +fn empty_query_windows_are_valid_but_have_no_accessible_rows() { + for window in [ + AttentionWindow::causal(0, 0, 0).unwrap(), + AttentionWindow::causal(0, 65, 65).unwrap(), + AttentionWindow::noncausal(0, 0).unwrap(), + AttentionWindow::noncausal(0, 65).unwrap(), + ] { + assert_eq!(window.query_tokens(), 0); + assert_eq!(window.visible_keys(0), Err(AttentionError::InvalidQueryRow)); + } + assert_eq!( + AttentionWindow::causal(0, 65, 66), + Err(AttentionError::InvalidWindow) + ); +} + +#[test] +fn window_rejects_queries_without_keys_and_out_of_bounds_causal_ranges() { + assert_eq!( + AttentionWindow::noncausal(1, 0), + Err(AttentionError::MissingKeys) + ); + assert_eq!( + AttentionWindow::causal(1, 0, 0), + Err(AttentionError::MissingKeys) + ); + assert_eq!( + AttentionWindow::causal(3, 19, 17), + Err(AttentionError::InvalidWindow) + ); + assert_eq!( + AttentionWindow::causal(1, usize::MAX, usize::MAX), + Err(AttentionError::ExtentOverflow) + ); + for window in [ + AttentionWindow::causal(3, 19, 16).unwrap(), + AttentionWindow::noncausal(3, 19).unwrap(), + ] { + assert_eq!(window.visible_keys(3), Err(AttentionError::InvalidQueryRow)); + assert_eq!( + window.visible_keys(usize::MAX), + Err(AttentionError::InvalidQueryRow) + ); + } +} + +#[test] +fn window_checks_token_arithmetic_without_claiming_allocation_extents() { + let boundary = AttentionWindow::causal(1, usize::MAX, usize::MAX - 1).unwrap(); + assert_eq!(boundary.visible_keys(0), Ok(usize::MAX)); + let empty = AttentionWindow::causal(0, usize::MAX, usize::MAX).unwrap(); + assert_eq!(empty.causal_start(), Some(usize::MAX)); + let noncausal = AttentionWindow::noncausal(1, usize::MAX).unwrap(); + assert_eq!(noncausal.visible_keys(0), Ok(usize::MAX)); + // Actual tensor extents remain a separate check at the operator boundary. + assert_eq!( + AttentionShape::<1, 1, 1>::new() + .unwrap() + .kv_len(noncausal.key_tokens()), + Err(AttentionError::ExtentOverflow) + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16.rs new file mode 100644 index 0000000000..ff517a9a92 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16.rs @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! BF16 storage and explicit conversion contracts. +use cute_rs::{BFloat16, RegisterTile, Tensor}; + +#[test] +fn storage_preserves_every_pattern_and_tensor_coordinates() { + assert_eq!(core::mem::size_of::(), 2); + assert_eq!(core::mem::align_of::(), 2); + for bits in 0..=u16::MAX { + let value = BFloat16::from_bits(bits); + assert_eq!(value.to_bits(), bits); + assert_eq!(value.to_f32().to_bits(), (bits as u32) << 16); + } + let data = [BFloat16::from_bits(0); 9]; + let tail = Tensor::from_slice(&data).zipped_divide::<8>().slice(1); + assert_eq!(tail.base(), 8); + assert!(!tail.is_full()); +} + +#[test] +fn conversion_rounds_ties_even_and_preserves_special_values() { + for (input, expected) in [ + (0x3f80_8000, 0x3f80), + (0x3f81_8000, 0x3f82), + (0xbf80_8000, 0xbf80), + (0xbf81_8000, 0xbf82), + (0, 0), + (0x8000_0000, 0x8000), + (0x7f80_0000, 0x7f80), + (0xff80_0000, 0xff80), + (0x0000_8000, 0), + (0x0001_8000, 2), + ] { + assert_eq!( + BFloat16::from_f32(f32::from_bits(input)).to_bits(), + expected + ); + } + for bits in [0x7f80_0001, 0xff80_0001] { + assert!(BFloat16::from_f32(f32::from_bits(bits)).to_f32().is_nan()); + } +} + +#[test] +fn register_arithmetic_requires_explicit_widening_and_narrowing() { + let tile = RegisterTile::from_array([BFloat16::from_f32(1.0), BFloat16::from_f32(2.0)]); + let output = tile + .map(BFloat16::to_f32) + .map(|v| v * 0.5) + .map(BFloat16::from_f32); + assert_eq!(output.into_array().map(BFloat16::to_f32), [0.5, 1.0]); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16_mma.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16_mma.rs new file mode 100644 index 0000000000..30e39da9e2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/bf16_mma.rs @@ -0,0 +1,216 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent checks of BF16 operand coverage, register order and layout evaluation. +use cute_rs::BFloat16; + +#[test] +fn static_atom_layouts_match_cute_algebra_for_every_lane_and_value() { + use cute_rs::mma::bf16::{ALayout, BLayout, CLayout}; + use cute_rs::static_layout::StaticLayout; + fn check(description: &str) { + let oracle: cute_rs::layout::Layout = description.parse().unwrap(); + for flat in 0..L::SIZE { + assert_eq!( + L::crd2idx(flat), + oracle + .checked_call(&cute_rs::layout::IntTuple::Leaf(flat as i64)) + .unwrap() as usize + ); + } + } + check::("((4,8),((2,2),2)):((2,16),((1,128),8))"); + check::("((4,8),(2,2)):((2,16),(1,8))"); + check::("((4,8),(2,2)):((2,8),(1,64))"); +} + +#[test] +fn packed_output_partition_visits_only_live_decode_rows() { + use cute_rs::mma::bf16::PackedDecode; + #[repr(align(4))] + struct Aligned([BFloat16; 256]); + let data = Aligned([BFloat16::from_bits(0); 256]); + for rows in [1, 2] { + let source = PackedDecode::new(&data.0, &data.0, rows, 16, 8, 16, 16).unwrap(); + let mut seen = vec![0; rows * 8]; + for lane in 0..32 { + Bf16Mma16x8x16::get_slice(lane) + .partition_c() + .store_decode_with(FragmentC::ZERO, &source, |r, c, _| { + seen[r * 8 + c] += 1; + }); + } + assert_eq!(seen, vec![1; rows * 8]); + } +} +use cute_rs::markers::{ComposedLayout, RowMajor, Swizzle}; +use cute_rs::mma::bf16::{Bf16Mma16x8x16, FragmentC}; + +#[test] +fn packed_decode_proof_rejects_incomplete_or_misaligned_operands() { + use cute_rs::mma::bf16::PackedDecode; + #[repr(align(4))] + struct Aligned([BFloat16; 1024]); + let data = Aligned([BFloat16::from_bits(0); 1024]); + assert!(PackedDecode::new(&data.0, &data.0, 1, 16, 8, 16, 16).is_some()); + assert!(PackedDecode::new(&data.0, &data.0, 2, 32, 16, 34, 34).is_some()); + for (rows, width, outputs, is, ws) in [ + (0, 16, 8, 16, 16), + (3, 16, 8, 16, 16), + (1, 15, 8, 16, 16), + (1, 16, 7, 16, 16), + (1, 16, 8, 15, 16), + (1, 16, 8, 16, 17), + ] { + assert!(PackedDecode::new(&data.0, &data.0, rows, width, outputs, is, ws).is_none()); + } + assert!(PackedDecode::new(&data.0[1..], &data.0, 1, 16, 8, 16, 16).is_none()); + assert!(PackedDecode::new(&data.0[..15], &data.0, 1, 16, 8, 16, 16).is_none()); + assert!(PackedDecode::new(&data.0, &data.0[..127], 1, 16, 8, 16, 16).is_none()); +} + +#[test] +fn validated_decode_callback_runs_once_only_for_proven_storage() { + use cute_rs::mma::bf16::PackedDecode; + #[repr(align(4))] + struct Aligned([BFloat16; 128]); + let data = Aligned([BFloat16::from_bits(0); 128]); + let mut calls = 0; + assert!(PackedDecode::with_validated( + &data.0, + &data.0, + 1, + 16, + 8, + 16, + 16, + |proof| { + calls += 1; + let mut count = 0; + Bf16Mma16x8x16::get_slice(0) + .partition_c() + .store_decode_with(FragmentC::ZERO, proof, |_, _, _| count += 1); + assert_eq!(count, 2); + } + )); + assert!(!PackedDecode::with_validated( + &data.0[..15], + &data.0, + 1, + 16, + 8, + 16, + 16, + |_| calls += 1 + )); + assert_eq!(calls, 1); +} + +#[test] +fn operand_partition_matches_ptx_coordinates_and_covers_every_element() { + let mut a_seen = [0; 256]; + let mut b_seen = [0; 128]; + let mut c_seen = [0; 128]; + for lane in 0..32 { + let mma = Bf16Mma16x8x16::get_slice(lane); + let mut a_coords = Vec::new(); + let mut b_coords = Vec::new(); + mma.partition_a().load_with(|r, k| { + a_coords.push((r, k)); + a_seen[r * 16 + k] += 1; + BFloat16::from_bits(0) + }); + mma.partition_b().load_with(|n, k| { + b_coords.push((n, k)); + b_seen[n * 16 + k] += 1; + BFloat16::from_bits(0) + }); + let row = lane / 4; + let col = lane % 4 * 2; + assert_eq!( + a_coords, + [ + (row, col), + (row, col + 1), + (row + 8, col), + (row + 8, col + 1), + (row, col + 8), + (row, col + 9), + (row + 8, col + 8), + (row + 8, col + 9) + ] + ); + assert_eq!( + b_coords, + [(row, col), (row, col + 1), (row, col + 8), (row, col + 9)] + ); + let mut c_coords = Vec::new(); + mma.partition_c().store_with(FragmentC::ZERO, |r, c, v| { + c_coords.push((r, c)); + c_seen[r * 8 + c] += 1; + assert_eq!(v.to_bits(), 0); + }); + assert_eq!( + c_coords, + [ + (row, col), + (row, col + 1), + (row + 8, col), + (row + 8, col + 1) + ] + ); + } + assert_eq!(a_seen, [1; 256]); + assert_eq!(b_seen, [1; 128]); + assert_eq!(c_seen, [1; 128]); +} + +#[test] +fn shared_layout_offset_matches_independent_cute_layout_oracle() { + type L = ComposedLayout, 0, RowMajor<64, 128>>; + let layout: cute_rs::layout::Layout = "(64,128):(128,1)".parse().unwrap(); + let swizzle = cute_rs::layout::Swizzle::new(3, 3, 3); + for r in 0..64 { + for c in 0..128 { + let plain = layout + .checked_call(&cute_rs::layout::IntTuple::Tuple(vec![ + cute_rs::layout::IntTuple::Leaf(r), + cute_rs::layout::IntTuple::Leaf(c), + ])) + .unwrap(); + assert_eq!( + cute_rs::static_layout::crd2idx_2d::(r as usize, c as usize), + swizzle.apply(plain) as usize + ); + } + } +} + +#[test] +fn static_composition_applies_nonzero_offset_before_negative_swizzle() { + use cute_rs::static_layout::StaticLayout; + type L = ComposedLayout, 7, RowMajor<3, 5>>; + let inner: cute_rs::layout::Layout = "(3,5):(5,1)".parse().unwrap(); + let swizzle = cute_rs::layout::Swizzle::new(2, 1, -3); + for flat in 0..L::SIZE { + let plain = inner + .checked_call(&cute_rs::layout::IntTuple::Leaf(flat as i64)) + .unwrap(); + let expected = swizzle.apply(plain + 7) as usize; + assert_eq!(L::crd2idx(flat), expected); + assert_eq!( + cute_rs::static_layout::crd2idx_2d::(flat % 3, flat / 3), + expected + ); + } +} + +#[test] +fn mma_rejects_swapped_operand_roles() { + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/bf16_mma_partition_roles.rs"); + tests.compile_fail("tests/compile_fail/packed_decode_borrow_escape.rs"); + tests.compile_fail("tests/compile_fail/packed_decode_callback_escape.rs"); + tests.compile_fail("tests/compile_fail/packed_decode_callback_mutation.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.rs new file mode 100644 index 0000000000..078e234118 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.rs @@ -0,0 +1,9 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{BFloat16, RegisterTile}; +fn main() { + let tile = RegisterTile::from_array([BFloat16::from_bits(0); 4]); + let _: RegisterTile = tile; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.stderr new file mode 100644 index 0000000000..c4f8b3e24f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_implicit_conversion.stderr @@ -0,0 +1,10 @@ +error[E0308]: mismatched types + --> tests/compile_fail/bf16_implicit_conversion.rs:8:35 + | +8 | let _: RegisterTile = tile; + | -------------------- ^^^^ expected `RegisterTile`, found `RegisterTile` + | | + | expected due to this + | + = note: expected struct `RegisterTile` + found struct `RegisterTile` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.rs new file mode 100644 index 0000000000..263d66eda1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![feature(f16)] +use cute_rs::{BFloat16, Contiguous1D, Tensor}; +fn main() { + let values = [1.0_f16; 8]; + let _ = Tensor::::from_storage(&values, Contiguous1D::new(8)); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.stderr new file mode 100644 index 0000000000..cf8452f452 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_invalid_storage.stderr @@ -0,0 +1,35 @@ +error[E0277]: the trait bound `BFloat16: TensorElement` is not satisfied + --> tests/compile_fail/bf16_invalid_storage.rs:9:13 + | +9 | let _ = Tensor::::from_storage(&values, Contiguous1D::new(8)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `TensorElement` is not implemented for `BFloat16` + | +help: the trait `TensorElement` is not implemented for `BFloat16` + but trait `TensorElement` is implemented for it + --> src/numeric.rs + | + | impl TensorElement for BFloat16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = help: for that trait implementation, expected `BFloat16`, found `f16` +note: required by a bound in `Tensor` + --> src/tensor/mod.rs + | + | pub struct Tensor<'a, Element, Layout, Storage = Element> + | ------ required by a bound in this struct + | where + | Element: TensorElement, + | ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `Tensor` + +error[E0599]: the associated function or constant `from_storage` exists for struct `Tensor<'_, BFloat16, _, f16>`, but its trait bounds were not satisfied + --> tests/compile_fail/bf16_invalid_storage.rs:9:41 + | +9 | let _ = Tensor::::from_storage(&values, Contiguous1D::new(8)); + | ^^^^^^^^^^^^ associated function or constant cannot be called on `Tensor<'_, BFloat16, _, f16>` due to unsatisfied trait bounds + | + ::: src/numeric.rs + | + | pub struct BFloat16(#[doc = "The unchanged IEEE-style BF16 storage bits."] u16); + | ------------------- doesn't satisfy `BFloat16: TensorElement` + | + = note: the following trait bounds were not satisfied: + `BFloat16: TensorElement` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.rs new file mode 100644 index 0000000000..af3e0022f5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{ + BFloat16, + mma::bf16::{Bf16Mma16x8x16, FragmentC}, +}; +fn main() { + let mma = Bf16Mma16x8x16::get_slice(0); + let a = mma.partition_a().load_with(|_, _| BFloat16::from_bits(0)); + let b = mma.partition_b().load_with(|_, _| BFloat16::from_bits(0)); + let _ = unsafe { mma.mma_atom_call(FragmentC::ZERO, b, a) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.stderr new file mode 100644 index 0000000000..cd7c30670d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/bf16_mma_partition_roles.stderr @@ -0,0 +1,18 @@ +error[E0308]: arguments to this method are incorrect + --> tests/compile_fail/bf16_mma_partition_roles.rs:13:26 + | +13 | let _ = unsafe { mma.mma_atom_call(FragmentC::ZERO, b, a) }; + | ^^^^^^^^^^^^^ - - expected `cute_rs::mma::bf16::FragmentB`, found `cute_rs::mma::bf16::FragmentA` + | | + | expected `cute_rs::mma::bf16::FragmentA`, found `cute_rs::mma::bf16::FragmentB` + | +note: method defined here + --> src/mma/bf16.rs + | + | pub unsafe fn mma_atom_call(self, c: FragmentC, a: FragmentA, b: FragmentB) -> FragmentC { + | ^^^^^^^^^^^^^ +help: swap these arguments + | +13 - let _ = unsafe { mma.mma_atom_call(FragmentC::ZERO, b, a) }; +13 + let _ = unsafe { mma.mma_atom_call(FragmentC::ZERO, a, b) }; + | diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.rs new file mode 100644 index 0000000000..ab8ba94123 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.rs @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::tensor::partition::{Cyclic, Strided1D, StridedTensor}; +fn warp_partition(_: StridedTensor<'_, f32, Cyclic>) {} +fn main() { + let values = [0.0_f32; 128]; + warp_partition(StridedTensor::from_slice(&values).local_partition::<16>(0)); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.stderr new file mode 100644 index 0000000000..6cf1f89cd9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/cyclic_partition_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/cyclic_partition_mismatch.rs:9:20 + | +9 | warp_partition(StridedTensor::from_slice(&values).local_partition::<16>(0)); + | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `32`, found `16` + | | + | arguments to this function are incorrect + | + = note: expected struct `StridedTensor<'_, f32, Cyclic>` + found struct `StridedTensor<'_, f32, Cyclic>` +note: function defined here + --> tests/compile_fail/cyclic_partition_mismatch.rs:6:4 + | +6 | fn warp_partition(_: StridedTensor<'_, f32, Cyclic>) {} + | ^^^^^^^^^^^^^^ ------------------------------------------------ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.rs new file mode 100644 index 0000000000..fd313cfed9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject FP64 carriers in the supported global-store tile formats. +/// Reject FP64 carriers in the supported global-store tile formats. +fn main() { + let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.stderr new file mode 100644 index 0000000000..21ec26f822 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_storage.stderr @@ -0,0 +1,43 @@ +error[E0277]: the trait bound `f64: GlobalStoreElement` is not satisfied + --> tests/compile_fail/global_store_storage.rs:7:38 + | +7 | let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); + | ------------------------ ^^^^^^^^^^^ the trait `GlobalStoreElement` is not implemented for `f64` + | | + | required by a bound introduced by this call + | +help: the following other types implement trait `GlobalStoreElement` + --> src/copy/vector.rs + | + | impl GlobalStoreElement for f16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f16` + | impl GlobalStoreElement for f32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f32` + | impl GlobalStoreElement for u32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `u32` +note: required by a bound in `GlobalStoreTile` + --> src/copy/vector.rs + | + | pub struct GlobalStoreTile( + | ^^^^^^^^^^^^^^^^^^ required by this bound in `GlobalStoreTile` + +error[E0277]: the trait bound `f64: GlobalStoreElement` is not satisfied + --> tests/compile_fail/global_store_storage.rs:7:13 + | +7 | let _ = cute_rs::GlobalStoreTile([0.0f64; 4]); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `GlobalStoreElement` is not implemented for `f64` + | +help: the following other types implement trait `GlobalStoreElement` + --> src/copy/vector.rs + | + | impl GlobalStoreElement for f16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f16` + | impl GlobalStoreElement for f32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f32` + | impl GlobalStoreElement for u32 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `u32` +note: required by a bound in `GlobalStoreTile` + --> src/copy/vector.rs + | + | pub struct GlobalStoreTile( + | ^^^^^^^^^^^^^^^^^^ required by this bound in `GlobalStoreTile` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.rs new file mode 100644 index 0000000000..469e76e7cb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: require an unsafe block when issuing a raw-pointer global store. +/// Require an unsafe block when issuing a raw-pointer global store. +fn main() { + let values = cute_rs::GlobalStoreTile([0.0f32; 8]); + cute_rs::store_global_no_allocate(core::ptr::null_mut(), &values); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.stderr new file mode 100644 index 0000000000..1b5e588dc7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/global_store_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `store_global_no_allocate` is unsafe and requires unsafe block + --> tests/compile_fail/global_store_unsafe.rs:8:5 + | +8 | cute_rs::store_global_no_allocate(core::ptr::null_mut(), &values); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.rs new file mode 100644 index 0000000000..763765ab0c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject stores through an immutable tensor tile. +#![feature(f16)] + +use cute_rs::Tensor; + +/// Reject stores through an immutable tensor tile. + +fn main() { + let storage = [0.0f32; 4]; + let tile = Tensor::from_slice(&storage) + .zipped_divide::<4>() + .slice(0); + let values = unsafe { tile.load() }; + unsafe { tile.store(values) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.stderr new file mode 100644 index 0000000000..e86959ec90 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/immutable_store.stderr @@ -0,0 +1,11 @@ +error[E0599]: no method named `store` found for struct `Tensor<'a, Element, Layout, Storage>` in the current scope + --> tests/compile_fail/immutable_store.rs:17:19 + | +17 | unsafe { tile.store(values) }; + | ^^^^^ + | +help: there is a method `storage` with a similar name, but with different arguments + --> src/tensor/mod.rs + | + | pub const fn storage(&self) -> &'a [Storage] { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.rs new file mode 100644 index 0000000000..a78883ef24 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject an FP32 carrier for packed E2M1 tensor storage. +#![feature(f16)] + +use cute_rs::{Float4E2M1FN, KMajor, Mkl, Tensor}; + +/// Reject an FP32 carrier for packed E2M1 tensor storage. + +fn main() { + let wrong_carrier = [0.0f32; 32]; + let _ = Tensor::::from_storage( + &wrong_carrier, + KMajor::::new(1, 64), + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.stderr new file mode 100644 index 0000000000..0c20794013 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/invalid_storage.stderr @@ -0,0 +1,35 @@ +error[E0277]: the trait bound `Float4E2M1FN: TensorElement` is not satisfied + --> tests/compile_fail/invalid_storage.rs:13:13 + | +13 | let _ = Tensor::::from_storage( + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `TensorElement` is not implemented for `Float4E2M1FN` + | +help: the trait `TensorElement` is not implemented for `Float4E2M1FN` + but trait `TensorElement` is implemented for it + --> src/numeric.rs + | + | impl TensorElement for Float4E2M1FN {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = help: for that trait implementation, expected `u8`, found `f32` +note: required by a bound in `Tensor` + --> src/tensor/mod.rs + | + | pub struct Tensor<'a, Element, Layout, Storage = Element> + | ------ required by a bound in this struct + | where + | Element: TensorElement, + | ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `Tensor` + +error[E0599]: the associated function or constant `from_storage` exists for struct `Tensor<'_, Float4E2M1FN, _, f32>`, but its trait bounds were not satisfied + --> tests/compile_fail/invalid_storage.rs:13:45 + | +13 | let _ = Tensor::::from_storage( + | ^^^^^^^^^^^^ associated function or constant cannot be called on `Tensor<'_, Float4E2M1FN, _, f32>` due to unsatisfied trait bounds + | + ::: src/numeric.rs + | + | pub struct Float4E2M1FN; + | ----------------------- doesn't satisfy `Float4E2M1FN: TensorElement` + | + = note: the following trait bounds were not satisfied: + `Float4E2M1FN: TensorElement` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs new file mode 100644 index 0000000000..5a34277a55 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.rs @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject a B-role shared tile passed to the A-fragment loader. +#![feature(f16)] + +use core::ptr; + +use cute_rs::{Mxf4E2M1, Mxf4TiledMma, Nkl, SharedTensor}; + +/// Reject a B-role shared tile passed to the A-fragment loader. + +fn main() { + let shared_b = unsafe { + SharedTensor::::from_raw_parts(ptr::null_mut(), 0) + }; + let mma = Mxf4TiledMma::<()>::get_slice(0); + let _ = unsafe { mma.load_a(&shared_b, (0, 0)) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr new file mode 100644 index 0000000000..96c73695b6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_operand_role_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/mma_operand_role_mismatch.rs:18:33 + | +18 | let _ = unsafe { mma.load_a(&shared_b, (0, 0)) }; + | ------ ^^^^^^^^^ expected `&SharedTensor`, found `&SharedTensor` + | | + | arguments to this method are incorrect + | + = note: expected reference `&SharedTensor` + found reference `&SharedTensor` +note: method defined here + --> src/mma/block_scaled/tiled.rs + | + | pub unsafe fn load_a( + | ^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs new file mode 100644 index 0000000000..6ee18a6b63 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.rs @@ -0,0 +1,27 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject swapped A and B scale roles in block-scaled MMA. +#![feature(f16)] + +use cute_rs::{FragmentC, FragmentA, FragmentB, Mkl, Mxf4BlockScaledTensor, Mxf4TiledMma, Nkl}; + +/// Reject swapped A and B scale roles in block-scaled MMA. + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxf4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let vector = Mxf4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let mma = Mxf4TiledMma::<()>::get_slice(0); + + let _ = unsafe { + mma.gemm( + FragmentA([0; 4]), + vector.load_scale_pair(0, 0), + FragmentB([0; 2]), + matrix.load_scale_pair(0, 0), + FragmentC::ZERO, + ) + }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr new file mode 100644 index 0000000000..d2d23e63a5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/mma_scale_role_mismatch.stderr @@ -0,0 +1,27 @@ +error[E0308]: arguments to this method are incorrect + --> tests/compile_fail/mma_scale_role_mismatch.rs:19:13 + | +19 | mma.gemm( + | ^^^^ +20 | FragmentA([0; 4]), +21 | vector.load_scale_pair(0, 0), + | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` +22 | FragmentB([0; 2]), +23 | matrix.load_scale_pair(0, 0), + | ---------------------------- expected `MmaScalePair`, found `MmaScalePair` + | +note: method defined here + --> src/mma/block_scaled/tiled.rs + | + | pub unsafe fn gemm( + | ^^^^ +help: swap these arguments + | +19 ~ mma.gemm( +20 + FragmentA([0; 4]), +21 + matrix.load_scale_pair(0, 0), +22 + FragmentB([0; 2]), +23 + vector.load_scale_pair(0, 0), +24 + FragmentC::ZERO, +25 + ) + | diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.rs new file mode 100644 index 0000000000..89593d328e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{BFloat16, mma::bf16::PackedDecode}; +fn escape<'a>() -> Option> { + let storage = [BFloat16::from_bits(0); 128]; + PackedDecode::new(&storage, &storage, 1, 16, 8, 16, 16) +} +fn main() { + let _ = escape(); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.stderr new file mode 100644 index 0000000000..9995464b32 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_borrow_escape.stderr @@ -0,0 +1,17 @@ +error[E0515]: cannot return value referencing local variable `storage` + --> tests/compile_fail/packed_decode_borrow_escape.rs:8:5 + | +8 | PackedDecode::new(&storage, &storage, 1, 16, 8, 16, 16) + | ^^^^^^^^^^^^^^^^^^--------^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | | | + | | `storage` is borrowed here + | returns a value referencing data owned by the current function + +error[E0515]: cannot return value referencing local variable `storage` + --> tests/compile_fail/packed_decode_borrow_escape.rs:8:5 + | +8 | PackedDecode::new(&storage, &storage, 1, 16, 8, 16, 16) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^--------^^^^^^^^^^^^^^^^^^^ + | | | + | | `storage` is borrowed here + | returns a value referencing data owned by the current function diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.rs new file mode 100644 index 0000000000..8e2027d9fd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.rs @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{BFloat16, mma::bf16::PackedDecode}; +fn main() { + let storage = [BFloat16::from_bits(0); 128]; + let mut escaped = None; + PackedDecode::with_validated(&storage, &storage, 1, 16, 8, 16, 16, |proof| { + escaped = Some(proof); + }); + let _ = escaped; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.stderr new file mode 100644 index 0000000000..2fb1d17433 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_escape.stderr @@ -0,0 +1,9 @@ +error[E0521]: borrowed data escapes outside of closure + --> tests/compile_fail/packed_decode_callback_escape.rs:10:9 + | + 8 | let mut escaped = None; + | ----------- `escaped` declared here, outside of the closure body + 9 | PackedDecode::with_validated(&storage, &storage, 1, 16, 8, 16, 16, |proof| { + | ----- `proof` is a reference that is only valid in the closure body +10 | escaped = Some(proof); + | ^^^^^^^^^^^^^^^^^^^^^ `proof` escapes the closure body here diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.rs new file mode 100644 index 0000000000..edf1eddfa1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{BFloat16, mma::bf16::PackedDecode}; +fn main() { + let mut storage = [BFloat16::from_bits(0); 128]; + PackedDecode::with_validated(&storage, &storage, 1, 16, 8, 16, 16, |_| { + storage[0] = BFloat16::from_bits(1); + }); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.stderr new file mode 100644 index 0000000000..72743041e2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/packed_decode_callback_mutation.stderr @@ -0,0 +1,9 @@ +error[E0502]: cannot borrow `storage` as mutable because it is also borrowed as immutable + --> tests/compile_fail/packed_decode_callback_mutation.rs:8:72 + | +8 | PackedDecode::with_validated(&storage, &storage, 1, 16, 8, 16, 16, |_| { + | ---------------------------- -------- immutable borrow occurs here ^^^ mutable borrow occurs here + | | + | immutable borrow later used by call +9 | storage[0] = BFloat16::from_bits(1); + | ------- second borrow occurs due to use of `storage` in closure diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.rs new file mode 100644 index 0000000000..9583c481c8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.rs @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![feature(f16)] + +use cute_rs::tensor::RegisterTile; + +fn main() { + let half = RegisterTile::from_array([1.0_f16; 4]); + let single = RegisterTile::from_array([2.0_f32; 4]); + let _ = half.zip_map(single, |a, b| a + b); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.stderr new file mode 100644 index 0000000000..2347b336f2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_implicit_conversion.stderr @@ -0,0 +1,33 @@ +error[E0308]: mismatched types + --> tests/compile_fail/register_tile_implicit_conversion.rs:13:45 + | +13 | let _ = half.zip_map(single, |a, b| a + b); + | ^ expected `f16`, found `f32` + +error[E0277]: cannot add `f32` to `f16` + --> tests/compile_fail/register_tile_implicit_conversion.rs:13:43 + | +13 | let _ = half.zip_map(single, |a, b| a + b); + | ^ no implementation for `f16 + f32` + | + = help: the trait `Add` is not implemented for `f16` +help: the following other types implement trait `Add` + --> $RUST/core/src/ops/arith.rs + | + | const impl Add for $t { + | ^^^^^^^^^^^^^^^^^^^^^ `f16` implements `Add` +... + | add_impl! { usize u8 u16 u32 u64 u128 isize i8 i16 i32 i64 i128 f16 f32 f64 f128 } + | ---------------------------------------------------------------------------------- in this macro invocation + | + ::: $RUST/core/src/internal_macros.rs + | + | const impl $imp<$u> for &$t { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ `&f16` implements `Add` +... + | const impl $imp<&$u> for $t { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ `f16` implements `Add<&f16>` +... + | const impl $imp<&$u> for &$t { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `&f16` implements `Add` + = note: this error originates in the macro `add_impl` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.rs new file mode 100644 index 0000000000..e20234c3c0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use cute_rs::tensor::RegisterTile; + +fn main() { + let four = RegisterTile::from_array([1.0_f32; 4]); + let eight = RegisterTile::from_array([2.0_f32; 8]); + let _ = four.zip_map(eight, |a, b| a + b); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.stderr new file mode 100644 index 0000000000..79463810b2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/register_tile_lane_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/register_tile_lane_mismatch.rs:11:26 + | +11 | let _ = four.zip_map(eight, |a, b| a + b); + | ------- ^^^^^ expected `4`, found `8` + | | + | arguments to this method are incorrect + | + = note: expected struct `RegisterTile<_, 4>` + found struct `RegisterTile` +note: method defined here + --> src/tensor/mod.rs + | + | pub fn zip_map( + | ^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs new file mode 100644 index 0000000000..b578fd4ac2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.rs @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a K16 shared/shared attention configuration. +use cute_rs::attention::sm100::{MmaSs, SharedLayout}; +/// K-major shared layout too narrow for this attention MMA path. +type UnsupportedK16 = SharedLayout<128, 16>; +/// Reject a K16 shared/shared attention configuration. +const _: MmaSs = MmaSs::new(); +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr new file mode 100644 index 0000000000..5e352eb438 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_invalid_layout.stderr @@ -0,0 +1,23 @@ +error[E0080]: evaluation panicked: assertion failed: contiguous >= 32 && contiguous.is_multiple_of(32) + --> src/attention/sm100/shared.rs + | + | assert!(contiguous >= 32 && contiguous.is_multiple_of(32)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::SharedLayout::<128, 16>::VALID` failed here + +note: erroneous constant encountered + --> src/attention/sm100/shared.rs + | + | const VALID: () = SharedLayout::::VALID; + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +note: erroneous constant encountered + --> src/attention/sm100/mma.rs + | + | let () = AL::VALID; + | ^^^^^^^^^ + +note: erroneous constant encountered + --> src/attention/sm100/mma.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs new file mode 100644 index 0000000000..a3927df6c3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.rs @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a B-role tile in the shared/shared MMA A position. +#![feature(f16)] +use cute_rs::attention::sm100::{MmaSs, SharedLayout, SharedTile}; +use cute_rs::OperandB; +/// Supported 128-by-128 layout shared by both role-checked operands. +type L = SharedLayout<128, 128>; +/// Reject a B-role tile in the shared/shared MMA A position. +fn main() { + let a = unsafe { SharedTile::::from_raw(core::ptr::null_mut()) }; + let b = unsafe { SharedTile::::from_raw(core::ptr::null_mut()) }; + let mma = MmaSs::::new(); + unsafe { mma.gemm(0, &a, &b, false) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr new file mode 100644 index 0000000000..faf40faa19 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_operand_role.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/sm100_attention_operand_role.rs:14:26 + | +14 | unsafe { mma.gemm(0, &a, &b, false) }; + | ---- ^^ expected `OperandA`, found `OperandB` + | | + | arguments to this method are incorrect + | + = note: expected reference `&SharedTile, OperandA>` + found reference `&SharedTile, OperandB>` +note: method defined here + --> src/attention/sm100/mma.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs new file mode 100644 index 0000000000..581bbd6852 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject empty and byte-overflowing shared tile arrays at compile time. +#![feature(f16)] +use cute_rs::OperandB; +use cute_rs::attention::sm100::{SharedLayout, SharedTileArray}; +/// Small 8-by-32 shared layout whose stage-array extent is checked. +type L = SharedLayout<8, 32>; +/// Force const evaluation of an invalid zero-tile array. +const EMPTY: SharedTileArray = + unsafe { SharedTileArray::from_raw(core::ptr::null_mut()) }; +/// Force const evaluation of an array whose byte extent exceeds u32. +const OVERFLOW: SharedTileArray = + unsafe { SharedTileArray::from_raw(core::ptr::null_mut()) }; +/// Reject empty and byte-overflowing shared tile arrays at compile time. +fn main() { + let _ = (EMPTY, OVERFLOW); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr new file mode 100644 index 0000000000..e1bf6b751e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_extent.stderr @@ -0,0 +1,17 @@ +error[E0080]: evaluation panicked: assertion failed: TILES > 0 + --> src/attention/sm100/shared.rs + | + | assert!(TILES > 0); + | ^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::SharedTileArray::, cute_rs::OperandB, 0>::VALID` failed here + +note: erroneous constant encountered + --> src/attention/sm100/shared.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ + +error[E0080]: evaluation panicked: assertion failed: TILES <= u32::MAX / Self::TILE_BYTES + --> src/attention/sm100/shared.rs + | + | assert!(TILES <= u32::MAX / Self::TILE_BYTES); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::SharedTileArray::, cute_rs::OperandB, u32::MAX>::VALID` failed here diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs new file mode 100644 index 0000000000..994d8ccc24 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.rs @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: preserve B-role typing when a tile is borrowed from a shared tile array. +#![feature(f16)] +use cute_rs::OperandB; +use cute_rs::attention::sm100::{MmaSs, SharedLayout, SharedTileArray}; +/// Supported per-CTA 64-by-64 tile layout for group-two MMA. +type L = SharedLayout<64, 64>; +/// Preserve B-role typing when a tile is borrowed from a shared tile array. +fn main() { + let tiles = unsafe { SharedTileArray::::from_raw(core::ptr::null_mut()) }; + let tile = unsafe { tiles.tile_unchecked(1) }; + let mma = MmaSs::::new(); + unsafe { mma.gemm(0, &tile, &tile, false) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr new file mode 100644 index 0000000000..81d381fc45 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tile_array_role.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/sm100_attention_tile_array_role.rs:14:26 + | +14 | unsafe { mma.gemm(0, &tile, &tile, false) }; + | ---- ^^^^^ expected `OperandA`, found `OperandB` + | | + | arguments to this method are incorrect + | + = note: expected reference `&SharedTile, OperandA>` + found reference `&SharedTile, OperandB>` +note: method defined here + --> src/attention/sm100/mma.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs new file mode 100644 index 0000000000..f0324ee7b8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a TMA load whose destination has a different operand-major layout. +#![feature(f16)] +use cute_rs::attention::sm100::{SharedLayout, SharedTile, TmaCopy}; +use cute_rs::{OperandB, TensorMap}; +/// K-major 128-by-128 descriptor layout. +type Qk = SharedLayout<128, 128>; +/// MN-major 128-by-128 destination layout. +type Pv = SharedLayout<128, 128, true>; +/// Reject a TMA load whose destination has a different operand-major layout. +fn mismatch(desc: &TensorMap, tile: &mut SharedTile) { + unsafe { TmaCopy::::load(desc, tile, 0, 0, core::ptr::null_mut(), 0) }; +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr new file mode 100644 index 0000000000..edf1faf2e2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_tma_layout.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/sm100_attention_tma_layout.rs:13:43 + | +13 | unsafe { TmaCopy::::load(desc, tile, 0, 0, core::ptr::null_mut(), 0) }; + | ---------------------- ^^^^ expected `false`, found `true` + | | + | arguments to this function are incorrect + | + = note: expected mutable reference `&mut SharedTile, _>` + found mutable reference `&mut SharedTile, OperandB>` +note: associated function defined here + --> src/attention/sm100/tma.rs + | + | pub unsafe fn load( + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs new file mode 100644 index 0000000000..e619c74ae8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject K16 TMEM/shared MMA even when the B layout itself is valid. +use cute_rs::attention::sm100::{MmaTs, SharedLayout}; +// MN-major shared B alone permits K16, but the TMEM-source importer does not. +/// MN-major B layout valid on its own but unsupported by TMEM/shared K16 MMA. +type Value = SharedLayout<128, 16, true>; +/// Reject K16 TMEM/shared MMA even when the B layout itself is valid. +const _: MmaTs = MmaTs::new(); +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr new file mode 100644 index 0000000000..f283248157 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_attention_ts_k16.stderr @@ -0,0 +1,11 @@ +error[E0080]: evaluation panicked: assertion failed: K == 32 || (K > 0 && K.is_multiple_of(64)) + --> src/attention/sm100/mma.rs + | + | assert!(K == 32 || (K > 0 && K.is_multiple_of(64))); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::MmaTs::, 128, 128, 16, 1>::VALID` failed here + +note: erroneous constant encountered + --> src/attention/sm100/mma.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs new file mode 100644 index 0000000000..24f94f117c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.rs @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject swapped N192 and N160 shared B partitions. +#![feature(f16)] +use cute_rs::{ComposedLayout, RowMajor, Sm100TiledMma, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. +type A = ComposedLayout, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. +type B0 = ComposedLayout, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. +type B1 = ComposedLayout, 0, RowMajor<80, 64>>; +/// Reject swapped N192 and N160 shared B partitions. +const _: Sm100TiledMma = Sm100TiledMma::new(); +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr new file mode 100644 index 0000000000..73a9c66de9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_mma_shape_mismatch.stderr @@ -0,0 +1,11 @@ +error[E0080]: evaluation panicked: assertion failed: B0L::SHAPE.0 == (N0 / 2) as i64 && B0L::SHAPE.1 == K as i64 + --> src/sm100/mma.rs + | + | assert!(B0L::SHAPE.0 == (N0 / 2) as i64 && B0L::SHAPE.1 == K as i64); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::Sm100TiledMma::, 0, cute_rs::RowMajor<128, 64>>, cute_rs::ComposedLayout, 0, cute_rs::RowMajor<80, 64>>, cute_rs::ComposedLayout, 0, cute_rs::RowMajor<96, 64>>, 192, 160>::VALID` failed here + +note: erroneous constant encountered + --> src/sm100/mma.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs new file mode 100644 index 0000000000..bafe5fb2a5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a B-role tile supplied as the paired SM100 MMA A operand. +#![feature(f16)] +use cute_rs::{ComposedLayout, OperandB, RowMajor, Sm100SharedTile, Sm100TiledMma, Swizzle}; +/// Local 128-by-64 FP16 A tile with the B128 swizzle. +type A = ComposedLayout, 0, RowMajor<128, 64>>; +/// Local 96-by-64 B tile belonging to the collective N192 partition. +type B0 = ComposedLayout, 0, RowMajor<96, 64>>; +/// Local 80-by-64 B tile belonging to the collective N160 partition. +type B1 = ComposedLayout, 0, RowMajor<80, 64>>; +/// Reject a B-role tile supplied as the paired SM100 MMA A operand. +fn main() { + let a = unsafe { Sm100SharedTile::::from_raw(core::ptr::null_mut()) }; + let b0 = unsafe { Sm100SharedTile::::from_raw(core::ptr::null_mut()) }; + let b1 = unsafe { Sm100SharedTile::::from_raw(core::ptr::null_mut()) }; + let mma = Sm100TiledMma::::new(); + unsafe { + mma.gemm(0, &a, &b0, &b1, false); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr new file mode 100644 index 0000000000..6be516aba7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_operand_role_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/sm100_operand_role_mismatch.rs:19:21 + | +19 | mma.gemm(0, &a, &b0, &b1, false); + | ---- ^^ expected `OperandA`, found `OperandB` + | | + | arguments to this method are incorrect + | + = note: expected reference `&Sm100SharedTile, 0, RowMajor<128, 64>>, OperandA>` + found reference `&Sm100SharedTile, 0, RowMajor<128, 64>>, OperandB>` +note: method defined here + --> src/sm100/mma.rs + | + | pub unsafe fn gemm( + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs new file mode 100644 index 0000000000..1f166115cf --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.rs @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Compile-fail probe: reject a cluster TMA copy whose descriptor and destination shapes differ. +#![feature(f16)] +use cute_rs::{ + ComposedLayout, OperandA, RowMajor, Sm100ClusterTmaCopy, Sm100SharedTile, Swizzle, TensorMap, +}; +/// B128-swizzled 128-by-64 descriptor layout. +type A = ComposedLayout, 0, RowMajor<128, 64>>; +/// B128-swizzled 96-by-64 destination layout. +type B = ComposedLayout, 0, RowMajor<96, 64>>; +/// Reject a cluster TMA copy whose descriptor and destination shapes differ. +fn main() { + let desc: *const TensorMap = core::ptr::null(); + let mut dst = unsafe { Sm100SharedTile::::from_raw(core::ptr::null_mut()) }; + unsafe { + Sm100ClusterTmaCopy::::copy(desc, &mut dst, 0, 0, core::ptr::null_mut(), 0); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr new file mode 100644 index 0000000000..9913dc7216 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/sm100_tma_layout_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/sm100_tma_layout_mismatch.rs:17:46 + | +17 | Sm100ClusterTmaCopy::::copy(desc, &mut dst, 0, 0, core::ptr::null_mut(), 0); + | ------------------------------ ^^^^^^^^ expected `128`, found `96` + | | + | arguments to this function are incorrect + | + = note: expected mutable reference `&mut Sm100SharedTile, 0, RowMajor<128, 64>>, _>` + found mutable reference `&mut Sm100SharedTile, 0, RowMajor<96, 64>>, OperandA>` +note: associated function defined here + --> src/sm100/tma.rs + | + | pub unsafe fn copy( + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.rs new file mode 100644 index 0000000000..311b270545 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{ + markers::{C, Layout}, + static_layout::max_offset_2d, +}; +type Invalid = Layout<(C<4>, C<1>), (C<{ i64::MAX }>, C<1>)>; +const _: usize = max_offset_2d::(); +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.stderr new file mode 100644 index 0000000000..9bd8c82f82 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_2d_overflow.stderr @@ -0,0 +1,21 @@ +error[E0080]: evaluation panicked: static layout physical offset overflow + --> tests/compile_fail/static_layout_2d_overflow.rs:10:18 + | +10 | const _: usize = max_offset_2d::(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `_` failed inside this call + | +note: inside `max_offset_2d::, C<1>), (C, C<1>)>>` + --> src/static_layout.rs + | + | let row_max = checked_offset_mul(parameter(L::SHAPE.0 - 1), parameter(L::STRIDE.0)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: inside `static_layout::checked_offset_mul` + --> $RUST/core/src/panic.rs + | + | $crate::panicking::panic_fmt($crate::const_format_args!($($t)+)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the failure occurred here + | + ::: src/static_layout.rs + | + | None => panic!("static layout physical offset overflow"), + | ------------------------------------------------ in this macro invocation diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.rs new file mode 100644 index 0000000000..114a6acbc7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{ + markers::{C, ComposedLayout, Layout, Swizzle}, + static_layout::StaticLayout, +}; +type Invalid = ComposedLayout, 1, Layout, C<{ i64::MAX }>>>; +const _: usize = Invalid::SIZE; +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.stderr new file mode 100644 index 0000000000..15d2a10c90 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_composed_overflow.stderr @@ -0,0 +1,31 @@ +error[E0080]: evaluation panicked: static layout offset exceeds i64 + --> src/static_layout.rs + | + | swizzled_bound(checked_offset_add(Outer::MAX_OFFSET, parameter(O)), B, M, S); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `, 1, cute_rs::Layout, cute_rs::C>> as cute_rs::static_layout::StaticLayout>::MAX_OFFSET` failed inside this call + | +note: inside `static_layout::checked_offset_add` + --> $RUST/core/src/panic.rs + | + | $crate::panicking::panic_fmt($crate::const_format_args!($($t)+)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the failure occurred here + | + ::: src/static_layout.rs + | + | / assert!( + | | result as u128 <= i64::MAX as u128, + | | "static layout offset exceeds i64" + | | ); + | |_____- in this macro invocation + +note: erroneous constant encountered + --> src/static_layout.rs + | + | let _ = Self::MAX_OFFSET; + | ^^^^^^^^^^^^^^^^ + +note: erroneous constant encountered + --> tests/compile_fail/static_layout_composed_overflow.rs:10:18 + | +10 | const _: usize = Invalid::SIZE; + | ^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.rs new file mode 100644 index 0000000000..9c457af09f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{ + markers::{C, Layout}, + static_layout::StaticLayout, +}; +type Invalid = Layout, C<{ i64::MAX }>>; +const _: usize = Invalid::SIZE; +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.stderr new file mode 100644 index 0000000000..2d33e2e32e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_leaf_overflow.stderr @@ -0,0 +1,28 @@ +error[E0080]: evaluation panicked: static layout physical offset overflow + --> src/static_layout.rs + | + | checked_offset_mul(parameter(N - 1), parameter(S)) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `, cute_rs::C> as cute_rs::static_layout::StaticLayout>::MAX_OFFSET` failed inside this call + | +note: inside `static_layout::checked_offset_mul` + --> $RUST/core/src/panic.rs + | + | $crate::panicking::panic_fmt($crate::const_format_args!($($t)+)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the failure occurred here + | + ::: src/static_layout.rs + | + | None => panic!("static layout physical offset overflow"), + | ------------------------------------------------ in this macro invocation + +note: erroneous constant encountered + --> src/static_layout.rs + | + | let _ = Self::MAX_OFFSET; + | ^^^^^^^^^^^^^^^^ + +note: erroneous constant encountered + --> tests/compile_fail/static_layout_leaf_overflow.rs:10:18 + | +10 | const _: usize = Invalid::SIZE; + | ^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.rs new file mode 100644 index 0000000000..889cad361e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.rs @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::{ + markers::{C, Layout}, + static_layout::StaticLayout, +}; +type Invalid = Layout<(C<2>, C<2>), (C<{ i64::MAX }>, C<1>)>; +const _: usize = Invalid::SIZE; +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.stderr new file mode 100644 index 0000000000..7912c1777e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/static_layout_sum_overflow.stderr @@ -0,0 +1,31 @@ +error[E0080]: evaluation panicked: static layout offset exceeds i64 + --> src/static_layout.rs + | + | checked_offset_add(Layout::::MAX_OFFSET, Layout::::MAX_OFFSET); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `, cute_rs::C<2>), (cute_rs::C, cute_rs::C<1>)> as cute_rs::static_layout::StaticLayout>::MAX_OFFSET` failed inside this call + | +note: inside `static_layout::checked_offset_add` + --> $RUST/core/src/panic.rs + | + | $crate::panicking::panic_fmt($crate::const_format_args!($($t)+)); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the failure occurred here + | + ::: src/static_layout.rs + | + | / assert!( + | | result as u128 <= i64::MAX as u128, + | | "static layout offset exceeds i64" + | | ); + | |_____- in this macro invocation + +note: erroneous constant encountered + --> src/static_layout.rs + | + | let _ = Self::MAX_OFFSET; + | ^^^^^^^^^^^^^^^^ + +note: erroneous constant encountered + --> tests/compile_fail/static_layout_sum_overflow.rs:10:18 + | +10 | const _: usize = Invalid::SIZE; + | ^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.rs new file mode 100644 index 0000000000..0d9bbf71ab --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::tensor::partition::IndexLayout; +#[derive(Clone, Copy)] +struct Escaping; +impl IndexLayout for Escaping { + fn len(&self) -> usize { + 1 + } + fn crd2idx(&self, _: usize) -> usize { + usize::MAX + } +} +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.stderr new file mode 100644 index 0000000000..a533e23595 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_layout_is_sealed.stderr @@ -0,0 +1,32 @@ +error[E0277]: the trait bound `Escaping: partition::sealed::Sealed` is not satisfied + --> tests/compile_fail/strided_layout_is_sealed.rs:8:22 + | +8 | impl IndexLayout for Escaping { + | ^^^^^^^^ unsatisfied trait bound + | +help: the trait `partition::sealed::Sealed` is not implemented for `Escaping` + --> tests/compile_fail/strided_layout_is_sealed.rs:7:1 + | +7 | struct Escaping; + | ^^^^^^^^^^^^^^^ +help: the following other types implement trait `partition::sealed::Sealed` + --> src/tensor/partition.rs + | + | impl sealed::Sealed for Strided1D {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `Strided1D` +... + | impl sealed::Sealed for Subrange {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `Subrange` +... + | impl sealed::Sealed for Cyclic {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `Cyclic` +note: required by a bound in `IndexLayout` + --> src/tensor/partition.rs + | + | pub trait IndexLayout: sealed::Sealed + Copy { + | ^^^^^^^^^^^^^^ required by this bound in `IndexLayout` + = note: `IndexLayout` is a "sealed trait", because to implement it you also need to implement `cute_rs::tensor::partition::sealed::Sealed`, which is not accessible; this is usually done to force you to use one of the provided types that already implement it + = help: the following types implement the trait: + cute_rs::tensor::partition::Strided1D + cute_rs::tensor::partition::Subrange + cute_rs::tensor::partition::Cyclic diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.rs new file mode 100644 index 0000000000..75b6ccfb53 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::tensor::partition::StridedTensorMut; +fn main() { + let mut values = [0.0_f32; 4]; + let mut first = StridedTensorMut::from_mut_slice(&mut values); + let mut second = first; + unsafe { + first.store_at(0, 1.0); + } + unsafe { + second.store_at(1, 2.0); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.stderr new file mode 100644 index 0000000000..a7b5ccc4f0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_move.stderr @@ -0,0 +1,10 @@ +error[E0382]: borrow of moved value: `first` + --> tests/compile_fail/strided_tensor_mut_move.rs:11:9 + | + 8 | let mut first = StridedTensorMut::from_mut_slice(&mut values); + | --------- move occurs because `first` has type `StridedTensorMut<'_, f32>`, which does not implement the `Copy` trait + 9 | let mut second = first; + | ----- value moved here +10 | unsafe { +11 | first.store_at(0, 1.0); + | ^^^^^ value borrowed here after move diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.rs new file mode 100644 index 0000000000..8cda978659 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.rs @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +use cute_rs::tensor::partition::StridedTensorMut; +fn main() { + let mut values = [0.0_f32; 4]; + let mut parent = StridedTensorMut::from_mut_slice(&mut values); + let mut child = parent.reborrow(); + unsafe { + parent.store_at(0, 1.0); + } + unsafe { + child.store_at(1, 2.0); + } +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.stderr new file mode 100644 index 0000000000..e21e928e29 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_mut_reborrow.stderr @@ -0,0 +1,11 @@ +error[E0499]: cannot borrow `parent` as mutable more than once at a time + --> tests/compile_fail/strided_tensor_mut_reborrow.rs:11:9 + | + 9 | let mut child = parent.reborrow(); + | ------ first mutable borrow occurs here +10 | unsafe { +11 | parent.store_at(0, 1.0); + | ^^^^^^ second mutable borrow occurs here +... +14 | child.store_at(1, 2.0); + | ----- first borrow later used here diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.rs new file mode 100644 index 0000000000..94bb246900 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.rs @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#![feature(f16)] +use cute_rs::{ + BFloat16, + tensor::partition::{Strided1D, StridedTensor}, +}; +fn main() { + let _invalid: Option> = None; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.stderr new file mode 100644 index 0000000000..07bc48eb01 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/strided_tensor_storage_mismatch.stderr @@ -0,0 +1,43 @@ +error[E0277]: the trait bound `BFloat16: TensorElement` is not satisfied + --> tests/compile_fail/strided_tensor_storage_mismatch.rs:11:19 + | +11 | let _invalid: Option> = None; + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `TensorElement` is not implemented for `BFloat16` + | +help: the trait `TensorElement` is not implemented for `BFloat16` + but trait `TensorElement` is implemented for it + --> src/numeric.rs + | + | impl TensorElement for BFloat16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = help: for that trait implementation, expected `BFloat16`, found `f16` +note: required by a bound in `StridedTensor` + --> src/tensor/partition.rs + | + | pub struct StridedTensor<'a, Element, Layout = Strided1D, Storage = Element> + | ------------- required by a bound in this struct + | where + | Element: TensorElement, + | ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `StridedTensor` + +error[E0277]: the trait bound `BFloat16: TensorElement` is not satisfied + --> tests/compile_fail/strided_tensor_storage_mismatch.rs:11:73 + | +11 | let _invalid: Option> = None; + | ^^^^ the trait `TensorElement` is not implemented for `BFloat16` + | +help: the trait `TensorElement` is not implemented for `BFloat16` + but trait `TensorElement` is implemented for it + --> src/numeric.rs + | + | impl TensorElement for BFloat16 {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + = help: for that trait implementation, expected `BFloat16`, found `f16` +note: required by a bound in `StridedTensor` + --> src/tensor/partition.rs + | + | pub struct StridedTensor<'a, Element, Layout = Strided1D, Storage = Element> + | ------------- required by a bound in this struct + | where + | Element: TensorElement, + | ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `StridedTensor` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.rs new file mode 100644 index 0000000000..bd158c04b3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: keep writable tensor views non-Clone so their exclusive borrow cannot be duplicated. +#![feature(f16)] + +use cute_rs::{Contiguous1D, TensorMut}; + +/// Require Clone so passing a writable tensor view exposes the missing bound. + +fn require_clone() {} + +/// Keep writable tensor views non-Clone so their exclusive borrow cannot be duplicated. + +fn main() { + require_clone::>(); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.stderr new file mode 100644 index 0000000000..6867d505d9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_clone.stderr @@ -0,0 +1,11 @@ +error[E0277]: the trait bound `TensorMut<'static, f32, Contiguous1D>: Clone` is not satisfied + --> tests/compile_fail/tensor_mut_clone.rs:16:21 + | +16 | require_clone::>(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `Clone` is not implemented for `TensorMut<'static, f32, Contiguous1D>` + | +note: required by a bound in `require_clone` + --> tests/compile_fail/tensor_mut_clone.rs:11:21 + | +11 | fn require_clone() {} + | ^^^^^ required by this bound in `require_clone` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.rs new file mode 100644 index 0000000000..2cdd40d0f1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: keep writable tensor views non-Copy so their exclusive borrow cannot be duplicated. +#![feature(f16)] + +use cute_rs::{Contiguous1D, TensorMut}; + +/// Require Copy so passing a writable tensor view exposes the missing bound. + +fn require_copy() {} + +/// Keep writable tensor views non-Copy so their exclusive borrow cannot be duplicated. + +fn main() { + require_copy::>(); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.stderr new file mode 100644 index 0000000000..2d15b48954 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tensor_mut_copy.stderr @@ -0,0 +1,11 @@ +error[E0277]: the trait bound `TensorMut<'static, f32, Contiguous1D>: Copy` is not satisfied + --> tests/compile_fail/tensor_mut_copy.rs:16:20 + | +16 | require_copy::>(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `Copy` is not implemented for `TensorMut<'static, f32, Contiguous1D>` + | +note: required by a bound in `require_copy` + --> tests/compile_fail/tensor_mut_copy.rs:11:20 + | +11 | fn require_copy() {} + | ^^^^ required by this bound in `require_copy` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs new file mode 100644 index 0000000000..100e6acdf9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.rs @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject copying packed E2M1 values into a plain FP16 shared tensor. +#![feature(f16)] + +use core::ptr; + +use cute_rs::{CopyG2S, Mkl, Mxf4BlockScaledTensor, SharedTensor, TiledCopy}; + +/// Reject copying packed E2M1 values into a plain FP16 shared tensor. + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxf4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let source = matrix.values_for_copy(); + let mut wrong_encoding = unsafe { + SharedTensor::::from_raw_parts(ptr::null_mut(), 0) + }; + let copy = TiledCopy::, (), (), ()>::new(); + + unsafe { copy.get_slice(0).copy(&source, (0, 0), &mut wrong_encoding) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr new file mode 100644 index 0000000000..8686cafa94 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_encoding_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/tiled_copy_encoding_mismatch.rs:23:54 + | +23 | unsafe { copy.get_slice(0).copy(&source, (0, 0), &mut wrong_encoding) }; + | ---- ^^^^^^^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor` + | | + | arguments to this method are incorrect + | + = note: expected mutable reference `&mut SharedTensor` + found mutable reference `&mut SharedTensor` +note: method defined here + --> src/copy/tiled.rs + | + | pub unsafe fn copy< + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs new file mode 100644 index 0000000000..a348c305bc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.rs @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject an M-role source paired with an N-role shared destination. +#![feature(f16)] + +use core::ptr; + +use cute_rs::{ + CopyG2S, Mkl, Mxf4E2M1, Mxf4BlockScaledTensor, Nkl, SharedTensor, TiledCopy, +}; + +/// Reject an M-role source paired with an N-role shared destination. + +fn main() { + let values = [f16::from_bits(0); 32]; + let scales = [127u8; 2]; + let matrix = Mxf4BlockScaledTensor::::from_slices(&values, &scales, 1, 64); + let source = matrix.values_for_copy(); + let mut shared_b = unsafe { + SharedTensor::::from_raw_parts(ptr::null_mut(), 0) + }; + let copy = TiledCopy::, (), (), ()>::new(); + + unsafe { copy.get_slice(0).copy(&source, (0, 0), &mut shared_b) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr new file mode 100644 index 0000000000..ffb2722a31 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tiled_copy_role_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/tiled_copy_role_mismatch.rs:25:54 + | +25 | unsafe { copy.get_slice(0).copy(&source, (0, 0), &mut shared_b) }; + | ---- ^^^^^^^^^^^^^ expected `&mut SharedTensor`, found `&mut SharedTensor<_, f16, (), Nkl>` + | | + | arguments to this method are incorrect + | + = note: expected mutable reference `&mut SharedTensor` + found mutable reference `&mut SharedTensor` +note: method defined here + --> src/copy/tiled.rs + | + | pub unsafe fn copy< + | ^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs new file mode 100644 index 0000000000..57f911ea3f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: require an unsafe block for device-only descriptor prefetch. +use cute_rs::TensorMap; + +/// Require an unsafe block for device-only descriptor prefetch. + +fn prefetch(descriptor: &TensorMap) { + descriptor.prefetch_descriptor(); +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr new file mode 100644 index 0000000000..dfbcd6b598 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `TensorMap::::prefetch_descriptor` is unsafe and requires unsafe block + --> tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs:12:5 + | +12 | descriptor.prefetch_descriptor(); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs new file mode 100644 index 0000000000..c31dd4eb46 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.rs @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject a TMA store whose descriptor and source have different logical shapes. +#![feature(f16)] + +use core::{marker::PhantomData, ptr}; + +use cute_rs::{RowMajor, SmemTile, TensorMap, copy_tma_s2g_2d}; + +/// Descriptor shape with eight rows and eight contiguous elements. + +type DescriptorLayout = RowMajor<8, 8>; +/// Equal-capacity source with four rows and sixteen contiguous elements. +type SourceLayout = RowMajor<4, 16>; + +/// Reject a TMA store whose descriptor and source have different logical shapes. + +fn main() { + let descriptor = ptr::null::>(); + let source = SmemTile:: { + base: ptr::null_mut(), + capacity: 64, + layout: PhantomData, + }; + + unsafe { copy_tma_s2g_2d(descriptor, (0, 0), source) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr new file mode 100644 index 0000000000..62795a8f5f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/tma_store_layout_mismatch.stderr @@ -0,0 +1,15 @@ +error[E0308]: mismatched types + --> tests/compile_fail/tma_store_layout_mismatch.rs:27:50 + | +27 | unsafe { copy_tma_s2g_2d(descriptor, (0, 0), source) }; + | --------------- ^^^^^^ expected `8`, found `4` + | | + | arguments to this function are incorrect + | + = note: expected struct `SmemTile>` + found struct `SmemTile>` +note: function defined here + --> src/copy/tma/mod.rs + | + | pub unsafe fn copy_tma_s2g_2d( + | ^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs new file mode 100644 index 0000000000..4d9698ddf0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.rs @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject consumer cursors on producer operations and producer cursors on consumer operations. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +/// Reject consumer cursors on producer operations and producer cursors on consumer operations. + +unsafe fn wrong_roles(ring: TwoCtaRing<15, AsyncCompletion>) { + unsafe { + ring.acquire_cursor(RingCursor::<15, Consumer>::new()); + ring.wait_cursor(RingCursor::<15, Producer>::new()); + ring.mma_full_cursor(RingCursor::<15, Consumer>::new()); + ring.mma_empty_cursor(RingCursor::<15, Producer>::new()); + } +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr new file mode 100644 index 0000000000..c7fa3becd4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_role.stderr @@ -0,0 +1,63 @@ +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:15:29 + | +15 | ring.acquire_cursor(RingCursor::<15, Consumer>::new()); + | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Producer>`, found `RingCursor<15, Consumer>` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Producer>` + found struct `RingCursor<15, Consumer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn acquire_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:16:26 + | +16 | ring.wait_cursor(RingCursor::<15, Producer>::new()); + | ----------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Consumer>`, found `RingCursor<15, Producer>` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Consumer>` + found struct `RingCursor<15, Producer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn wait_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:17:30 + | +17 | ring.mma_full_cursor(RingCursor::<15, Consumer>::new()); + | --------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Producer>`, found `RingCursor<15, Consumer>` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Producer>` + found struct `RingCursor<15, Consumer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_full_cursor(self, cursor: RingCursor) + | ^^^^^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_role.rs:18:31 + | +18 | ring.mma_empty_cursor(RingCursor::<15, Producer>::new()); + | ---------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `RingCursor<15, Consumer>`, found `RingCursor<15, Producer>` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Consumer>` + found struct `RingCursor<15, Producer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_empty_cursor(self, cursor: RingCursor) + | ^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs new file mode 100644 index 0000000000..d1c7821fca --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject cursors whose stage count differs from the ring. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +/// Reject cursors whose stage count differs from the ring. + +unsafe fn wrong_stages(ring: TwoCtaRing<15, AsyncCompletion>) { + unsafe { + ring.acquire_cursor(RingCursor::<2, Producer>::new()); + ring.wait_cursor(RingCursor::<2, Consumer>::new()); + } +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr new file mode 100644 index 0000000000..7d24b03e11 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_wrong_stages.stderr @@ -0,0 +1,31 @@ +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:15:29 + | +15 | ring.acquire_cursor(RingCursor::<2, Producer>::new()); + | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `15`, found `2` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Producer>` + found struct `RingCursor<2, Producer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn acquire_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^^^^ + +error[E0308]: mismatched types + --> tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs:16:26 + | +16 | ring.wait_cursor(RingCursor::<2, Consumer>::new()); + | ----------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `15`, found `2` + | | + | arguments to this method are incorrect + | + = note: expected struct `RingCursor<15, Consumer>` + found struct `RingCursor<2, Consumer>` +note: method defined here + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn wait_cursor(self, cursor: RingCursor) { + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs new file mode 100644 index 0000000000..f29abcc55f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.rs @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject a zero-stage cursor during const evaluation. +use cute_rs::attention::sm100::pipeline::{Producer, RingCursor}; + +/// Reject a zero-stage cursor during const evaluation. + +const _: RingCursor<0, Producer> = RingCursor::new(); +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr new file mode 100644 index 0000000000..c6ad8e625e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_cursor_zero_stages.stderr @@ -0,0 +1,11 @@ +error[E0080]: evaluation panicked: assertion failed: STAGES > 0 + --> src/attention/sm100/pipeline.rs + | + | const VALID: () = assert!(STAGES > 0); + | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::pipeline::RingCursor::<0, cute_rs::Producer>::VALID` failed here + +note: erroneous constant encountered + --> src/attention/sm100/pipeline.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs new file mode 100644 index 0000000000..511864086c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject completion signals unavailable for async-copy reuse and local-shared policies. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, AsyncCopyReuse, Consumer, LocalShared, RingCursor, TwoCtaRing, +}; + +/// Reject completion signals unavailable for async-copy reuse and local-shared policies. + +unsafe fn wrong_non_tc_signals( + copies: TwoCtaRing<15, AsyncCompletion, AsyncCopyReuse>, + reuse: TwoCtaRing<2, AsyncCopyReuse>, + pages: TwoCtaRing<4, LocalShared>, +) { + unsafe { + copies.thread_empty(0); + reuse.producer_expect_tx(0, 16); + pages.producer_expect_tx(0, 16); + pages.mma_full(0); + pages.mma_empty_cursor(RingCursor::<4, Consumer>::new()); + pages.thread_full(0); + } +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr new file mode 100644 index 0000000000..831db8f291 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_non_tc_wrong_signal.stderr @@ -0,0 +1,143 @@ +error[E0277]: the trait bound `AsyncCopyReuse: SignalPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:19:16 + | +19 | copies.thread_empty(0); + | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCopyReuse` + | +help: the following other types implement trait `SignalPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_empty` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn thread_empty(self, ticket: u32) + | ------------ required by a bound in this associated function + | where + | Empty: SignalPolicy, + | ^^^^^^^^^^^^ required by this bound in `TwoCtaRing::::thread_empty` + +error[E0277]: the trait bound `AsyncCopyReuse: TmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:20:15 + | +20 | reuse.producer_expect_tx(0, 16); + | ^^^^^^^^^^^^^^^^^^ the trait `TmaPolicy` is not implemented for `AsyncCopyReuse` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/attention/sm100/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::producer_expect_tx` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn producer_expect_tx(self, ticket: u32, bytes: u32) + | ------------------ required by a bound in this associated function + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::producer_expect_tx` + +error[E0277]: the trait bound `LocalShared: TmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:21:15 + | +21 | pages.producer_expect_tx(0, 16); + | ^^^^^^^^^^^^^^^^^^ the trait `TmaPolicy` is not implemented for `LocalShared` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/attention/sm100/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::producer_expect_tx` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn producer_expect_tx(self, ticket: u32, bytes: u32) + | ------------------ required by a bound in this associated function + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::producer_expect_tx` + +error[E0277]: the trait bound `LocalShared: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:22:15 + | +22 | pages.mma_full(0); + | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalShared` + | +help: the following other types implement trait `MmaPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl MmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCompletion` +... + | impl MmaPolicy for AsyncCopyReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCopyReuse` +... + | impl MmaPolicy for TensorMemory {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl MmaPolicy for SharedReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedReuse` +note: required by a bound in `TwoCtaRing::::mma_full` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_full(self, ticket: u32) + | -------- required by a bound in this associated function + | where + | Full: MmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::mma_full` + +error[E0277]: the trait bound `LocalShared: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:23:15 + | +23 | pages.mma_empty_cursor(RingCursor::<4, Consumer>::new()); + | ^^^^^^^^^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalShared` + | +help: the following other types implement trait `MmaPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl MmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCompletion` +... + | impl MmaPolicy for AsyncCopyReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCopyReuse` +... + | impl MmaPolicy for TensorMemory {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl MmaPolicy for SharedReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedReuse` +note: required by a bound in `TwoCtaRing::::mma_empty_cursor` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_empty_cursor(self, cursor: RingCursor) + | ---------------- required by a bound in this associated function + | where + | Empty: MmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::mma_empty_cursor` + +error[E0277]: the trait bound `LocalShared: SignalPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs:24:15 + | +24 | pages.thread_full(0); + | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `LocalShared` + | +help: the following other types implement trait `SignalPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_full` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn thread_full(self, ticket: u32) + | ----------- required by a bound in this associated function + | where + | Full: SignalPolicy, + | ^^^^^^^^^^^^ required by this bound in `TwoCtaRing::::thread_full` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs new file mode 100644 index 0000000000..bff31594c6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.rs @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: prevent external code from defining a wait policy that omits required ordering. +use cuda_device::barrier::Barrier; +use cute_rs::attention::sm100::pipeline::WaitPolicy; + +/// Invalid external wait policy with no ordering guarantees. + +#[derive(Clone, Copy)] +struct Unordered; + +impl WaitPolicy for Unordered { + unsafe fn wait(_: *const Barrier, _: u32) {} +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr new file mode 100644 index 0000000000..47f2203d7d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_policy_sealed.stderr @@ -0,0 +1,33 @@ +error[E0277]: the trait bound `Unordered: cute_rs::attention::sm100::pipeline::sealed::Wait` is not satisfied + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:15:21 + | +15 | impl WaitPolicy for Unordered { + | ^^^^^^^^^ unsatisfied trait bound + | +help: the trait `cute_rs::attention::sm100::pipeline::sealed::Wait` is not implemented for `Unordered` + --> tests/compile_fail/two_cta_ring_policy_sealed.rs:13:1 + | +13 | struct Unordered; + | ^^^^^^^^^^^^^^^^ + = help: the following other types implement trait `cute_rs::attention::sm100::pipeline::sealed::Wait`: + AsyncCompletion + AsyncCopyReuse + LocalMemory + LocalShared + SharedReuse + SharedToAsync + cute_rs::attention::sm100::pipeline::TensorMemory +note: required by a bound in `WaitPolicy` + --> src/attention/sm100/pipeline.rs + | + | pub trait WaitPolicy: sealed::Wait + Copy { + | ^^^^^^^^^^^^ required by this bound in `WaitPolicy` + = note: `WaitPolicy` is a "sealed trait", because to implement it you also need to implement `cute_rs::attention::sm100::pipeline::sealed::Wait`, which is not accessible; this is usually done to force you to use one of the provided types that already implement it + = help: the following types implement the trait: + cute_rs::attention::sm100::pipeline::LocalMemory + cute_rs::attention::sm100::pipeline::LocalShared + cute_rs::attention::sm100::pipeline::AsyncCompletion + cute_rs::attention::sm100::pipeline::AsyncCopyReuse + cute_rs::attention::sm100::pipeline::TensorMemory + cute_rs::attention::sm100::pipeline::SharedToAsync + cute_rs::attention::sm100::pipeline::SharedReuse diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs new file mode 100644 index 0000000000..4e25a4c417 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject signaling mechanisms incompatible with each ring policy. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, LocalMemory, SharedReuse, SharedToAsync, TensorMemory, TwoCtaRing, +}; + +/// Reject signaling mechanisms incompatible with each ring policy. + +unsafe fn wrong_completion_mechanisms( + tensor: TwoCtaRing<2, TensorMemory>, + shared: TwoCtaRing<2, SharedToAsync, SharedReuse>, + copies: TwoCtaRing<2, AsyncCompletion>, + local: TwoCtaRing<2, LocalMemory>, +) { + unsafe { + tensor.producer_expect_tx(0, 16); + shared.mma_full(0); + shared.thread_empty(0); + copies.thread_full(0); + local.mma_empty(0); + } +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr new file mode 100644 index 0000000000..52fec6c5be --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_wrong_signal.stderr @@ -0,0 +1,123 @@ +error[E0277]: the trait bound `TensorMemory: TmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:20:16 + | +20 | tensor.producer_expect_tx(0, 16); + | ^^^^^^^^^^^^^^^^^^ the trait `TmaPolicy` is not implemented for `TensorMemory` + | +help: the trait `TmaPolicy` is implemented for `AsyncCompletion` + --> src/attention/sm100/pipeline.rs + | + | impl TmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +note: required by a bound in `TwoCtaRing::::producer_expect_tx` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn producer_expect_tx(self, ticket: u32, bytes: u32) + | ------------------ required by a bound in this associated function + | where + | Full: TmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::producer_expect_tx` + +error[E0277]: the trait bound `SharedToAsync: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:21:16 + | +21 | shared.mma_full(0); + | ^^^^^^^^ the trait `MmaPolicy` is not implemented for `SharedToAsync` + | +help: the following other types implement trait `MmaPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl MmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCompletion` +... + | impl MmaPolicy for AsyncCopyReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCopyReuse` +... + | impl MmaPolicy for TensorMemory {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl MmaPolicy for SharedReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedReuse` +note: required by a bound in `TwoCtaRing::::mma_full` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_full(self, ticket: u32) + | -------- required by a bound in this associated function + | where + | Full: MmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::mma_full` + +error[E0277]: the trait bound `SharedReuse: SignalPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:22:16 + | +22 | shared.thread_empty(0); + | ^^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `SharedReuse` + | +help: the following other types implement trait `SignalPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_empty` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn thread_empty(self, ticket: u32) + | ------------ required by a bound in this associated function + | where + | Empty: SignalPolicy, + | ^^^^^^^^^^^^ required by this bound in `TwoCtaRing::::thread_empty` + +error[E0277]: the trait bound `AsyncCompletion: SignalPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:23:16 + | +23 | copies.thread_full(0); + | ^^^^^^^^^^^ the trait `SignalPolicy` is not implemented for `AsyncCompletion` + | +help: the following other types implement trait `SignalPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl SignalPolicy for TensorMemory { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl SignalPolicy for SharedToAsync { + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedToAsync` +note: required by a bound in `TwoCtaRing::::thread_full` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn thread_full(self, ticket: u32) + | ----------- required by a bound in this associated function + | where + | Full: SignalPolicy, + | ^^^^^^^^^^^^ required by this bound in `TwoCtaRing::::thread_full` + +error[E0277]: the trait bound `LocalMemory: MmaPolicy` is not satisfied + --> tests/compile_fail/two_cta_ring_wrong_signal.rs:24:15 + | +24 | local.mma_empty(0); + | ^^^^^^^^^ the trait `MmaPolicy` is not implemented for `LocalMemory` + | +help: the following other types implement trait `MmaPolicy` + --> src/attention/sm100/pipeline.rs + | + | impl MmaPolicy for AsyncCompletion {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCompletion` +... + | impl MmaPolicy for AsyncCopyReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `AsyncCopyReuse` +... + | impl MmaPolicy for TensorMemory {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `TensorMemory` +... + | impl MmaPolicy for SharedReuse {} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `SharedReuse` +note: required by a bound in `TwoCtaRing::::mma_empty` + --> src/attention/sm100/pipeline.rs + | + | pub unsafe fn mma_empty(self, ticket: u32) + | --------- required by a bound in this associated function + | where + | Empty: MmaPolicy, + | ^^^^^^^^^ required by this bound in `TwoCtaRing::::mma_empty` diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs new file mode 100644 index 0000000000..71407906d4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.rs @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-fail probe: reject a zero-stage ring during const evaluation. +use cute_rs::attention::sm100::pipeline::{AsyncCompletion, TwoCtaRing}; + +/// Reject a zero-stage ring during const evaluation. + +const _: TwoCtaRing<0, AsyncCompletion> = unsafe { TwoCtaRing::from_raw(core::ptr::null_mut(), 0) }; + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr new file mode 100644 index 0000000000..7d26893413 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/two_cta_ring_zero_stages.stderr @@ -0,0 +1,11 @@ +error[E0080]: evaluation panicked: assertion failed: STAGES > 0 + --> src/attention/sm100/pipeline.rs + | + | const VALID: () = assert!(STAGES > 0); + | ^^^^^^^^^^^^^^^^^^^ evaluation of `cute_rs::attention::sm100::pipeline::TwoCtaRing::<0, cute_rs::attention::sm100::pipeline::AsyncCompletion>::VALID` failed here + +note: erroneous constant encountered + --> src/attention/sm100/pipeline.rs + | + | let () = Self::VALID; + | ^^^^^^^^^^^ diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.rs new file mode 100644 index 0000000000..ea89bcbd2c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.rs @@ -0,0 +1,7 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +fn main() { + let _ = cute_rs::arch::warp_reduction_sum_f32(1.0); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.stderr new file mode 100644 index 0000000000..906ad8aa5c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_sum_requires_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `warp_reduction_sum_f32` is unsafe and requires unsafe block + --> tests/compile_fail/warp_sum_requires_unsafe.rs:6:13 + | +6 | let _ = cute_rs::arch::warp_reduction_sum_f32(1.0); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs new file mode 100644 index 0000000000..30e633ac6c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: require an unsafe assertion before treating a value as warp-uniform. +/// Require an unsafe assertion before treating a value as warp-uniform. +fn main() { + let _ = cute_rs::arch::make_warp_uniform(0); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr new file mode 100644 index 0000000000..1706ac8634 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_unsafe.stderr @@ -0,0 +1,7 @@ +error[E0133]: call to unsafe function `make_warp_uniform` is unsafe and requires unsafe block + --> tests/compile_fail/warp_uniform_unsafe.rs:7:13 + | +7 | let _ = cute_rs::arch::make_warp_uniform(0); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ call to unsafe function + | + = note: consult the function's documentation for information on how to avoid undefined behavior diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.rs new file mode 100644 index 0000000000..bddb090ab7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Compile-fail probe: reject a 64-bit value where the uniformity marker accepts only 32 bits. +/// Reject a 64-bit value where the uniformity marker accepts only 32 bits. +fn main() { + let _ = unsafe { cute_rs::arch::make_warp_uniform(0u64) }; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.stderr b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.stderr new file mode 100644 index 0000000000..fb94ea17f0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_fail/warp_uniform_width.stderr @@ -0,0 +1,18 @@ +error[E0308]: mismatched types + --> tests/compile_fail/warp_uniform_width.rs:7:55 + | +7 | let _ = unsafe { cute_rs::arch::make_warp_uniform(0u64) }; + | -------------------------------- ^^^^ expected `u32`, found `u64` + | | + | arguments to this function are incorrect + | +note: function defined here + --> src/arch.rs + | + | pub unsafe fn make_warp_uniform(value: u32) -> u32 { + | ^^^^^^^^^^^^^^^^^ +help: change the type of the numeric literal from `u64` to `u32` + | +7 - let _ = unsafe { cute_rs::arch::make_warp_uniform(0u64) }; +7 + let _ = unsafe { cute_rs::arch::make_warp_uniform(0u32) }; + | diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs new file mode 100644 index 0000000000..79cfe8f0ca --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/tma_descriptor_prefetch.rs @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-pass probe: type-check descriptor prefetch for arbitrary element and layout markers. +use cute_rs::TensorMap; + +// Prefetching metadata preserves the existing element and layout types and +// requires no extra layout trait or conversion at the call site. +/// Type-check descriptor prefetch for arbitrary element and layout markers. This device-only probe is never executed. +unsafe fn prefetch(descriptor: &TensorMap) { + unsafe { descriptor.prefetch_descriptor() }; +} + +/// Type-check descriptor prefetch for arbitrary element and layout markers. + +fn main() { + /// Layout marker with no extra traits, proving prefetch adds no layout bounds. + struct Layout; + let _: unsafe fn(&TensorMap) = prefetch; + let _: unsafe fn(&TensorMap) = prefetch; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs new file mode 100644 index 0000000000..5cf3404495 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_non_tc_receivers.rs @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-pass probe: type-check raw and cursor operations supported by async-copy and local-shared receivers. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, AsyncCopyReuse, Consumer, LocalShared, Producer, RingCursor, TwoCtaRing, +}; + +/// Type-check raw and cursor operations supported by async-copy and local-shared receivers. This device-only probe is never executed. + +unsafe fn non_tc_receivers( + copies: TwoCtaRing<15, AsyncCompletion, AsyncCopyReuse>, + pages: TwoCtaRing<4, LocalShared>, +) { + let producer = RingCursor::<15, Producer>::new(); + let consumer = RingCursor::<15, Consumer>::new(); + let page_producer = RingCursor::<4, Producer>::new(); + let page_consumer = RingCursor::<4, Consumer>::new(); + unsafe { + copies.acquire(0); + copies.producer_expect_tx(0, 16); + copies.wait(0); + copies.mma_empty(0); + copies.acquire_cursor(producer); + copies.producer_expect_tx_cursor(producer, 16); + copies.wait_cursor(consumer); + copies.mma_empty_cursor(consumer); + pages.acquire(0); + pages.wait(0); + pages.acquire_cursor(page_producer); + pages.wait_cursor(page_consumer); + let _ = pages.full_at_cursor(page_producer); + let _ = pages.empty_at_cursor(page_consumer); + } +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs new file mode 100644 index 0000000000..92c248c7f8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_cursor.rs @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-pass probe: type-check producer and consumer cursors with matching ring policies and stage counts. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, SharedReuse, SharedToAsync, TensorMemory, + TwoCtaRing, +}; + +/// Type-check producer and consumer cursors with matching ring policies and stage counts. This device-only probe is never executed. + +unsafe fn valid_cursors( + copies: TwoCtaRing<15, AsyncCompletion>, + tensor: TwoCtaRing<2, TensorMemory>, + shared: TwoCtaRing<2, SharedToAsync, SharedReuse>, +) { + let producer = RingCursor::<15, Producer>::new(); + let consumer = RingCursor::<15, Consumer>::new(); + let tensor_producer = RingCursor::<2, Producer>::new(); + let tensor_consumer = RingCursor::<2, Consumer>::new(); + unsafe { + copies.acquire_cursor(producer); + copies.producer_expect_tx_cursor(producer, 16); + copies.mma_full_cursor(producer); + copies.wait_cursor(consumer); + copies.mma_empty_cursor(consumer); + tensor.thread_full_cursor(tensor_producer); + tensor.thread_empty_cursor(tensor_consumer); + shared.thread_full_cursor(tensor_producer); + shared.mma_empty_cursor(tensor_consumer); + } +} +/// Unused host entry point; the device or rejection probe is only type-checked. +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs new file mode 100644 index 0000000000..69eb604114 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/compile_pass/two_cta_ring_valid_signal.rs @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Compile-pass probe: type-check the completion signals permitted by each ring policy. +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, SharedReuse, SharedToAsync, TensorMemory, TwoCtaRing, +}; + +/// Type-check the completion signals permitted by each ring policy. This device-only probe is never executed. + +unsafe fn valid_completion_mechanisms( + tensor: TwoCtaRing<2, TensorMemory>, + shared: TwoCtaRing<2, SharedToAsync, SharedReuse>, + copies: TwoCtaRing<2, AsyncCompletion>, +) { + unsafe { + tensor.mma_full(0); + tensor.mma_empty(0); + tensor.thread_full(0); + tensor.thread_empty(0); + shared.thread_full(0); + shared.mma_empty(0); + copies.producer_expect_tx(0, 16); + copies.mma_full(0); + copies.mma_empty(0); + } +} + +/// Unused host entry point; the device or rejection probe is only type-checked. + +fn main() {} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/decoder.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/decoder.rs new file mode 100644 index 0000000000..8350f1d5c1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/decoder.rs @@ -0,0 +1,154 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Decoder geometry and embedding indices can be checked without a GPU. +#![cfg(feature = "host")] + +use cute_rs::attention::AttentionError; +use cute_rs::decoder::{DecoderError, DecoderShape, EmbeddingShape, TokenIds}; +use cute_rs::ffn::{FfnError, RowMajorShape}; + +#[test] +fn talker_and_predictor_keep_hidden_query_and_kv_widths_distinct() { + let talker = DecoderShape::<2048, 6144, 16, 8, 128>::new(3, 2051).unwrap(); + let predictor = DecoderShape::<1024, 3072, 16, 8, 128>::new(3, 1024).unwrap(); + assert_eq!(talker.input(), RowMajorShape::new(3, 2048, 2051).unwrap()); + assert_eq!(talker.hidden(), RowMajorShape::dense(3, 2048).unwrap()); + assert_eq!(predictor.hidden(), RowMajorShape::dense(3, 1024).unwrap()); + assert_eq!(talker.query(), predictor.query()); + assert_eq!(talker.key_value(), predictor.key_value()); + assert_eq!(predictor.query(), RowMajorShape::dense(3, 2048).unwrap()); + assert_eq!( + predictor.key_value(), + RowMajorShape::dense(3, 1024).unwrap() + ); + assert_eq!(talker.query_norm(), RowMajorShape::dense(48, 128).unwrap()); + assert_eq!(talker.key_norm(), RowMajorShape::dense(24, 128).unwrap()); + assert_eq!( + talker.intermediate(), + RowMajorShape::dense(3, 6144).unwrap() + ); + assert_eq!( + predictor.intermediate(), + RowMajorShape::dense(3, 3072).unwrap() + ); +} + +#[test] +fn empty_tokens_preserve_geometry_but_require_no_activation_storage() { + let shape = DecoderShape::<1024, 3072, 16, 8, 128>::new(0, 1025).unwrap(); + for matrix in [ + shape.input(), + shape.hidden(), + shape.query(), + shape.key_value(), + shape.query_norm(), + shape.key_norm(), + shape.intermediate(), + ] { + assert_eq!(matrix.required_len(), 0); + } + assert_eq!( + ( + shape.tokens(), + shape.query_heads(), + shape.kv_heads(), + shape.head_dim() + ), + (0, 16, 8, 128) + ); +} + +#[test] +fn invalid_head_grouping_and_matrix_extents_fail_before_allocation() { + assert_eq!( + DecoderShape::<8, 16, 3, 2, 4>::new(1, 8), + Err(DecoderError::Attention(AttentionError::InvalidHeadGrouping)) + ); + assert_eq!( + DecoderShape::<8, 16, 0, 2, 4>::new(1, 8), + Err(DecoderError::Attention(AttentionError::ZeroDimension)) + ); + assert_eq!( + DecoderShape::<0, 16, 4, 2, 4>::new(0, 8), + Err(DecoderError::Matrix(FfnError::ZeroWidth)) + ); + assert_eq!( + DecoderShape::<8, 0, 4, 2, 4>::new(0, 8), + Err(DecoderError::Matrix(FfnError::ZeroWidth)) + ); + assert_eq!( + DecoderShape::<8, 16, 4, 2, 4>::new(1, 7), + Err(DecoderError::Matrix(FfnError::StrideTooSmall)) + ); + assert!(DecoderShape::<1, 1, 4, 2, 4>::new(isize::MAX as usize / 2, 1).is_err()); + assert!(DecoderShape::<1, 1, { usize::MAX }, 1, 2>::new(0, 1).is_err()); +} + +#[test] +fn projection_widths_are_checked_against_head_geometry() { + let shape = DecoderShape::<1024, 3072, 16, 8, 128>::new(1, 1024).unwrap(); + assert_eq!(shape.validate_projection_widths::<2048, 1024>(), Ok(())); + assert_eq!( + shape.validate_projection_widths::<1024, 1024>(), + Err(DecoderError::ProjectionWidthMismatch) + ); + assert_eq!( + shape.validate_projection_widths::<2048, 512>(), + Err(DecoderError::ProjectionWidthMismatch) + ); +} + +#[test] +fn embedding_tables_keep_vocab_width_and_pitch_independent() { + let shape = EmbeddingShape::<17, 2048>::new(2050).unwrap(); + assert_eq!(shape.table(), RowMajorShape::new(17, 2048, 2050).unwrap()); + assert_eq!( + shape.output(3).unwrap(), + RowMajorShape::dense(3, 2048).unwrap() + ); + assert_eq!(shape.output(0).unwrap().required_len(), 0); + assert_eq!( + EmbeddingShape::<0, 2048>::new(2048), + Err(DecoderError::EmptyVocabulary) + ); + assert_eq!( + EmbeddingShape::<17, 0>::new(0), + Err(DecoderError::Matrix(FfnError::ZeroWidth)) + ); + assert_eq!( + EmbeddingShape::<17, 2048>::new(2047), + Err(DecoderError::Matrix(FfnError::StrideTooSmall)) + ); + assert!(EmbeddingShape::<{ usize::MAX }, 2>::new(2).is_err()); + assert!(shape.output(usize::MAX).is_err()); +} + +#[test] +fn token_ids_validate_the_exact_vocab_boundary_and_retain_the_borrow() { + let ids = [0, 16, 3, 16]; + let checked = TokenIds::<17>::new(&ids).unwrap(); + assert_eq!(checked.as_slice(), &ids); + assert_eq!(checked.len(), 4); + assert!(!checked.is_empty()); + assert!(TokenIds::<17>::new(&[]).unwrap().is_empty()); + assert_eq!( + TokenIds::<17>::new(&[0, 17, 16]), + Err(DecoderError::InvalidToken { + index: 1, + token: 17, + vocabulary: 17 + }) + ); + assert_eq!( + TokenIds::<17>::new(&[u32::MAX]), + Err(DecoderError::InvalidToken { + index: 0, + token: u32::MAX, + vocabulary: 17 + }) + ); + assert_eq!(TokenIds::<0>::new(&[]), Err(DecoderError::EmptyVocabulary)); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/ffn.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/ffn.rs new file mode 100644 index 0000000000..31149b07cd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/ffn.rs @@ -0,0 +1,97 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Checked BF16 matrix extents and FFN dimensions require no device allocation. +#![cfg(feature = "host")] + +use cute_rs::ffn::{FfnError, FfnShape, RowMajorShape}; + +#[test] +fn dense_and_padded_rows_preserve_logical_shape_and_minimum_span() { + let dense = RowMajorShape::dense(3, 5).unwrap(); + assert_eq!((dense.rows(), dense.cols(), dense.row_stride()), (3, 5, 5)); + assert_eq!(dense.required_len(), 15); + let padded = RowMajorShape::new(3, 5, 8).unwrap(); + assert_eq!(padded.required_len(), 21); + assert_eq!(padded.validate_len(20), Err(FfnError::BufferTooShort)); + assert_eq!(padded.validate_len(21), Ok(())); + assert_eq!(padded.validate_len(24), Ok(())); +} + +#[test] +fn zero_rows_need_no_storage_but_width_and_stride_remain_valid() { + let empty = RowMajorShape::new(0, 7, 9).unwrap(); + assert_eq!(empty.required_len(), 0); + assert_eq!(empty.validate_len(0), Ok(())); + assert_eq!(empty.validate_len(9), Ok(())); + assert_eq!(RowMajorShape::new(0, 0, 0), Err(FfnError::ZeroWidth)); + assert_eq!(RowMajorShape::new(0, 7, 6), Err(FfnError::StrideTooSmall)); +} + +#[test] +fn strided_shapes_reject_short_pitch_and_overflow_before_allocation() { + assert_eq!(RowMajorShape::new(2, 9, 8), Err(FfnError::StrideTooSmall)); + assert_eq!( + RowMajorShape::new(3, 1, usize::MAX), + Err(FfnError::ExtentOverflow) + ); + assert_eq!( + RowMajorShape::new(2, 2, usize::MAX), + Err(FfnError::ExtentOverflow) + ); + assert_eq!( + RowMajorShape::dense(usize::MAX, 2), + Err(FfnError::ExtentOverflow) + ); +} + +#[test] +fn bf16_slice_byte_extent_must_fit_isize_even_when_element_count_fits() { + let max_elements = isize::MAX as usize / 2; + let largest = RowMajorShape::dense(1, max_elements).unwrap(); + assert_eq!(largest.required_len(), max_elements); + assert_eq!(largest.validate_len(max_elements), Ok(())); + assert_eq!( + largest.validate_len(max_elements + 1), + Err(FfnError::ExtentOverflow) + ); + assert_eq!( + RowMajorShape::dense(1, max_elements + 1), + Err(FfnError::ExtentOverflow) + ); + assert_eq!( + RowMajorShape::dense(1, 1).unwrap().validate_len(usize::MAX), + Err(FfnError::ExtentOverflow) + ); +} + +#[test] +fn model_shapes_keep_input_pitch_separate_from_dense_workspace() { + let talker = FfnShape::<2048, 6144>::new(3, 2051).unwrap(); + assert_eq!(talker.rows(), 3); + assert_eq!(talker.input(), RowMajorShape::new(3, 2048, 2051).unwrap()); + assert_eq!(talker.hidden(), RowMajorShape::dense(3, 2048).unwrap()); + assert_eq!( + talker.intermediate(), + RowMajorShape::dense(3, 6144).unwrap() + ); + let predictor = FfnShape::<1024, 3072>::new(1, 1024).unwrap(); + assert_eq!(predictor.hidden().required_len(), 1024); + assert_eq!(predictor.intermediate().required_len(), 3072); +} + +#[test] +fn ffn_rejects_zero_dimensions_short_inputs_and_workspace_overflow() { + assert_eq!(FfnShape::<0, 7>::new(1, 1), Err(FfnError::ZeroWidth)); + assert_eq!(FfnShape::<7, 0>::new(1, 7), Err(FfnError::ZeroWidth)); + assert_eq!(FfnShape::<7, 13>::new(1, 6), Err(FfnError::StrideTooSmall)); + assert_eq!( + FfnShape::<1, 2>::new(isize::MAX as usize / 2, 1), + Err(FfnError::ExtentOverflow) + ); + let empty = FfnShape::<7, 13>::new(0, 7).unwrap(); + assert_eq!(empty.hidden().required_len(), 0); + assert_eq!(empty.intermediate().required_len(), 0); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/global_store.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/global_store.rs new file mode 100644 index 0000000000..ea4a1aae27 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/global_store.rs @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![doc = "Checks global-store type contracts and guaranteed carrier alignment."] +/// Store function pointers must preserve scalar type and compile-time vector width. +#[test] +fn global_store_preserves_element_and_vector_width_in_its_signature() { + let _: unsafe fn(*mut f32, &cute_rs::GlobalStoreTile) = + cute_rs::store_global_no_allocate::; + let _: unsafe fn(*mut u32, &cute_rs::GlobalStoreTile) = + cute_rs::store_global_no_allocate::; +} + +/// Every supported register-tile store carrier must provide its full vector alignment. +#[test] +fn global_store_carrier_guarantees_the_full_source_alignment() { + use core::mem::{align_of, size_of}; + assert_eq!(align_of::>(), 32); + assert_eq!(size_of::>(), 32); + let tile = cute_rs::GlobalStoreTile([0x7fc0_0001u32, 1, 2, 3, 4, 5, 6, 0x8000_0000]); + assert_eq!(tile.0[0], 0x7fc0_0001); + assert_eq!(tile.0.as_ptr() as usize % 32, 0); + assert_eq!((&tile as *const _ as usize), tile.0.as_ptr() as usize); +} + +/// Compile checks must reject unsupported store elements and calls outside unsafe blocks. +#[test] +fn global_store_rejects_unsupported_storage_and_requires_unsafe() { + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/global_store_unsafe.rs"); + tests.compile_fail("tests/compile_fail/global_store_storage.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/module_paths.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/module_paths.rs new file mode 100644 index 0000000000..62f536caa6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/module_paths.rs @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Root and parent-module re-exports must name the same public types as their defining paths. + +macro_rules! same_type { + ($import:ty, $defining:ty) => { + let _: fn($import) -> $defining = core::convert::identity; + }; +} + +#[test] +fn parent_module_re_exports_preserve_type_identity() { + same_type!( + cute_rs::epilogue::Sm120Epilogue128x128, + cute_rs::epilogue::sm120::Sm120Epilogue128x128 + ); +} + +#[test] +fn root_exports_preserve_type_identity() { + same_type!(cute_rs::FragmentA, cute_rs::mma::FragmentA); + same_type!(cute_rs::RegisterTile, cute_rs::tensor::RegisterTile); + same_type!( + cute_rs::Mxf4TiledMma<()>, + cute_rs::mma::block_scaled::Mxf4TiledMma<()> + ); + same_type!( + cute_rs::Sm120Epilogue128x128, + cute_rs::epilogue::sm120::Sm120Epilogue128x128 + ); + same_type!(cute_rs::SmemTile, cute_rs::copy::cooperative::SmemTile); + same_type!(cute_rs::SharedTensor, cute_rs::copy::tiled::SharedTensor); + same_type!(cute_rs::GlobalStoreTile, cute_rs::copy::vector::GlobalStoreTile); + same_type!(cute_rs::TensorMap, cute_rs::copy::tma::TensorMap); + same_type!(cute_rs::Layout<(), ()>, cute_rs::markers::Layout<(), ()>); + same_type!(cute_rs::ComposedLayout<(), 0, ()>, cute_rs::markers::ComposedLayout<(), 0, ()>); + same_type!(cute_rs::CopyG2S<16>, cute_rs::markers::CopyG2S<16>); + same_type!(cute_rs::ThrCopy<(), (), (), ()>, cute_rs::copy::tiled::ThrCopy<(), (), (), ()>); +} + +#[cfg(not(feature = "host"))] +#[test] +fn device_attention_is_available_without_host_metadata() { + use cute_rs::attention::sm100::{SharedLayout, SmemLayout}; + + assert_eq!(SharedLayout::<64, 64>::crd2idx(0, 1), 1); + assert_eq!(SharedLayout::<64, 64>::ROWS, 64); +} + +#[cfg(feature = "host")] +#[test] +fn host_attention_metadata_keeps_parent_imports() { + same_type!( + cute_rs::attention::AttentionError, + cute_rs::attention::metadata::AttentionError + ); + same_type!(cute_rs::attention::AttentionShape<16, 8, 128>, cute_rs::attention::metadata::AttentionShape<16, 8, 128>); + same_type!( + cute_rs::attention::AttentionWindow, + cute_rs::attention::metadata::AttentionWindow + ); + same_type!(cute_rs::attention::CacheShape<8, 128>, cute_rs::attention::metadata::CacheShape<8, 128>); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/register_tile.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/register_tile.rs new file mode 100644 index 0000000000..a2444af196 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/register_tile.rs @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Behavioral contracts for typed register-tile composition."] +#![feature(f16)] + +use cute_rs::tensor::RegisterTile; + +#[test] +fn construction_preserves_lane_order_and_supports_constants() { + const TILE: RegisterTile = RegisterTile::from_array([3, -2, 7, 1]); + assert_eq!(TILE.as_array(), &[3, -2, 7, 1]); + assert_eq!(TILE.into_array(), [3, -2, 7, 1]); +} + +#[test] +fn producer_constructs_each_register_once_in_order() { + let mut coordinates = Vec::new(); + let tile = RegisterTile::::from_fn(|i| { + coordinates.push(i); + 7 + i as i32 * 3 + }); + assert_eq!(coordinates, [0, 1, 2, 3]); + assert_eq!(tile.into_array(), [7, 10, 13, 16]); + let singleton = RegisterTile::::from_fn(|i| i + 9); + assert_eq!(singleton.into_array(), [9]); +} + +#[test] +fn empty_producer_never_calls_its_callback() { + let empty = RegisterTile::::from_fn(|_| panic!("empty producer invoked")); + assert_eq!(empty.into_array(), []); +} + +#[test] +fn map_visits_each_lane_once_in_order() { + let mut visited = Vec::new(); + let tile = RegisterTile::from_array([3, -2, 7, 1]); + let output: RegisterTile = tile.map(|value| { + visited.push(value); + i64::from(value) * visited.len() as i64 + }); + assert_eq!(visited, [3, -2, 7, 1]); + assert_eq!(output.into_array(), [3, -4, 21, 4]); +} + +#[test] +fn zip_map_pairs_lanes_once_and_allows_explicit_type_conversion() { + let lhs = RegisterTile::from_array([1.5_f16, -2.0, 0.25, 8.0]); + let rhs = RegisterTile::from_array([2_i32, 3, 4, -1]); + let mut visited = Vec::new(); + let output: RegisterTile = lhs.zip_map(rhs, |left, right| { + visited.push((left as f32, right)); + (left as f32) * (right as f32) + }); + assert_eq!(visited, [(1.5, 2), (-2.0, 3), (0.25, 4), (8.0, -1)]); + assert_eq!(output.into_array(), [3.0, -6.0, 1.0, -8.0]); +} + +#[test] +fn empty_tiles_do_not_invoke_transforms() { + let tile = RegisterTile::::from_array([]); + let mapped: RegisterTile = tile.map(|_| panic!("no lanes to map")); + let zipped: RegisterTile = mapped + .zip_map(RegisterTile::::from_array([]), |_, _| { + panic!("no lanes to zip") + }); + assert_eq!(zipped.into_array(), []); +} + +#[test] +fn fold_preserves_order_with_an_explicit_accumulator_type() { + let tile = RegisterTile::from_array([1_i16, 2, 3]); + let mut visited = Vec::new(); + let result = tile.fold(5_i64, |acc, value| { + visited.push(value); + acc * 10 + i64::from(value) + }); + assert_eq!(result, 5123); + assert_eq!(visited, [1, 2, 3]); + assert_eq!( + RegisterTile::::from_array([]).fold(7, |_, _| panic!("empty")), + 7 + ); +} + +#[test] +fn arithmetic_operators_apply_per_lane_in_order() { + let lhs = RegisterTile::from_array([8.0_f32, -6.0, 1.5, 0.0]); + let rhs = RegisterTile::from_array([2.0_f32, 3.0, -0.5, 4.0]); + assert_eq!((lhs + rhs).into_array(), [10.0, -3.0, 1.0, 4.0]); + assert_eq!((lhs - rhs).into_array(), [6.0, -9.0, 2.0, -4.0]); + assert_eq!((lhs * rhs).into_array(), [16.0, -18.0, -0.75, 0.0]); + assert_eq!((lhs / rhs).into_array(), [4.0, -2.0, -3.0, 0.0]); + let negated = (-lhs).into_array(); + assert_eq!(negated, [-8.0, 6.0, -1.5, 0.0]); + assert!(negated[3].is_sign_negative()); +} + +#[test] +fn integer_arithmetic_keeps_the_element_type_and_truncates_quotients() { + let lhs = RegisterTile::from_array([7_i32, -4, 9]); + let rhs = RegisterTile::from_array([2_i32, 3, -3]); + let difference: RegisterTile = lhs - rhs; + assert_eq!(difference.into_array(), [5, -7, 12]); + assert_eq!((lhs * rhs).into_array(), [14, -12, -27]); + assert_eq!((lhs / rhs).into_array(), [3, -1, -3]); + assert_eq!((-rhs).into_array(), [-2, -3, 3]); +} + +#[test] +fn empty_tiles_support_every_operator() { + let empty = RegisterTile::::from_array([]); + assert_eq!((empty + empty).into_array(), []); + assert_eq!((empty - empty).into_array(), []); + assert_eq!((empty * empty).into_array(), []); + assert_eq!((empty / empty).into_array(), []); + assert_eq!((-empty).into_array(), []); +} + +#[test] +fn warp_reduction_sum_exposes_an_unsafe_collective_boundary() { + let _: unsafe fn(f32) -> f32 = cute_rs::arch::warp_reduction_sum_f32; +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/shared_row.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/shared_row.rs new file mode 100644 index 0000000000..eda406a12b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/shared_row.rs @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Ordered streaming reductions over independently laid-out tensor views. + +use cute_rs::tensor::partition::StridedTensor; +use cute_rs::{ComposedLayout, RowMajor, SmemTile, Swizzle}; + +#[test] +fn zipped_row_reduction_preserves_coordinate_order_and_both_layouts() { + type Layout = ComposedLayout, 0, RowMajor<4, 4>>; + let mut shared = [0.0_f32; 16]; + for row in 0..4 { + for column in 0..4 { + shared[cute_rs::static_layout::crd2idx_2d::(row, column)] = + (row * 10 + column) as f32; + } + } + let tile = SmemTile:: { + base: shared.as_mut_ptr(), + capacity: shared.len(), + layout: core::marker::PhantomData, + }; + let global = [1.0_f32, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0]; + let global = StridedTensor::from_slice(&global).local_partition::<2>(0); + let mut visited = Vec::new(); + // SAFETY: the entire logical row is initialized, in bounds and has no writer. + let row = unsafe { tile.row_view::<4>(2, 0) }; + let result = row.zip_fold(global, 0.0, |accumulator, shared, global| { + visited.push((shared, global)); + accumulator + shared * global + }); + assert_eq!( + visited, + [(20.0, 1.0), (21.0, 2.0), (22.0, 3.0), (23.0, 4.0)] + ); + assert_eq!(result, 220.0); +} + +#[test] +fn zipped_row_reduction_does_not_reassociate() { + let mut shared = [1.0e20_f32, 1.0, -1.0e20, 3.0]; + let tile = SmemTile::> { + base: shared.as_mut_ptr(), + capacity: 4, + layout: core::marker::PhantomData, + }; + let ones = [1.0_f32; 4]; + let row = unsafe { tile.row_view::<4>(0, 0) }; + assert_eq!( + row.zip_fold(StridedTensor::from_slice(&ones), 0.0, |a, b, c| a + b * c), + 3.0 + ); +} + +#[test] +fn empty_shared_row_never_calls_its_reduction() { + let tile = SmemTile::> { + base: core::ptr::null_mut(), + capacity: 0, + layout: core::marker::PhantomData, + }; + // SAFETY: an empty row has no coordinates, so it never accesses the pointer. + let row = unsafe { tile.row_view::<0>(0, 0) }; + assert_eq!( + row.zip_fold( + StridedTensor::from_slice(&[]), + 9.0, + |_: f32, _: f32, _: f32| panic!("empty") + ), + 9.0 + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention.rs new file mode 100644 index 0000000000..68c4260db5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention.rs @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks attention shared layouts, typed MMA shapes and complete tile-array contracts."] +use cute_rs::attention::sm100::{MmaSs, MmaTs, SharedLayout, SharedTileArray, SmemLayout}; +use cute_rs::{OperandA, OperandB}; + +/// Keep explicit SM100 instructions behind CTM in both kernels and their CuTe helpers. +/// CuTe compiler markers for structured MMA/TMA remain a separate supported path. +#[test] +fn sm100_attention_uses_ctm_for_explicit_device_operations() { + let cute_rs_preview = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap(); + let paths = [ + "examples/sm100_fmha_prefill_cute/src/kernel.rs", + "examples/sm100_mla_decode_fp16_cute/src/kernel.rs", + "cute-rs/src/arch.rs", + "cute-rs/src/attention/sm100/pipeline.rs", + "cute-rs/src/attention/sm100/tmem.rs", + "cute-rs/src/attention/sm100/tma.rs", + "cute-rs/src/attention/sm100/__compiler.rs", + ]; + for path in paths { + let source = std::fs::read_to_string(cute_rs_preview.join(path)).unwrap(); + assert!( + !source.contains("cuda_device::"), + "{path} bypasses the CTM instruction/storage interface" + ); + } +} + +/// Enumerate one shared layout and require each storage element to have exactly one owner. +fn check_bijection() { + let mut offsets = Vec::new(); + for row in 0..L::ROWS { + for col in 0..L::COLS { + offsets.push(L::crd2idx(row, col)); + } + } + offsets.sort_unstable(); + assert_eq!( + offsets, + (0..(L::ROWS * L::COLS) as usize).collect::>() + ); +} + +/// Supported K-major and MN-major tiled swizzles must cover each element exactly once. +#[test] +fn tiled_swizzles_cover_storage_without_overlap() { + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); + check_bijection::>(); +} + +/// TMA chunk coordinates and MN-major coordinates must select identical physical FP16 offsets. +#[test] +fn tma_contiguous_chunks_and_transposed_v_share_the_same_storage() { + /// K-major 128-by-128 tile used to compare chunked TMA and transposed coordinates. + type K = SharedLayout<128, 128>; + assert_eq!(K::crd2idx(0, 0), 0); + assert_eq!(K::crd2idx(0, 64), 8192); + assert_eq!(K::crd2idx(1, 0), 72); + for row in 0..32 { + for col in 0..128 { + assert_eq!( + SharedLayout::<32, 128>::crd2idx(row, col), + SharedLayout::<128, 32, true>::crd2idx(col, row) + ); + } + } + assert_eq!(SharedLayout::<64, 32>::SWIZZLE_BITS, 2); + assert_eq!(SharedLayout::<128, 128>::SWIZZLE_BITS, 3); +} + +/// Supported shared/shared and tensor/shared attention shapes must construct without runtime data. +#[test] +fn independent_qk_and_pv_configurations_are_constructible() { + let _ = MmaSs::, SharedLayout<128, 128>, 128, 128, 128, 1>::new(); + let _ = MmaTs::, 128, 128, 128, 1>::new(); + let _ = MmaSs::, SharedLayout<64, 64>, 128, 128, 64, 2>::new(); + let _ = MmaSs::, SharedLayout<128, 32, true>, 128, 256, 32, 2>::new(); +} + +// Host storage models the aligned allocation without accessing it as GPU memory. +/// Full-swizzle-period-aligned buffer for three complete shared-layout test tiles. +#[repr(align(1024))] +struct SharedStorage( + #[doc = "FP16-sized storage words backing the complete tile array."] [u16; 3 * 16384], +); + +/// Check each complete tile’s base address and swizzle mapping within an aligned array. +fn check_tile_array() { + let mut storage = SharedStorage([0; 3 * 16384]); + let base = storage.0.as_mut_ptr(); + // SAFETY: storage covers three complete tiles and meets either alignment. + let tiles = unsafe { SharedTileArray::::from_raw(base.cast()) }; + let tile_elements = (L::ROWS * L::COLS) as usize; + let period_elements = 1usize << (L::SWIZZLE_BITS + 6); + assert_eq!(tile_elements % period_elements, 0); + let mut seen = vec![false; 3 * tile_elements]; + for tile_index in 0..3 { + // SAFETY: the index is in the array's static extent. + let tile = unsafe { tiles.tile_unchecked(tile_index) }; + assert_eq!( + tile.as_ptr().cast::(), + base.wrapping_add(tile_index as usize * tile_elements) + ); + for row in 0..L::ROWS { + for col in 0..L::COLS { + let offset = tile_index as usize * tile_elements + L::crd2idx(row, col); + assert!(!seen[offset], "overlapping physical tile elements"); + seen[offset] = true; + // The XOR swizzle is unchanged when translated by whole periods. + let mask = (1usize << L::SWIZZLE_BITS) - 1; + let unswizzled = offset ^ (((offset >> 6) & mask) << 3); + let local = L::crd2idx(row, col); + let local_unswizzled = local ^ (((local >> 6) & mask) << 3); + assert_eq!( + unswizzled, + tile_index as usize * tile_elements + local_unswizzled + ); + } + } + } + assert!(seen.into_iter().all(|value| value)); +} + +/// Every array element must preserve the complete tile mapping and expected byte displacement. +#[test] +fn complete_tile_arrays_preserve_addresses_and_swizzles() { + check_tile_array::, OperandB>(); + check_tile_array::, OperandA>(); + check_tile_array::, OperandB>(); + check_tile_array::, OperandA>(); + check_tile_array::, OperandB>(); + check_tile_array::, OperandA>(); + check_tile_array::, OperandA>(); +} + +/// Compile checks must reject wrong A/B roles, zero tile counts and overflowing arrays. +#[test] +fn tile_arrays_reject_operand_role_changes_and_invalid_extents() { + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/sm100_attention_tile_array_role.rs"); + tests.compile_fail("tests/compile_fail/sm100_attention_tile_array_extent.rs"); +} + +// Const evaluation rejects inbounds pointer arithmetic after an allocation +// expires. A safe raw-address accessor must not require that allocation live. +/// Require shared tile-array construction and selection to remain usable in const evaluation. +const _: () = { + let tile = { + let mut storage = SharedStorage([0; 3 * 16384]); + // SAFETY: two full tiles fit in this aligned, currently live storage. + let tiles = unsafe { + SharedTileArray::, OperandA, 2>::from_raw( + storage.0.as_mut_ptr().cast(), + ) + }; + // SAFETY: index one is within the two-tile extent; no asynchronous use. + unsafe { tiles.tile_unchecked(1) } + }; + let _ = tile.as_ptr(); +}; diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention_pipeline.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention_pipeline.rs new file mode 100644 index 0000000000..338a0737eb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/sm100_attention_pipeline.rs @@ -0,0 +1,93 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Checks two-CTA barrier-ring addressing, persistent cursors and policy rejection."] +use core::mem::MaybeUninit; +use cuda_device::barrier::Barrier; +use cute_rs::attention::sm100::pipeline::{ + AsyncCompletion, Consumer, Producer, RingCursor, TwoCtaRing, +}; + +/// Full and empty arrays must be disjoint, and ticket indices must wrap within each array. +#[test] +fn barrier_views_select_disjoint_arrays_and_wrap_slots() { + // Address calculation is CPU-safe; this test never initializes or uses GPU barriers. + let mut storage = [const { MaybeUninit::::uninit() }; 8]; + let base = storage.as_mut_ptr().cast::(); + let ring = unsafe { TwoCtaRing::<3, AsyncCompletion>::from_raw(base, 1) }; + for (ticket, full_offset, empty_offset) in + [(0, 1, 4), (1, 2, 5), (2, 3, 6), (3, 1, 4), (7, 2, 5)] + { + unsafe { + assert_eq!(ring.full_at(ticket), base.add(full_offset)); + assert_eq!(ring.empty_at(ticket), base.add(empty_offset)); + } + } +} + +/// Compile checks must reject signals that do not match the ring edge’s memory-ordering policy. +#[test] +fn completion_policies_reject_incompatible_signals() { + let tests = trybuild::TestCases::new(); + tests.pass("tests/compile_pass/two_cta_ring_valid_signal.rs"); + tests.pass("tests/compile_pass/two_cta_ring_non_tc_receivers.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_non_tc_wrong_signal.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_wrong_signal.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_zero_stages.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_policy_sealed.rs"); +} + +/// Cursors must retain slot, parity and barrier addresses across persistent work boundaries. +#[test] +fn cursors_preserve_phase_and_addresses_across_persistent_work() { + let mut storage = [const { MaybeUninit::::uninit() }; 32]; + let base = storage.as_mut_ptr().cast::(); + let ring = unsafe { TwoCtaRing::<15, AsyncCompletion>::from_raw(base, 1) }; + let mut producer = RingCursor::<15, Producer>::new(); + let mut consumer = RingCursor::<15, Consumer>::new(); + let mut ticket = 0; + for work_stages in [17, 1, 256, 4097] { + for _ in 0..work_stages { + assert_eq!(producer.index(), ticket % 15); + assert_eq!(consumer.index(), ticket % 15); + assert_eq!(producer.phase(), ((ticket / 15) & 1) ^ 1); + assert_eq!(consumer.phase(), (ticket / 15) & 1); + unsafe { + assert_eq!(ring.full_at_cursor(producer), ring.full_at(ticket)); + assert_eq!(ring.full_at_cursor(consumer), ring.full_at(ticket)); + assert_eq!(ring.empty_at_cursor(producer), ring.empty_at(ticket)); + assert_eq!(ring.empty_at_cursor(consumer), ring.empty_at(ticket)); + } + producer.advance(); + consumer.advance(); + ticket += 1; + } + } + assert_eq!((producer.index(), producer.phase()), (6, 0)); + assert_eq!((consumer.index(), consumer.phase()), (6, 1)); +} + +/// Producer and consumer parities must flip after every one-stage handoff and remain opposite. +#[test] +fn one_stage_cursors_flip_opposite_phases_on_every_advance() { + let mut producer = RingCursor::<1, Producer>::default(); + let mut consumer = RingCursor::<1, Consumer>::default(); + for (producer_phase, consumer_phase) in [(1, 0), (0, 1), (1, 0), (0, 1)] { + assert_eq!((producer.index(), producer.phase()), (0, producer_phase)); + assert_eq!((consumer.index(), consumer.phase()), (0, consumer_phase)); + producer.advance(); + consumer.advance(); + } +} + +/// Compile checks must reject cursor/ring stage mismatches, role swaps and empty rings. +#[test] +fn cursors_reject_wrong_stages_roles_and_zero_stages() { + let tests = trybuild::TestCases::new(); + tests.pass("tests/compile_pass/two_cta_ring_valid_cursor.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_cursor_wrong_stages.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_cursor_wrong_role.rs"); + tests.compile_fail("tests/compile_fail/two_cta_ring_cursor_zero_stages.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/static_layout.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/static_layout.rs new file mode 100644 index 0000000000..f1027b4982 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/static_layout.rs @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Static offset-span validation, including layouts with a tiny logical extent. +use cute_rs::markers::{C, ComposedLayout, Layout, RowMajor, Swizzle}; +use cute_rs::static_layout::StaticLayout; + +#[test] +fn valid_boundary_offsets_and_swizzle_bounds_are_representable() { + type Boundary = Layout, C<{ i64::MAX }>>; + assert_eq!(Boundary::crd2idx(1), i64::MAX as usize); + assert_eq!(Boundary::MAX_OFFSET, i64::MAX as usize); + type Pitched = Layout<(C<2>, C<3>), (C<17>, C<4>)>; + assert_eq!(Pitched::MAX_OFFSET, 25); + type ComposedMap = ComposedLayout, 7, RowMajor<3, 5>>; + for i in 0..ComposedMap::SIZE { + assert!(ComposedMap::crd2idx(i) <= ComposedMap::MAX_OFFSET); + } +} + +#[test] +fn overflowing_physical_layouts_fail_before_device_execution() { + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/static_layout_leaf_overflow.rs"); + tests.compile_fail("tests/compile_fail/static_layout_sum_overflow.rs"); + tests.compile_fail("tests/compile_fail/static_layout_composed_overflow.rs"); + tests.compile_fail("tests/compile_fail/static_layout_2d_overflow.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/strided_partition.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/strided_partition.rs new file mode 100644 index 0000000000..4bcfa9595c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/strided_partition.rs @@ -0,0 +1,193 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Partition coverage, ordered reductions and disjoint stores on ordinary memory. + +use cute_rs::tensor::partition::{StridedTensor, StridedTensorMut}; + +#[test] +fn nested_partitions_retain_their_types_and_compose_coordinates() { + use cute_rs::tensor::partition::{Cyclic, Strided1D, Subrange}; + fn typed(_: &StridedTensor<'_, f32, Cyclic>, 3>>) {} + let storage: Vec = (0..100).map(|i| i as f32).collect(); + let selected = StridedTensor::from_slice(&storage) + .local_partition::<4>(1) + .subview(2, 12) + .local_partition::<3>(2); + typed(&selected); + assert_eq!(selected.load::<4>().into_array(), [17.0, 29.0, 41.0, 53.0]); + let mut output = vec![-1.0; 100]; + let mut destination = StridedTensorMut::from_mut_slice(&mut output) + .local_partition::<4>(1) + .subview(2, 12) + .local_partition::<3>(2); + unsafe { + selected.map_to(&mut destination, |value| value); + } + for (i, value) in output.iter().enumerate() { + assert_eq!( + *value, + if [17, 29, 41, 53].contains(&i) { + i as f32 + } else { + -1.0 + } + ); + } +} + +#[test] +fn singleton_and_empty_nested_partitions_do_not_multiply_physical_stride() { + let storage = [5.0_f32]; + let singleton = StridedTensor::from_slice(&storage) + .local_partition::<{ usize::MAX }>(0) + .local_partition::<{ usize::MAX }>(0); + assert_eq!(singleton.get(0), 5.0); + let empty = StridedTensor::from_slice(&storage) + .local_partition::<{ usize::MAX }>(usize::MAX - 1) + .subview(0, 0); + assert_eq!( + empty.fold(9.0, |_, _| panic!("empty layout must not access storage")), + 9.0 + ); +} + +#[test] +fn partition_layout_and_storage_identity_are_compile_time_contracts() { + let tests = trybuild::TestCases::new(); + tests.compile_fail("tests/compile_fail/cyclic_partition_mismatch.rs"); + tests.compile_fail("tests/compile_fail/strided_tensor_storage_mismatch.rs"); + tests.compile_fail("tests/compile_fail/strided_tensor_mut_reborrow.rs"); + tests.compile_fail("tests/compile_fail/strided_tensor_mut_move.rs"); + tests.compile_fail("tests/compile_fail/strided_layout_is_sealed.rs"); +} + +#[test] +fn encoded_storage_retains_its_element_contract_and_bits() { + use cute_rs::tensor::partition::Strided1D; + let bytes = [0x12_u8, 0x34, 0x56, 0x78]; + let encoded = StridedTensor::::from_storage(&bytes); + assert_eq!( + encoded.local_partition::<2>(1).load::<2>().into_array(), + [0x34, 0x78] + ); +} + +#[test] +fn lane_partition_covers_tails_once_in_original_order() { + let input: Vec = (0..70).map(|i| i as f32).collect(); + let mut output = vec![-1.0; input.len()]; + for lane in 0..32 { + let source = StridedTensor::from_slice(&input).local_partition::<32>(lane); + let mut destination = + StridedTensorMut::from_mut_slice(&mut output).local_partition::<32>(lane); + let expected: Vec<_> = (lane..70).step_by(32).map(|i| i as f32).collect(); + let actual = source.fold(Vec::new(), |mut values, value| { + values.push(value); + values + }); + assert_eq!(actual, expected); + // SAFETY: each lane writes its own positions; iterations are sequential. + unsafe { source.map_to(&mut destination, |value| value + 1.0) }; + } + assert_eq!(output, input.iter().map(|v| v + 1.0).collect::>()); +} + +#[test] +fn register_gather_and_scatter_keep_stride_and_guards() { + let input: Vec = (0..128).map(|i| i as f32).collect(); + let mut output = [-9.0; 130]; + let gathered = StridedTensor::from_slice(&input) + .local_partition::<32>(7) + .load::<4>(); + assert_eq!(gathered.into_array(), [7.0, 39.0, 71.0, 103.0]); + let mut destination = StridedTensorMut::from_mut_slice(&mut output) + .subview(1, 128) + .local_partition::<32>(7); + // SAFETY: this is the sole writable view and its four coordinates are distinct. + unsafe { destination.store(gathered) }; + for (i, value) in output.iter().enumerate() { + assert_eq!( + *value, + if [8, 40, 72, 104].contains(&i) { + (i - 1) as f32 + } else { + -9.0 + } + ); + } +} + +#[test] +fn fold_does_not_reassociate_floating_point_values() { + let input = [1.0e20_f32, 1.0, -1.0e20, 3.0]; + assert_eq!( + StridedTensor::from_slice(&input).fold(0.0, |a, b| a + b), + 3.0 + ); +} + +#[test] +fn empty_and_outside_lanes_never_access_storage() { + let input = [7.0_f32]; + for lane in 1..32 { + let view = StridedTensor::from_slice(&input).local_partition::<32>(lane); + assert!(view.is_empty()); + assert_eq!(view.fold(5.0, |a, b| a + b), 5.0); + } + let empty = StridedTensor::from_slice(&input).subview(1, 0); + assert!(empty.local_partition::<32>(31).is_empty()); +} + +#[test] +#[should_panic] +fn subview_rejects_overflowing_extent() { + let input = [1.0_f32]; + let _ = StridedTensor::from_slice(&input).subview(usize::MAX, 2); +} + +#[test] +#[should_panic] +fn fixed_load_rejects_short_partition() { + let input = [1.0_f32; 63]; + let _ = StridedTensor::from_slice(&input) + .local_partition::<32>(31) + .load::<2>(); +} + +#[test] +fn exclusive_update_changes_only_the_selected_coordinate() { + let mut data = [3_u32, 8, 10]; + let mut view = StridedTensorMut::from_mut_slice(&mut data); + unsafe { view.update_at(1, |value| value | 4) }; + assert_eq!(data, [3, 12, 10]); +} + +#[test] +fn register_gather_expresses_permuted_logical_coordinates() { + let tile = cute_rs::RegisterTile::from_array([3.0_f32, 7.0]); + assert_eq!(tile.gather([1, 0]).into_array(), [7.0, 3.0]); +} + +#[test] +fn shared_tile_loads_follow_its_declared_layout() { + type Layout = cute_rs::ComposedLayout, 0, cute_rs::RowMajor<4, 4>>; + let mut storage = [0.0_f32; 16]; + for row in 0..4 { + for col in 0..4 { + let offset = cute_rs::static_layout::crd2idx_2d::(row, col); + storage[offset] = (10 * row + col) as f32; + } + } + let tile = cute_rs::SmemTile:: { + base: storage.as_mut_ptr(), + capacity: storage.len(), + layout: core::marker::PhantomData, + }; + // SAFETY: host-backed storage has the same layout, is initialized and has no writer. + assert_eq!( + unsafe { tile.load::<4>(2, 0) }.into_array(), + [20.0, 21.0, 22.0, 23.0] + ); +} diff --git a/cuda-oxide/cute-rs-preview/cute-rs/tests/type_safety.rs b/cuda-oxide/cute-rs-preview/cute-rs/tests/type_safety.rs new file mode 100644 index 0000000000..afc1675bcc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/cute-rs/tests/type_safety.rs @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![doc = "Compile-fail coverage for tensor storage, operand roles and writable-view ownership."] +/// Compile checks must reject invalid storage encodings, operand roles and mutable-view duplication. +#[test] +fn tensor_and_block_scaled_types_fail_closed() { + let tests = trybuild::TestCases::new(); + tests.pass("tests/compile_pass/tma_descriptor_prefetch.rs"); + tests.compile_fail("tests/compile_fail/tma_descriptor_prefetch_requires_unsafe.rs"); + tests.compile_fail("tests/compile_fail/invalid_storage.rs"); + tests.compile_fail("tests/compile_fail/immutable_store.rs"); + tests.compile_fail("tests/compile_fail/tensor_mut_copy.rs"); + tests.compile_fail("tests/compile_fail/tensor_mut_clone.rs"); + tests.compile_fail("tests/compile_fail/register_tile_lane_mismatch.rs"); + tests.compile_fail("tests/compile_fail/register_tile_implicit_conversion.rs"); + tests.compile_fail("tests/compile_fail/bf16_implicit_conversion.rs"); + tests.compile_fail("tests/compile_fail/bf16_invalid_storage.rs"); + tests.compile_fail("tests/compile_fail/warp_sum_requires_unsafe.rs"); + tests.compile_fail("tests/compile_fail/mma_operand_role_mismatch.rs"); + tests.compile_fail("tests/compile_fail/mma_scale_role_mismatch.rs"); + tests.compile_fail("tests/compile_fail/tiled_copy_role_mismatch.rs"); + tests.compile_fail("tests/compile_fail/tiled_copy_encoding_mismatch.rs"); + tests.compile_fail("tests/compile_fail/tma_store_layout_mismatch.rs"); + tests.compile_fail("tests/compile_fail/sm100_operand_role_mismatch.rs"); + tests.compile_fail("tests/compile_fail/sm100_tma_layout_mismatch.rs"); + tests.compile_fail("tests/compile_fail/sm100_mma_shape_mismatch.rs"); + tests.compile_fail("tests/compile_fail/sm100_attention_operand_role.rs"); + tests.compile_fail("tests/compile_fail/sm100_attention_tma_layout.rs"); + tests.compile_fail("tests/compile_fail/sm100_attention_invalid_layout.rs"); + tests.compile_fail("tests/compile_fail/sm100_attention_ts_k16.rs"); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/Cargo.toml b/cuda-oxide/cute-rs-preview/dialect-cute/Cargo.toml new file mode 100644 index 0000000000..3358bc6ca9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "dialect-cute" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "CuTe dialect for Pliron IR framework (layout-algebra ops and types)" + +[dependencies] +pliron = { workspace = true } +pliron-derive = { workspace = true } +thiserror = { workspace = true } +linkme = { workspace = true } +combine = { workspace = true } +once_cell = { workspace = true } +cute-layout = { workspace = true } +# Semantic operations use ordinary MIR carrier types. Native and MLIR +# backends depend on this crate, not the other way around. +dialect-mir = { workspace = true } diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/README.md b/cuda-oxide/cute-rs-preview/dialect-cute/README.md new file mode 100644 index 0000000000..d234b0d74a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/README.md @@ -0,0 +1,44 @@ +# CuTe dialect + +`dialect-cute` defines backend-neutral CuTe operations for cuda-oxide using +Pliron IR. Operations retain typed layouts and collective contracts alongside +ordinary MIR scalar and pointer carriers. + +The public source boundaries are: + +| Module | Responsibility | +| --- | --- | +| [`ops/`](src/ops) | Operation families and their local verification contracts. | +| [`attributes/`](src/attributes) | Static layouts, formats, roles, and plans. | +| [`types/`](src/types) | Tensor descriptions, fragment shapes, and compiler-only handles. | +| [`verify/`](src/verify) | Function-level control flow, provenance, and operation protocol checks. | + +Prefer the grouped operation paths below. Each former root module is a public +alias to the same module, preserving its existing types and functions. + +| Preferred path | Compatibility alias | Contents | +| --- | --- | --- | +| `ops::copy` | `copy_ops` | Direct and cooperative copies, raw TMA transport, cache-aware global stores. | +| `ops::tensor` | `tensor_ops` | Tensor construction, partitioning, selection, and transfers. | +| `ops::gemv` | `gemv_ops` | Scaled tensor views and block-scaled dot products. | +| `ops::gemm_tma` | `gemm_tma_ops` | Typed global/shared transport views and GEMM TMA copies. | +| `ops::mma` | `smem_mma_ops` | Shared tensor overlays, fragments, matrix loads, and tiled GEMM. | +| `ops::epilogue` | `epilogue_ops` | Shared output tiles and asynchronous TMA stores. | +| `ops::pipeline` | `pipeline_ops` | Load-pipeline construction, scalar state, and lifecycle. | +| `ops::scheduler` | `scheduler_ops` | Persistent work grids, coordinates, and cursor advancement. | +| `ops::sm100` | `sm100_ops` | Two-CTA tensor-memory MMA, cluster TMA, and pipelines. | +| `ops::sm100_attention` | `sm100_attention_ops` | Independent SM100 attention MMA and TMA contracts. | + +The existing `ops::CuteCopyOp`, `ops::CuteCopyG2SOp`, `ops::CuteAssumeDivOp`, +`ops::CuteLdmatrixOp`, `ops::CuteTmaLoad2dOp`, and `ops::CuteTmaStore2dOp` imports +also remain available. Larger families separate production responsibilities +into private children and keep regression fixtures in `tests.rs`. + +Run the host checks from the owning workspace, `cuda-oxide/`: + +```bash +cargo test -p dialect-cute +cargo clippy -p dialect-cute --all-targets -- -D warnings +RUSTDOCFLAGS="-D warnings" cargo doc -p dialect-cute --no-deps +cargo fmt -p dialect-cute --check +``` diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/copy.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/copy.rs new file mode 100644 index 0000000000..c12e525d7d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/copy.rs @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed copy widths, divisibility promises, and warp matrix operand roles. + +use pliron_derive::pliron_attr; + +/// A promise that an integer is divisible by this positive number. +/// +/// This is deliberately a typed attribute instead of an unstructured string: +/// A later direct SSA consumer can ask the defining `cute.assume_div` op for +/// the exact divisor without rediscovering it from surrounding arithmetic. +#[pliron_attr(name = "cute.divisibility", format = "$0", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteDivisibilityAttr(#[doc = "Positive divisor promised by the caller."] pub u64); + +/// Width of one copy atom in bytes. +#[pliron_attr(name = "cute.copy_atom_bytes", format = "$0", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteCopyAtomAttr(#[doc = "Bytes transferred by one copy instruction."] pub u32); + +/// Which `mma.sync` operand a warp-cooperative matrix load feeds. +/// +/// The role fixes the window shape and the `ldmatrix` variant: +/// +/// ```text +/// A: 16x16 window, four 8x8 matrices (.x4), row-major fragment +/// B: 16x8 window, two 8x8 matrices (.x2), transposed in flight +/// BPair: 16x16 window, four 8x8 matrices (.x4), transposed in flight +/// BPairKMajor: 16x16 window, four 8x8 matrices (.x4), K-major B pair without transpose +/// ``` +#[pliron_attr(name = "cute.matrix_role", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteMatrixRoleAttr { + /// A-operand fragment loaded from four 8-by-8 matrices without transpose. + A, + /// B-operand fragment loaded from two 8-by-8 matrices with transpose. + B, + /// Two B fragments loaded from four 8-by-8 matrices with transpose. + BPair, + /// Two K-major B fragments loaded from four matrices without transpose. + BPairKMajor, +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/epilogue.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/epilogue.rs new file mode 100644 index 0000000000..1b37030f20 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/epilogue.rs @@ -0,0 +1,287 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared output placement and asynchronous store synchronization contracts. + +use cute_layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +use super::{layout::CuteComposedLayoutAttr, mma::CuteTiledMmaPlanAttr, tensor::CuteAlignmentAttr}; + +/// How many result buffers the asynchronous TMA store queue may reuse. +/// +/// The value has no runtime storage. It only explains the wait count used by +/// the three store-pipeline operations: +/// +/// ```text +/// acquire waits for at most stages - 1 old readers +/// commit closes the stores issued since the previous commit +/// tail waits for zero old readers +/// ``` +#[pliron_attr(name = "cute.tma_store_pipeline", format = "`<` $stages `>`")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteTmaStorePipelineAttr { + /// Number of reusable slots in the barrier or output-buffer ring. + pub stages: u32, +} + +impl CuteTmaStorePipelineAttr { + /// Record the number of output buffers tracked by asynchronous store groups. + #[must_use] + pub const fn new(stages: u32) -> Self { + Self { stages } + } + + /// Number passed to `cp.async.bulk.wait_group.read` before reuse. + #[must_use] + pub const fn max_pending(self) -> Option { + self.stages.checked_sub(1) + } +} + +impl Verify for CuteTmaStorePipelineAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.stages == 0 { + return verify_err_noloc!("cute.tma_store_pipeline needs at least one stage"); + } + Ok(()) + } +} + +/// One named CTA barrier used by a contiguous group of warps. +/// +/// The current epilogue uses warps 0 through 7 and leaves warp 8 out: +/// +/// ```text +/// CTA warps 0 1 2 3 4 5 6 7 | 8 +/// barrier ID 2 [--- 256 threads ---] | producer +/// ``` +#[pliron_attr( + name = "cute.counted_cta_barrier", + format = "`<` $barrier_id `,` $first_warp `,` $warp_count `,` $cta_warps `,` $lanes_per_warp `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteCountedCtaBarrierAttr { + /// Named CTA barrier identifier, in `0..16`. + pub barrier_id: u32, + /// First participating warp index within the CTA. + pub first_warp: u32, + /// Number of consecutive participating warps. + pub warp_count: u32, + /// Total warps in the CTA, including excluded producers. + pub cta_warps: u32, + /// Thread count contributed by each participating warp. + pub lanes_per_warp: u32, +} + +impl CuteCountedCtaBarrierAttr { + /// Record the barrier identifier and exact participating/excluded warp ranges. + #[must_use] + pub const fn new( + barrier_id: u32, + first_warp: u32, + warp_count: u32, + cta_warps: u32, + lanes_per_warp: u32, + ) -> Self { + Self { + barrier_id, + first_warp, + warp_count, + cta_warps, + lanes_per_warp, + } + } + + /// Threads that must arrive at the counted barrier. + #[must_use] + pub const fn participant_threads(self) -> Option { + self.warp_count.checked_mul(self.lanes_per_warp) + } + + /// Warps in the CTA that deliberately do not arrive. + #[must_use] + pub const fn excluded_warps(self) -> Option { + self.cta_warps.checked_sub(self.warp_count) + } +} + +impl Verify for CuteCountedCtaBarrierAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.barrier_id >= 16 { + return verify_err_noloc!("cute.counted_cta_barrier ID must be between 0 and 15"); + } + if self.warp_count == 0 || self.cta_warps == 0 || self.lanes_per_warp == 0 { + return verify_err_noloc!( + "cute.counted_cta_barrier warp and lane counts must be greater than zero" + ); + } + let Some(end_warp) = self.first_warp.checked_add(self.warp_count) else { + return verify_err_noloc!("cute.counted_cta_barrier warp range overflowed"); + }; + if end_warp > self.cta_warps { + return verify_err_noloc!( + "cute.counted_cta_barrier participating warps must lie inside the CTA" + ); + } + let Some(threads) = self.participant_threads() else { + return verify_err_noloc!("cute.counted_cta_barrier thread count overflowed"); + }; + let Some(cta_threads) = self.cta_warps.checked_mul(self.lanes_per_warp) else { + return verify_err_noloc!("cute.counted_cta_barrier CTA thread count overflowed"); + }; + if threads == 0 || threads > 1024 || cta_threads > 1024 { + return verify_err_noloc!( + "cute.counted_cta_barrier cannot describe more than 1024 CTA threads" + ); + } + Ok(()) + } +} + +/// Why the compute warps meet at the epilogue barrier. +#[pliron_attr(name = "cute.epilogue_sync_phase", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteEpilogueSyncPhaseAttr { + /// The previous TMA reader is done, so writers may reuse shared memory. + Reusable, + /// Publish generic-proxy writes, then release TMA to read shared memory. + ReadyForTma, +} + +/// Left or right 128x64 half of the current 128x128 result tile. +#[pliron_attr(name = "cute.epilogue_half_index", format = "$0")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteEpilogueHalfAttr(#[doc = "Zero-based left/right half index."] pub u32); + +impl Verify for CuteEpilogueHalfAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.0 >= 2 { + return verify_err_noloc!("cute.epilogue_half_index v0 must be 0 or 1"); + } + Ok(()) + } +} + +/// How one warp's FP32 accumulator reaches a shared FP16 result tile. +/// +/// ```text +/// 2 x 8 accumulator cells per warp +/// │ +/// ▼ +/// logical 128x128 f16 tile +/// ┌────── 128x64 ──────┬────── 128x64 ──────┐ +/// │ one B128 TMA store │ one B128 TMA store │ +/// └────────────────────┴────────────────────┘ +/// ``` +#[pliron_attr( + name = "cute.epilogue_plan", + format = "`<` $tiled_mma `,` $half_layout `,` $halves `,` $base_alignment `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteEpiloguePlanAttr { + /// Compute plan whose accumulator mapping feeds this output tile. + pub tiled_mma: CuteTiledMmaPlanAttr, + /// Physical shared placement of one output half. + pub half_layout: CuteComposedLayoutAttr, + /// Number of shared output halves in the complete result tile. + pub halves: u32, + /// Promised shared-base alignment, in bytes. + pub base_alignment: CuteAlignmentAttr, +} + +impl CuteEpiloguePlanAttr { + /// Current SM120 two-half FP16 epilogue for the 128x128 MXFP4 plan. + #[must_use] + pub fn sm120_mxf4_128x128(tiled_mma: CuteTiledMmaPlanAttr) -> Self { + let outer: Layout = "(128,64):(64,1)" + .parse() + .expect("fixed epilogue half layout is valid"); + let half_layout = + ComposedLayout::new(Swizzle::new(3, 3, 3), 0, outer, OffsetUnit::Elements) + .expect("fixed epilogue swizzle is valid"); + Self { + tiled_mma, + half_layout: CuteComposedLayoutAttr(half_layout), + halves: 2, + base_alignment: CuteAlignmentAttr(1024), + } + } + + /// Return the element capacity of one output half, or `None` on overflow. + #[must_use] + pub fn half_elements(&self) -> Option { + u64::try_from(self.half_layout.0.outer().checked_size()?).ok() + } + + /// Return the element capacity of the complete output tile, or `None` on overflow. + #[must_use] + pub fn full_elements(&self) -> Option { + self.half_elements()?.checked_mul(u64::from(self.halves)) + } + + /// Return one output half's byte capacity, or `None` for an invalid size. + #[must_use] + pub fn half_bytes(&self, element_bytes: u64) -> Option { + self.half_elements()?.checked_mul(element_bytes) + } + + /// Return the complete output tile's byte capacity, or `None` for an invalid size. + #[must_use] + pub fn full_bytes(&self, element_bytes: u64) -> Option { + self.full_elements()?.checked_mul(element_bytes) + } +} + +impl Verify for CuteEpiloguePlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.tiled_mma.verify(ctx)?; + self.base_alignment.verify(ctx)?; + + let data_outer: Layout = "(128,32):(32,1)" + .parse() + .expect("fixed SM120 MMA shared layout is valid"); + let data_layout = + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, data_outer, OffsetUnit::Elements) + .expect("fixed SM120 MMA shared swizzle is valid"); + let expected_mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(data_layout); + if self.tiled_mma != expected_mma { + return verify_err_noloc!( + "cute.epilogue_plan v0 needs the 128x128x128 SM120 MXFP4 tiled-MMA plan" + ); + } + let expected = Self::sm120_mxf4_128x128(self.tiled_mma.clone()); + if self.half_layout != expected.half_layout + || self.halves != expected.halves + || self.base_alignment != expected.base_alignment + { + return verify_err_noloc!( + "cute.epilogue_plan v0 needs two B128-swizzled 128x64 f16 halves aligned to 1024 bytes" + ); + } + let modes = self.half_layout.0.outer().modes(); + if modes.len() != 2 + || modes[0].checked_size() != Some(128) + || modes[1].checked_size() != Some(64) + || self.full_elements() != Some(128 * 128) + { + return verify_err_noloc!( + "cute.epilogue_plan halves must cover one 128x128 result tile" + ); + } + if cute_layout::tma_phase_alignment_bytes(&self.half_layout.0, 2) != Ok(1024) { + return verify_err_noloc!( + "cute.epilogue_plan half layout must use the 1024-byte B128 phase alignment" + ); + } + if self.tiled_mma.accumulator_registers_per_lane() != Some(64) { + return verify_err_noloc!( + "cute.epilogue_plan v0 needs 64 FP32 accumulator registers per lane" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/gemv.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/gemv.rs new file mode 100644 index 0000000000..15df89d103 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/gemv.rs @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Packed-value and scale-view contracts for the block-scaled GEMV flow. + +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +/// Which part of a block-scaled tensor is currently selected. +/// +/// ```text +/// Full ── select (batch, row) ──► Row ── select K tile ──► KTile<64> +/// ``` +#[pliron_attr(name = "cute.scaled_layout", format)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteScaledLayoutAttr { + /// The complete bound value/scale tensor before selecting a row. + Full, + /// One selected row across its complete reduction dimension. + Row, + /// One selected reduction tile with the stored logical K width. + KTile(#[doc = "Logical K elements in one selected fragment."] u64), +} + +impl Verify for CuteScaledLayoutAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if matches!(self, Self::KTile(0)) { + return verify_err_noloc!("cute.scaled_layout K tile width must be greater than zero"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/layout.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/layout.rs new file mode 100644 index 0000000000..ef4d3f30a5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/layout.rs @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Structural layouts and swizzles in the CuTe attribute syntax. + +use core::fmt; + +use combine::Parser; +use cute_layout::{ComposedLayout, Layout, OffsetUnit, ParseLayoutError, Swizzle}; +use pliron::{ + impl_printable_for_display, + irfmt::parsers::quoted_string_parser, + parsable::{Parsable, ParseResult, StateStream}, +}; +use pliron_derive::pliron_attr; + +/// A complete swizzled layout, stored as one typed attribute. +/// +/// The quoted payload is intentionally simple and round-trippable. Its +/// fields are CuTe's inner swizzle, offset, unit, and outer layout, in that +/// order: +/// +/// ```text +/// "3,4,3;0;elements;(8,32):(32,1)" +/// B M S off unit outer layout +/// ``` +#[pliron_attr(name = "cute.composed_layout", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteComposedLayoutAttr( + #[doc = "Structural inner swizzle, offset, and outer layout."] pub ComposedLayout, +); + +impl fmt::Display for CuteComposedLayoutAttr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let swizzle = self.0.inner(); + write!( + f, + "\"{},{},{};{};{};{}\"", + swizzle.num_bits, + swizzle.num_base, + swizzle.num_shift, + self.0.offset(), + self.0.unit(), + self.0.outer() + ) + } +} + +impl_printable_for_display!(CuteComposedLayoutAttr); + +/// Parse the structural inner swizzle, offset, unit, and outer layout from the printed attribute syntax. +pub(super) fn parse_composed_layout( + text: String, +) -> Result { + let mut fields = text.splitn(4, ';'); + let swizzle = fields.next().ok_or(ParseLayoutError)?; + let offset = fields + .next() + .ok_or(ParseLayoutError)? + .parse::() + .map_err(|_| ParseLayoutError)?; + let unit = match fields.next().ok_or(ParseLayoutError)? { + "elements" => OffsetUnit::Elements, + "bytes" => OffsetUnit::Bytes, + _ => return Err(ParseLayoutError), + }; + let outer = fields.next().ok_or(ParseLayoutError)?.parse::()?; + let swizzle_fields = swizzle.split(',').map(str::trim).collect::>(); + let [bits, base, shift] = swizzle_fields.as_slice() else { + return Err(ParseLayoutError); + }; + let bits = bits.parse::().map_err(|_| ParseLayoutError)?; + let base = base.parse::().map_err(|_| ParseLayoutError)?; + let shift = shift.parse::().map_err(|_| ParseLayoutError)?; + let inner = Swizzle::try_new(bits, base, shift).ok_or(ParseLayoutError)?; + ComposedLayout::new(inner, offset, outer, unit) + .map(CuteComposedLayoutAttr) + .map_err(|_| ParseLayoutError) +} + +impl Parsable for CuteComposedLayoutAttr { + type Arg = (); + type Parsed = CuteComposedLayoutAttr; + + fn parse<'a>( + state_stream: &mut StateStream<'a>, + _arg: Self::Arg, + ) -> ParseResult<'a, Self::Parsed> { + quoted_string_parser() + .and_then(parse_composed_layout) + .parse_stream(state_stream) + .into() + } +} + +/// A CuTe layout (`shape:stride`) as a first-class typed attribute. +#[pliron_attr(name = "cute.layout", verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteLayoutAttr(#[doc = "Structural shape and stride mapping."] pub Layout); + +impl fmt::Display for CuteLayoutAttr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "\"{}\"", self.0) + } +} + +impl_printable_for_display!(CuteLayoutAttr); + +impl Parsable for CuteLayoutAttr { + type Arg = (); + type Parsed = CuteLayoutAttr; + + fn parse<'a>( + state_stream: &mut StateStream<'a>, + _arg: Self::Arg, + ) -> ParseResult<'a, Self::Parsed> { + quoted_string_parser() + .and_then(|s: String| s.parse::().map(CuteLayoutAttr)) + .parse_stream(state_stream) + .into() + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mma.rs new file mode 100644 index 0000000000..b54d87facb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mma.rs @@ -0,0 +1,386 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared tensor placement, register carriers, and tiled matrix-instruction contracts. + +use cute_layout::{ComposedLayout, Layout, OffsetUnit}; +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +use super::{ + layout::{CuteComposedLayoutAttr, CuteLayoutAttr}, + tensor::CuteTensorFormatAttr, +}; + +/// Register type accumulated by one matrix-multiply atom. +/// +/// The first shared-memory MMA slice only needs FP32 accumulation. Keeping it +/// named makes the atom description extensible without hiding the choice in +/// an integer code. +#[pliron_attr(name = "cute.mma_accumulator", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteMmaAccumulatorAttr { + /// FP32 scalar storage or accumulation. + F32, +} + +/// What one ordinary MIR aggregate means at a semantic MMA operation. +/// +/// B is deliberately absent from this first list. A B partition is a lazy +/// pointer-and-coordinate view, not an eagerly loaded register fragment: +/// +/// ```text +/// ScaleStage ── K half ──► ScaleK64 +/// A shared tile ─────────► A +/// zero / previous step ──► Accumulator +/// B shared tile ─────────► pointer + capacity + coordinates +/// ``` +#[pliron_attr(name = "cute.mma_carrier_kind", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteMmaCarrierKindAttr { + /// All scale registers needed by one staged K region. + ScaleStage, + /// Scale registers selected for one K=64 matrix-instruction step. + ScaleK64, + /// Packed registers for one loaded A fragment. + A, + /// FP32 output registers carried between compute steps. + Accumulator, +} + +/// One warp-level matrix-multiply instruction shape. +/// +/// Register counts are per lane. They describe the physical MIR aggregates +/// consumed by the selected backend's matrix-instruction lowering: +/// +/// ```text +/// A: 4 x u32 B: 2 x u32 C: 4 x f32 +/// m16n8k64 +/// ``` +#[pliron_attr( + name = "cute.mma_atom", + format = "`<` $m `,` $n `,` $k `,` $a_format `,` $b_format `,` $scale_format `,` $values_per_scale `,` $accumulator `,` $threads `,` $a_registers `,` $b_registers `,` $accumulator_registers `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteMmaAtomAttr { + /// Collective output-row extent M. + pub m: u64, + /// Collective output-column extent N. + pub n: u64, + /// Reduction extent K, in logical elements. + pub k: u64, + /// Logical encoding of packed A values. + pub a_format: CuteTensorFormatAttr, + /// Logical encoding of packed B values. + pub b_format: CuteTensorFormatAttr, + /// Encoding used by block-scale values. + pub scale_format: CuteTensorFormatAttr, + /// Number of consecutive logical K values sharing one scale. + pub values_per_scale: u64, + /// Scalar format used by the accumulator registers. + pub accumulator: CuteMmaAccumulatorAttr, + /// Participating threads in one matrix-instruction atom. + pub threads: u32, + /// Packed u32 A registers per participating lane. + pub a_registers: u32, + /// Packed u32 B registers per participating lane. + pub b_registers: u32, + /// FP32 accumulator registers per lane. + pub accumulator_registers: u32, +} + +impl CuteMmaAtomAttr { + /// SM120's packed MXFP4 `m16n8k64` atom (`UE8M0` scales, `SF_VEC = 32`). + #[must_use] + pub const fn mxf4_m16n8k64() -> Self { + Self { + m: 16, + n: 8, + k: 64, + a_format: CuteTensorFormatAttr::E2M1, + b_format: CuteTensorFormatAttr::E2M1, + scale_format: CuteTensorFormatAttr::UE8M0, + values_per_scale: 32, + accumulator: CuteMmaAccumulatorAttr::F32, + threads: 32, + a_registers: 4, + b_registers: 2, + accumulator_registers: 4, + } + } + + /// SM120's packed NVFP4 `m16n8k64` atom (`E4M3` scales, `SF_VEC = 16`). + #[must_use] + pub const fn nvf4_m16n8k64() -> Self { + Self { + m: 16, + n: 8, + k: 64, + a_format: CuteTensorFormatAttr::E2M1, + b_format: CuteTensorFormatAttr::E2M1, + scale_format: CuteTensorFormatAttr::E4M3, + values_per_scale: 16, + accumulator: CuteMmaAccumulatorAttr::F32, + threads: 32, + a_registers: 4, + b_registers: 2, + accumulator_registers: 4, + } + } +} + +impl Verify for CuteMmaAtomAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.m == 0 || self.n == 0 || self.k == 0 { + return verify_err_noloc!("cute.mma_atom M, N, and K must be greater than zero"); + } + if self.a_format != CuteTensorFormatAttr::E2M1 + || self.b_format != CuteTensorFormatAttr::E2M1 + || !matches!( + self.scale_format, + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3 + ) + { + return verify_err_noloc!( + "cute.mma_atom v0 needs E2M1 A/B values with UE8M0 or E4M3 scales" + ); + } + if self.values_per_scale == 0 || !self.k.is_multiple_of(self.values_per_scale) { + return verify_err_noloc!( + "cute.mma_atom K must contain a whole positive number of scale groups" + ); + } + if self.threads == 0 + || self.a_registers == 0 + || self.b_registers == 0 + || self.accumulator_registers == 0 + { + return verify_err_noloc!( + "cute.mma_atom thread and per-lane register counts must be greater than zero" + ); + } + Ok(()) + } +} + +/// How warp-level MMA atoms cover one CTA output tile. +/// +/// `m_ownership` and `n_ownership` map a warp position plus a cell position +/// to the corresponding M or N atom number. The current 128 x 128 plan reads: +/// +/// ```text +/// M atoms: (warp_m, cell_m) (4,2):(1,4) +/// N atoms: ((within_pair,pair), ((2,4),2):((1,4),2) +/// warp_n) +/// +/// 4 warp-M positions x 2 warp-N positions = 8 compute warps +/// one warp owns 2 M atoms x 8 N atoms = 16 accumulator cells +/// ``` +/// +/// `b_load_group` is the number of B fragments loaded before immediately +/// consuming them. It keeps the performance-critical lazy-B schedule visible. +#[pliron_attr( + name = "cute.tiled_mma_plan", + format = "`<` $atom `,` $cta_m `,` $cta_n `,` $cta_k `,` $warp_m `,` $warp_n `,` $b_load_group `,` $shared_layout `,` $m_ownership `,` $n_ownership `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteTiledMmaPlanAttr { + /// Warp-level instruction shape and register contract. + pub atom: CuteMmaAtomAttr, + /// Output rows covered by one CTA tile. + pub cta_m: u64, + /// Output columns covered by one CTA tile. + pub cta_n: u64, + /// Reduction elements staged for one CTA compute step. + pub cta_k: u64, + /// Number of warp positions along the output M axis. + pub warp_m: u32, + /// Number of warp positions along the output N axis. + pub warp_n: u32, + /// B fragments loaded together before their immediate MMA use. + pub b_load_group: u32, + /// Physical shared placement expected by the tiled-MMA data loads. + pub shared_layout: CuteComposedLayoutAttr, + /// Layout mapping warp/cell coordinates to M-atom ownership. + pub m_ownership: CuteLayoutAttr, + /// Layout mapping warp/cell coordinates to N-atom ownership. + pub n_ownership: CuteLayoutAttr, +} + +impl CuteTiledMmaPlanAttr { + /// Describe the current 128 x 128 x 128 SM120 MXFP4 tiled MMA. + #[must_use] + pub fn mxf4_128x128x128(shared_layout: ComposedLayout) -> Self { + Self { + atom: CuteMmaAtomAttr::mxf4_m16n8k64(), + cta_m: 128, + cta_n: 128, + cta_k: 128, + warp_m: 4, + warp_n: 2, + b_load_group: 2, + shared_layout: CuteComposedLayoutAttr(shared_layout), + m_ownership: CuteLayoutAttr( + "(4,2):(1,4)" + .parse() + .expect("fixed M ownership layout is valid"), + ), + n_ownership: CuteLayoutAttr( + "((2,4),2):((1,4),2)" + .parse() + .expect("fixed N ownership layout is valid"), + ), + } + } + + /// Describe the 128 x 128 x 128 SM120 NVFP4 (`E4M3` scale, `SF_VEC = 16`) tiled MMA. + #[must_use] + pub fn nvf4_128x128x128(shared_layout: ComposedLayout) -> Self { + Self { + atom: CuteMmaAtomAttr::nvf4_m16n8k64(), + ..Self::mxf4_128x128x128(shared_layout) + } + } + + /// Number of compute warps in this plan. + #[must_use] + pub fn compute_warps(&self) -> Option { + u64::from(self.warp_m).checked_mul(u64::from(self.warp_n)) + } + + /// Number of M atoms owned by one warp. + #[must_use] + pub fn m_atoms_per_warp(&self) -> Option { + self.cta_m + .checked_div(self.atom.m)? + .checked_div(u64::from(self.warp_m)) + } + + /// Number of N atoms owned by one warp. + #[must_use] + pub fn n_atoms_per_warp(&self) -> Option { + self.cta_n + .checked_div(self.atom.n)? + .checked_div(u64::from(self.warp_n)) + } + + /// Packed scale words held by one lane for one K=128 stage. + #[must_use] + pub fn scale_words_per_lane(&self) -> Option { + self.m_atoms_per_warp()? + .checked_add(self.n_atoms_per_warp()?) + } + + /// A-fragment u32 registers held by one lane for one K=64 step. + #[must_use] + pub fn a_registers_per_lane(&self) -> Option { + self.m_atoms_per_warp()? + .checked_mul(u64::from(self.atom.a_registers)) + } + + /// FP32 accumulator registers held by one lane for the complete warp tile. + #[must_use] + pub fn accumulator_registers_per_lane(&self) -> Option { + self.m_atoms_per_warp()? + .checked_mul(self.n_atoms_per_warp()?)? + .checked_mul(u64::from(self.atom.accumulator_registers)) + } +} + +/// Check that all logical coordinates cover one dense physical interval exactly once. +fn layout_is_dense_permutation(layout: &Layout) -> bool { + let Some(size) = layout.checked_size() else { + return false; + }; + let mut offsets = (0..size) + .map(|coordinate| layout.checked_call(&cute_layout::IntTuple::Leaf(coordinate))) + .collect::>>(); + let Some(ref mut offsets) = offsets else { + return false; + }; + offsets.sort_unstable(); + offsets.iter().copied().eq(0..size) +} + +impl Verify for CuteTiledMmaPlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.atom.verify(ctx)?; + if self.cta_m == 0 || self.cta_n == 0 || self.cta_k == 0 { + return verify_err_noloc!( + "cute.tiled_mma_plan CTA M, N, and K must be greater than zero" + ); + } + if !self.cta_m.is_multiple_of(self.atom.m) + || !self.cta_n.is_multiple_of(self.atom.n) + || !self.cta_k.is_multiple_of(self.atom.k) + { + return verify_err_noloc!("cute.tiled_mma_plan CTA tile must contain whole MMA atoms"); + } + if self.warp_m == 0 || self.warp_n == 0 { + return verify_err_noloc!( + "cute.tiled_mma_plan warp grid sizes must be greater than zero" + ); + } + let m_atoms = self.cta_m / self.atom.m; + let n_atoms = self.cta_n / self.atom.n; + if !m_atoms.is_multiple_of(u64::from(self.warp_m)) + || !n_atoms.is_multiple_of(u64::from(self.warp_n)) + { + return verify_err_noloc!( + "cute.tiled_mma_plan atom grid must divide evenly across its warp grid" + ); + } + let Some(compute_warps) = self.compute_warps() else { + return verify_err_noloc!("cute.tiled_mma_plan compute-warp count overflowed"); + }; + if compute_warps == 0 || compute_warps > 32 { + return verify_err_noloc!( + "cute.tiled_mma_plan must use between 1 and 32 compute warps" + ); + } + let Some(n_per_warp) = self.n_atoms_per_warp() else { + return verify_err_noloc!("cute.tiled_mma_plan N ownership is not divisible"); + }; + if self.b_load_group == 0 || n_per_warp % u64::from(self.b_load_group) != 0 { + return verify_err_noloc!( + "cute.tiled_mma_plan B load group must divide one warp's N atoms" + ); + } + if self.shared_layout.0.unit() != OffsetUnit::Elements || self.shared_layout.0.offset() != 0 + { + return verify_err_noloc!( + "cute.tiled_mma_plan shared placement must use zero-based element offsets" + ); + } + let expected_shared_elements = self + .cta_m + .checked_mul(self.cta_k) + .and_then(|values| values.checked_div(4)); + if expected_shared_elements.and_then(|value| i64::try_from(value).ok()) + != self.shared_layout.0.outer().checked_size() + { + return verify_err_noloc!( + "cute.tiled_mma_plan shared placement must hold four packed E2M1 values per f16 carrier" + ); + } + if self.m_ownership.0.checked_size() != i64::try_from(m_atoms).ok() + || self.n_ownership.0.checked_size() != i64::try_from(n_atoms).ok() + || !layout_is_dense_permutation(&self.m_ownership.0) + || !layout_is_dense_permutation(&self.n_ownership.0) + { + return verify_err_noloc!( + "cute.tiled_mma_plan ownership layouts must cover every M/N atom exactly once" + ); + } + if self.m_ownership.0.shape.flatten().first().copied() != Some(i64::from(self.warp_m)) + || self.n_ownership.0.shape.flatten().last().copied() != Some(i64::from(self.warp_n)) + { + return verify_err_noloc!( + "cute.tiled_mma_plan ownership layouts must expose warp-M first and warp-N last" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mod.rs new file mode 100644 index 0000000000..522967231a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/mod.rs @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed attributes for the cute dialect. +//! +//! Layouts are stored structurally (the `cute-layout` types), never as +//! strings; only the *textual syntax* is the CuTe notation in quotes, e.g. +//! `"(8,32):(32,1)"`, so module dumps read like CuTe IR. + +/// Typed copy widths, divisibility promises, and warp matrix operand roles. +mod copy; +/// Shared output placement and asynchronous store synchronization contracts. +mod epilogue; +/// Packed-value and scale-view contracts for the block-scaled GEMV flow. +mod gemv; +/// Structural layouts and swizzles in the CuTe attribute syntax. +mod layout; +/// Shared tensor placement, register carriers, and tiled matrix-instruction contracts. +mod mma; +/// Barrier-ring configuration and producer/consumer cursor contracts. +mod pipeline; +/// Persistent output-tile grids and their ordinary cursor provenance. +mod scheduler; +/// Tensor storage, access, format, and selection metadata. +mod tensor; + +pub use copy::{CuteCopyAtomAttr, CuteDivisibilityAttr, CuteMatrixRoleAttr}; +pub use epilogue::{ + CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpiloguePlanAttr, + CuteEpilogueSyncPhaseAttr, CuteTmaStorePipelineAttr, +}; +pub use gemv::CuteScaledLayoutAttr; +pub use layout::{CuteComposedLayoutAttr, CuteLayoutAttr}; +pub use mma::{ + CuteMmaAccumulatorAttr, CuteMmaAtomAttr, CuteMmaCarrierKindAttr, CuteTiledMmaPlanAttr, +}; +pub use pipeline::{CutePipelineRoleAttr, CutePipelineStateAttr}; +pub use scheduler::CuteTileGridAttr; +pub use tensor::{ + CuteAlignmentAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, + CuteTensorLayoutAttr, CuteTensorRoleAttr, SM1XX_BLOCK_SCALE_ATOM_BYTES, + SM1XX_BLOCK_SCALE_GROUPS_PER_ATOM, SM1XX_BLOCK_SCALE_ROWS_PER_ATOM, +}; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/pipeline.rs new file mode 100644 index 0000000000..1207147f8b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/pipeline.rs @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Barrier-ring configuration and producer/consumer cursor contracts. + +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +/// Which side of a staged load owns one ring position. +/// +/// ```text +/// Producer waits for an empty stage, then starts TMA +/// Consumer waits for a full stage, then reads and releases it +/// ``` +#[pliron_attr(name = "cute.pipeline_role", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CutePipelineRoleAttr { + /// Cursor side waiting for empty storage before starting an asynchronous transfer. + Producer, + /// Cursor side waiting for completed input before reading and releasing storage. + Consumer, +} + +/// Compile-time facts attached to the ordinary `(slot, phase)` scalars. +/// +/// The attribute says how to interpret those scalars. It is not a runtime +/// state object: +/// +/// ```text +/// pipeline_state +/// │ +/// └── slot: u32, phase: u32 +/// ``` +#[pliron_attr(name = "cute.pipeline_state", format = "`<` $role `,` $stages `>`")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CutePipelineStateAttr { + /// Producer or consumer interpretation of the cursor. + pub role: CutePipelineRoleAttr, + /// Number of reusable slots in the barrier or output-buffer ring. + pub stages: u64, +} + +impl CutePipelineStateAttr { + /// Describe producer or consumer state for one fixed-size stage ring. + #[must_use] + pub const fn new(role: CutePipelineRoleAttr, stages: u64) -> Self { + Self { role, stages } + } + + /// Describe producer state for one fixed-size stage ring. + #[must_use] + pub const fn producer(stages: u64) -> Self { + Self::new(CutePipelineRoleAttr::Producer, stages) + } + + /// Describe consumer state for one fixed-size stage ring. + #[must_use] + pub const fn consumer(stages: u64) -> Self { + Self::new(CutePipelineRoleAttr::Consumer, stages) + } +} + +impl Verify for CutePipelineStateAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.stages == 0 || self.stages > u64::from(u32::MAX) { + return verify_err_noloc!( + "cute.pipeline_state stages must be between 1 and {}, got {}", + u32::MAX, + self.stages + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/scheduler.rs new file mode 100644 index 0000000000..e7d706c237 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/scheduler.rs @@ -0,0 +1,69 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent output-tile grids and their ordinary cursor provenance. + +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +/// The logical output-tile grid owned by a static scheduler. +/// +/// M changes first, then N, then batch (also called L): +/// +/// ```text +/// grid<3, 2, 1> +/// +/// linear: 0 1 2 3 4 5 +/// tile: 0,0 1,0 2,0 0,1 1,1 2,1 +/// ``` +/// +/// All three sizes are positive. Their product must fit in `u64`, which is +/// the scalar type carried through the scheduler loop. +#[pliron_attr( + name = "cute.tile_grid", + format = "`<` $m_tiles `,` $n_tiles `,` $batches `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteTileGridAttr { + /// Number of logical output tiles along M. + pub m_tiles: u64, + /// Number of logical output tiles along N. + pub n_tiles: u64, + /// Number of independent output batches. + pub batches: u64, +} + +impl CuteTileGridAttr { + /// Describe a grid where M changes first, then N, then batch. + #[must_use] + pub const fn new(m_tiles: u64, n_tiles: u64, batches: u64) -> Self { + Self { + m_tiles, + n_tiles, + batches, + } + } + + /// Return the number of logical work tiles when multiplication is safe. + #[must_use] + pub const fn total_tiles(self) -> Option { + let Some(mn) = self.m_tiles.checked_mul(self.n_tiles) else { + return None; + }; + mn.checked_mul(self.batches) + } +} + +impl Verify for CuteTileGridAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.m_tiles == 0 || self.n_tiles == 0 || self.batches == 0 { + return verify_err_noloc!("cute.tile_grid sizes must be greater than zero"); + } + if self.total_tiles().is_none() { + return verify_err_noloc!("cute.tile_grid tile count must fit in u64"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tensor.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tensor.rs new file mode 100644 index 0000000000..0112119c9c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tensor.rs @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor storage, access, format, and selection metadata. + +use pliron::{common_traits::Verify, context::Context, result::Error, verify_err_noloc}; +use pliron_derive::pliron_attr; + +/// Where a tensor view's storage lives. +/// +/// The first tensor flow starts in global memory. TMA adds a shared-memory +/// destination without changing the meaning of the existing global variant: +/// +/// ```text +/// host buffer ── kernel argument ──► global GPU memory +/// │ TMA +/// ▼ +/// CTA shared memory +/// ``` +/// +/// Keeping the choice typed avoids turning address-space numbers such as `1` +/// or `3` into a hidden contract. +#[pliron_attr(name = "cute.tensor_space", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteTensorAddressSpaceAttr { + /// Device global memory used by tensor inputs and outputs. + Gmem, + /// CTA-local shared memory, visible to every thread in one block. + Smem, +} + +/// Whether a tensor view may only read or may also write. +/// +/// ```text +/// ReadOnly storage ──► registers +/// ReadWrite storage ◄── registers +/// ``` +#[pliron_attr(name = "cute.tensor_access", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteTensorAccessAttr { + /// Storage may be read but must not be written through this view. + ReadOnly, + /// Storage may be both read and written through this view. + ReadWrite, +} + +/// How to read the bits stored by a tensor view. +/// +/// `Plain` is the elementwise path: one `f16` or `f32` storage value already +/// is the logical value. The packed GEMV path keeps the other two meanings +/// visible: +/// +/// ```text +/// E2M1 two four-bit values in global u8, or four in shared f16 +/// UE8M0 one scale in global u8, or four packed scales in shared u32 +/// ``` +#[pliron_attr(name = "cute.tensor_format", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteTensorFormatAttr { + /// Storage elements already have the tensor's logical scalar meaning. + Plain, + /// Packed four-bit floating-point values requiring decoding for scalar arithmetic. + E2M1, + /// Unsigned exponent-only scale values for block-scaled arithmetic. + UE8M0, + /// Eight-bit E4M3FN floating-point scale values for NVFP4 block-scaled arithmetic. + E4M3, +} + +/// Which logical row coordinate a tensor uses. +/// +/// ```text +/// Mkl rows follow M: A[m, k, l] +/// Nkl rows follow N: B[n, k, l] +/// Generic elementwise data +/// ``` +/// +/// `Nkl` is a single vector row when GEMV has `N = 1`, and a full B matrix +/// when GEMM has many N rows. Keeping the coordinate name avoids baking one +/// kernel profile into the tensor type. +/// +/// The role is a compile-time fact. It does not change the stored bytes. +#[pliron_attr(name = "cute.tensor_role", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteTensorRoleAttr { + /// No matrix-operand coordinate role is assigned to the storage. + Generic, + /// Rows address A's M coordinate, with reduction K and batch L. + Mkl, + /// Rows address B's N coordinate, with reduction K and batch L. + Nkl, +} + +/// The one-dimensional layout carried by a v0 tensor view. +/// +/// The tile width stays inside the layout instead of living in a nearby +/// untyped integer: +/// +/// ```text +/// Contiguous1D [0 1 2 3 4 5 6 7] +/// Zipped1D<4> [0 1 2 3] [4 5 6 7] +/// Tile1D<4> ^ one selected group ^ +/// ``` +#[pliron_attr(name = "cute.tensor_layout", format)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteTensorLayoutAttr { + /// A complete contiguous one-dimensional tensor before tiling. + Contiguous1D, + /// The tensor divided into groups of the stored element width. + Zipped1D(#[doc = "Elements in each logical tile."] u64), + /// One selected tile whose width is stored in this variant. + Tile1D(#[doc = "Elements in the selected logical tile."] u64), + /// Logical K values are adjacent. Packed formats may place more than one + /// logical value in each storage element. + KMajor, + /// Blackwell's canonical scale order. The number says how many logical + /// values share one scale. The fixed atom geometry is named by the + /// `SM1XX_BLOCK_SCALE_*` constants in this module. + BlockScaleKMajor(#[doc = "Logical K values sharing one scale."] u64), + /// A two-dimensional tile moved as raw carrier values by TMA. + /// + /// The exact row/column shape, stride, and swizzle live in the enclosing + /// `cute.tma_view` type as a composed layout. + Tma2D, +} + +impl CuteTensorLayoutAttr { + /// Return the tile width for a zipped or selected-tile layout. + #[must_use] + pub const fn tile_size(self) -> Option { + match self { + Self::Contiguous1D | Self::KMajor | Self::BlockScaleKMajor(_) | Self::Tma2D => None, + Self::Zipped1D(size) | Self::Tile1D(size) => Some(size), + } + } + + /// Return how many logical K values share one scale, when this is the + /// canonical block-scale layout. + #[must_use] + pub const fn values_per_scale(self) -> Option { + match self { + Self::BlockScaleKMajor(size) => Some(size), + _ => None, + } + } +} + +/// Logical rows inside one Blackwell canonical scale atom. +pub const SM1XX_BLOCK_SCALE_ROWS_PER_ATOM: u64 = 128; +/// Consecutive K-scale groups inside one Blackwell canonical scale atom. +pub const SM1XX_BLOCK_SCALE_GROUPS_PER_ATOM: u64 = 4; +/// Physical bytes reserved by one Blackwell canonical scale atom. +pub const SM1XX_BLOCK_SCALE_ATOM_BYTES: u64 = 512; + +impl Verify for CuteTensorLayoutAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.tile_size().is_some_and(|size| size == 0) { + return verify_err_noloc!("cute.tensor_layout tile size must be greater than zero"); + } + if matches!(self, Self::BlockScaleKMajor(0)) { + return verify_err_noloc!( + "cute.tensor_layout values per scale must be greater than zero" + ); + } + Ok(()) + } +} + +/// A byte-alignment promise attached to a tensor view or vector transfer. +/// +/// This is a promise, not padding owned by the IR. For example, `16` says +/// the address used by a full tile is safe for one 16-byte transaction. +#[pliron_attr(name = "cute.alignment_bytes", format = "$0")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteAlignmentAttr(#[doc = "Promised alignment in bytes."] pub u64); + +impl Verify for CuteAlignmentAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.0 == 0 || !self.0.is_power_of_two() { + return verify_err_noloc!( + "cute.alignment_bytes must be a positive power of two, got {}", + self.0 + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tests.rs new file mode 100644 index 0000000000..6389b2979a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/attributes/tests.rs @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe operation, type, and semantic verification contracts. + +use cute_layout::{ComposedLayout, IntTuple, Layout, OffsetUnit, Swizzle}; +use pliron::{common_traits::Verify, context::Context}; + +use super::{layout::parse_composed_layout, *}; + +/// Printing and parsing preserve the swizzle, offset, unit, shape, and stride together. +#[test] +fn composed_layout_text_round_trips_every_addressing_piece() { + let outer = Layout::new( + IntTuple::Tuple(vec![IntTuple::Leaf(8), IntTuple::Leaf(32)]), + IntTuple::Tuple(vec![IntTuple::Leaf(32), IntTuple::Leaf(1)]), + ); + let composed = + ComposedLayout::new(Swizzle::new(3, 4, 3), 8, outer, OffsetUnit::Elements).unwrap(); + let attr = CuteComposedLayoutAttr(composed); + let printed = attr.to_string(); + let reparsed = parse_composed_layout(printed.trim_matches('"').to_string()).unwrap(); + assert_eq!(reparsed, attr); +} + +/// Parsed swizzle bit fields cannot overlap in ways that violate the layout contract. +#[test] +fn composed_layout_text_rejects_overlapping_swizzle_fields() { + // B=4 cannot fit into a shift distance of 3: + // + // [ source ] + // [ target ] <- overlap, so XOR is not reversible + assert!(parse_composed_layout("4,3,3;0;bytes;16:4".to_string()).is_err()); +} + +/// The textual attribute syntax preserves valid negative swizzle shifts without normalization loss. +#[test] +fn composed_layout_text_preserves_negative_cute_shifts() { + let attr = parse_composed_layout("2,0,-3;0;bytes;16:4".to_string()).unwrap(); + assert_eq!(attr.0.inner(), Swizzle::new(2, 0, -3)); + assert_eq!(attr.to_string(), "\"2,0,-3;0;bytes;16:4\""); +} + +/// Zero tile/scale widths and invalid alignment promises fail attribute verification. +#[test] +fn tensor_structure_rejects_impossible_sizes() { + let ctx = Context::new(); + assert!(CuteTensorLayoutAttr::Contiguous1D.verify(&ctx).is_ok()); + assert!(CuteTensorLayoutAttr::Zipped1D(4).verify(&ctx).is_ok()); + assert!(CuteTensorLayoutAttr::Tile1D(0).verify(&ctx).is_err()); + assert!( + CuteTensorLayoutAttr::BlockScaleKMajor(16) + .verify(&ctx) + .is_ok() + ); + assert!( + CuteTensorLayoutAttr::BlockScaleKMajor(0) + .verify(&ctx) + .is_err() + ); + assert!(CuteScaledLayoutAttr::KTile(64).verify(&ctx).is_ok()); + assert!(CuteScaledLayoutAttr::KTile(0).verify(&ctx).is_err()); + assert!(CuteAlignmentAttr(16).verify(&ctx).is_ok()); + assert!(CuteAlignmentAttr(12).verify(&ctx).is_err()); + assert!(CuteAlignmentAttr(0).verify(&ctx).is_err()); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/lib.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/lib.rs new file mode 100644 index 0000000000..65abf25fda --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/lib.rs @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CuTe dialect definition. +//! +//! Represents CuTe layouts and operations in Pliron, alongside MIR control +//! flow and scalar operations. +//! +//! Layouts are stored as typed attributes backed by `cute-layout` (never as +//! strings). Layout operations compute their result types when the IR is +//! built. The verifier checks how tensors, copies, pipelines, and MMA +//! operations fit together before a backend lowers them. +//! +//! The implemented surface is deliberately narrow: elementwise tensor tiles, +//! one block-scaled GEMV fragment flow, raw-carrier TMA views for the four GEMM +//! stage copies, shared-tensor/tiled-MMA compute, and SM100 two-CTA tensor +//! memory with collector-aware MMA, cluster TMA, and asynchronous pipelines. +//! The selected backend lowers these operations to GPU code. +//! +//! Operation families live under [`ops`]. The root `*_ops` modules are +//! compatibility aliases to the same modules; [`attributes`], [`types`], +//! and [`verify`] retain their separate roles. + +pub mod attributes; +pub mod ops; +pub mod types; +pub mod verify; + +/// Compatibility alias for [`ops::copy`]. +pub use ops::copy as copy_ops; +/// Compatibility alias for [`ops::epilogue`]. +pub use ops::epilogue as epilogue_ops; +/// Compatibility alias for [`ops::gemm_tma`]. +pub use ops::gemm_tma as gemm_tma_ops; +/// Compatibility alias for [`ops::gemv`]. +pub use ops::gemv as gemv_ops; +/// Compatibility alias for [`ops::mma`]. +pub use ops::mma as smem_mma_ops; +/// Compatibility alias for [`ops::pipeline`]. +pub use ops::pipeline as pipeline_ops; +/// Compatibility alias for [`ops::scheduler`]. +pub use ops::scheduler as scheduler_ops; +/// Compatibility alias for [`ops::sm100`]. +pub use ops::sm100 as sm100_ops; +/// Compatibility alias for [`ops::sm100_attention`]. +pub use ops::sm100_attention as sm100_attention_ops; +/// Compatibility alias for [`ops::tensor`]. +pub use ops::tensor as tensor_ops; + +/// Layout math shared by the importer, verifier, and backends. +pub use cute_layout as layout; + +use pliron::attribute::Attribute; +use pliron::context::Context; +use pliron::dialect::{Dialect, DialectName}; +use pliron::op::Op; +use pliron::r#type::Type; + +/// Namespace used for all CuTe operation and attribute identifiers. +pub const CUTE_DIALECT_NAME: &str = "cute"; + +/// Register the CuTe dialect and its supported operation, attribute, and type families. +pub fn register(ctx: &mut Context) { + Dialect::register( + ctx, + &DialectName::try_new(CUTE_DIALECT_NAME).expect("valid dialect name"), + ); + + ops::copy::CuteStoreGlobalNoAllocateOp::register(ctx); + ops::copy::CuteGlobalStoreAttr::register(ctx); + ops::sm100::register(ctx); + ops::sm100_attention::register(ctx); + + // The #[pliron_op]/#[pliron_attr] macros auto-register via Context::default; + // registering them again is safe and matches the other dialects. + ops::CuteCopyOp::register(ctx); + ops::CuteAssumeDivOp::register(ctx); + ops::CuteCopyG2SOp::register(ctx); + ops::CuteLdmatrixOp::register(ctx); + ops::CuteTmaLoad2dOp::register(ctx); + ops::CuteTmaStore2dOp::register(ctx); + ops::tensor::CuteTensorMakeOp::register(ctx); + ops::tensor::CuteTensorZippedDivideOp::register(ctx); + ops::tensor::CuteTensorSliceOp::register(ctx); + ops::tensor::CuteTensorIsFullOp::register(ctx); + ops::tensor::CuteTensorBaseOp::register(ctx); + ops::tensor::CuteTensorLoadIntoOp::register(ctx); + ops::tensor::CuteTensorStoreFromOp::register(ctx); + ops::tensor::CuteTensorStoreElementAbsOp::register(ctx); + ops::gemv::CuteTensorMake2DOp::register(ctx); + ops::gemv::CuteScaledViewMakeOp::register(ctx); + ops::gemv::CuteScaledViewRowOp::register(ctx); + ops::gemv::CuteScaledViewKTileOp::register(ctx); + ops::gemv::CuteScaledViewLoadOp::register(ctx); + ops::gemv::CuteDotOp::register(ctx); + ops::gemm_tma::CuteTmaGmemViewOp::register(ctx); + ops::gemm_tma::CuteTmaSmemViewOp::register(ctx); + ops::gemm_tma::CuteTmaCopy2dOp::register(ctx); + ops::epilogue::CuteEpilogueSmemOverlayOp::register(ctx); + ops::epilogue::CuteEpilogueWarpSliceOp::register(ctx); + ops::epilogue::CuteEpilogueStoreFragmentOp::register(ctx); + ops::epilogue::CuteEpilogueSyncOp::register(ctx); + ops::epilogue::CuteEpilogueHalfOp::register(ctx); + ops::epilogue::CuteTmaStoreAcquireOp::register(ctx); + ops::epilogue::CuteTmaStoreCommitOp::register(ctx); + ops::epilogue::CuteTmaStoreTailOp::register(ctx); + ops::epilogue::CuteTmaStore2dSemanticOp::register(ctx); + ops::scheduler::CuteSchedulerNew1dOp::register(ctx); + ops::scheduler::CuteSchedulerHasWorkOp::register(ctx); + ops::scheduler::CuteSchedulerCurrentOp::register(ctx); + ops::scheduler::CuteWorkTileCoordinatesOp::register(ctx); + ops::scheduler::CuteSchedulerAdvanceOp::register(ctx); + ops::pipeline::CuteTmaLoadPipelineMakeOp::register(ctx); + ops::pipeline::CuteTmaLoadPipelineInitOp::register(ctx); + ops::pipeline::CutePipelineStateNewOp::register(ctx); + ops::pipeline::CutePipelineStateSlotOp::register(ctx); + ops::pipeline::CutePipelineStateAdvanceOp::register(ctx); + ops::pipeline::CutePipelineProducerAcquireOp::register(ctx); + ops::pipeline::CutePipelineProducerExpectTxOp::register(ctx); + ops::pipeline::CutePipelineConsumerWaitOp::register(ctx); + ops::pipeline::CutePipelineConsumerReleaseOp::register(ctx); + ops::pipeline::CutePipelineProducerTailOp::register(ctx); + ops::mma::CuteSmemTensorOverlayOp::register(ctx); + ops::mma::CuteTiledMmaSliceOp::register(ctx); + ops::mma::CuteFragmentFillOp::register(ctx); + ops::mma::CuteMmaLoadScalesOp::register(ctx); + ops::mma::CuteFragmentSliceKOp::register(ctx); + ops::mma::CuteMmaLoadAOp::register(ctx); + ops::mma::CuteMmaPartitionBOp::register(ctx); + ops::mma::CuteTiledGemmOp::register(ctx); + attributes::CuteLayoutAttr::register(ctx); + attributes::CuteDivisibilityAttr::register(ctx); + attributes::CuteCopyAtomAttr::register(ctx); + attributes::CuteComposedLayoutAttr::register(ctx); + attributes::CuteTensorAddressSpaceAttr::register(ctx); + attributes::CuteTensorAccessAttr::register(ctx); + attributes::CuteTensorFormatAttr::register(ctx); + attributes::CuteTensorRoleAttr::register(ctx); + attributes::CuteTensorLayoutAttr::register(ctx); + attributes::CuteScaledLayoutAttr::register(ctx); + attributes::CuteAlignmentAttr::register(ctx); + attributes::CuteMatrixRoleAttr::register(ctx); + attributes::CuteTileGridAttr::register(ctx); + attributes::CutePipelineRoleAttr::register(ctx); + attributes::CutePipelineStateAttr::register(ctx); + attributes::CuteMmaAccumulatorAttr::register(ctx); + attributes::CuteMmaCarrierKindAttr::register(ctx); + attributes::CuteMmaAtomAttr::register(ctx); + attributes::CuteTiledMmaPlanAttr::register(ctx); + attributes::CuteTmaStorePipelineAttr::register(ctx); + attributes::CuteCountedCtaBarrierAttr::register(ctx); + attributes::CuteEpilogueSyncPhaseAttr::register(ctx); + attributes::CuteEpilogueHalfAttr::register(ctx); + attributes::CuteEpiloguePlanAttr::register(ctx); + types::CuteTensorViewType::register(ctx); + types::CuteTmaViewType::register(ctx); + types::CuteScaledViewType::register(ctx); + types::CuteFragmentType::register(ctx); + types::CuteWorkTileType::register(ctx); + types::CuteTmaLoadPipelineType::register(ctx); + types::CuteSmemTensorType::register(ctx); + types::CuteEpilogueTileType::register(ctx); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/global.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/global.rs new file mode 100644 index 0000000000..50295989da --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/global.rs @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed per-thread cache-aware register-to-global copies. + +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::{ + builtin::types::{FP32Type, IntegerType}, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + value::Value, + verify_err, verify_err_noloc, +}; +use pliron_derive::{pliron_attr, pliron_op}; + +/// Number of adjacent scalar values copied by one thread. +#[pliron_attr(name = "cute.global_store", format = "`<` $elements `>`")] +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct CuteGlobalStoreAttr { + /// Number of contiguous storage elements transferred by this operation. + pub elements: u64, +} +impl Verify for CuteGlobalStoreAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.elements == 0 || self.elements > 16 { + return verify_err_noloc!("global store needs 1..16 elements"); + } + Ok(()) + } +} + +/// Byte size of a scalar accepted by the cache-aware global-store boundary. +pub fn global_store_element_bytes(ctx: &Context, element: TypeHandle) -> Option { + let ty = element.deref(ctx); + if ty.downcast_ref::().is_some() { + Some(2) + } else if ty.downcast_ref::().is_some() + || ty + .downcast_ref::() + .is_some_and(|i| i.width() == 32 && i.is_unsigned()) + { + Some(4) + } else { + None + } +} + +/// Weak global store with an L1 no-allocation cache hint. +/// +/// Operands are a writable global pointer and a local/generic register-carrier +/// pointer with identical element types. The destination's full-width alignment, +/// liveness, bounds, and absence of conflicting accesses are unsafe caller +/// obligations. The source carrier must be 32-byte aligned; the public Rust carrier type +/// guarantees this independently of the unsafe destination contract. +#[pliron_op(name="cute.store_global_no_allocate",format,attributes=(global_store:CuteGlobalStoreAttr))] +pub struct CuteStoreGlobalNoAllocateOp; +impl CuteStoreGlobalNoAllocateOp { + /// Build a detached cache-aware register-to-global store with a static element count. + pub fn new(ctx: &mut Context, destination: Value, source: Value, elements: u64) -> Self { + let op = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![destination, source], + vec![], + 0, + ), + }; + op.set_attr_global_store(ctx, CuteGlobalStoreAttr { elements }); + op + } + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + /// Return the static count of elements transferred by the store. + pub fn elements(&self, ctx: &Context) -> Option { + self.get_attr_global_store(ctx).map(|a| a.elements) + } +} +impl Verify for CuteStoreGlobalNoAllocateOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "global store needs two pointers and no results"); + } + let pointers: Vec<_> = op + .operands() + .map(|v| { + v.get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() + }) + .collect(); + let (Some(dst), Some(src)) = (&pointers[0], &pointers[1]) else { + return verify_err!(op.loc(), "global store operands must be typed pointers"); + }; + if !dst.is_mutable + || dst.address_space != address_space::GLOBAL + || !matches!( + src.address_space, + address_space::LOCAL | address_space::GENERIC + ) + || dst.pointee != src.pointee + { + return verify_err!( + op.loc(), + "global store requires matching writable global destination and local/generic source" + ); + } + let Some(bytes) = global_store_element_bytes(ctx, dst.pointee) else { + return verify_err!( + op.loc(), + "global store supports only f16, f32, or u32 storage" + ); + }; + let Some(n) = self.elements(ctx) else { + return verify_err!(op.loc(), "global store lost its static width"); + }; + if !matches!(n.checked_mul(bytes), Some(4 | 8 | 16 | 32)) { + return verify_err!(op.loc(), "global store needs 4, 8, 16, or 32 bytes"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/mod.rs new file mode 100644 index 0000000000..be5db09006 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/mod.rs @@ -0,0 +1,387 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Direct, cooperative, descriptor-backed, and cache-aware copies. + +mod global; +mod tma; + +pub use global::{CuteGlobalStoreAttr, CuteStoreGlobalNoAllocateOp, global_store_element_bytes}; +pub use tma::{CuteTmaLoad2dOp, CuteTmaStore2dOp}; + +use cute_layout::{ComposedLayout, Layout, validate_cooperative_copy_plan}; +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::{ + builtin::{ + attributes::TypeAttr, + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + value::Value, + verify_err, +}; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteComposedLayoutAttr, CuteCopyAtomAttr, CuteDivisibilityAttr, CuteLayoutAttr, +}; + +/// Frozen zero-result ABI for a block-wide global-to-shared copy. +/// +/// Runtime operands are flattened so backend lowering never depends on Rust +/// struct field layout or traces a value back to a constructor: +/// +/// ```text +/// 0 gmem base 3 leading dim 6 smem base +/// 1 rows 4 tile row 7 smem capacity +/// 2 columns 5 tile column 8 thread index +/// ``` +/// +/// Static layouts and the atom stay in typed attributes. Boundary handling +/// is fixed to uniform `cp.async` source-size predication; commit/wait/barrier +/// remain explicit operations outside this op. Operands 1..=5 and 7 are +/// NVPTX64 Rust `usize` (`u64`); operand 8 is Rust `u32`. The leading-dimension +/// divisor is also an attribute, so M3 never has to rediscover an alignment +/// promise hidden inside the source carrier. +#[pliron_op( + name = "cute.copy_g2s", + format, + interfaces = [NOpdsInterface<9>, NResultsInterface<0>], + attributes = ( + atom_bytes: CuteCopyAtomAttr, + thread_layout: CuteLayoutAttr, + value_layout: CuteLayoutAttr, + tile_layout: CuteLayoutAttr, + smem_layout: CuteComposedLayoutAttr, + leading_dim_divisor: CuteDivisibilityAttr, + copy_g2s_elem: TypeAttr + ) +)] +pub struct CuteCopyG2SOp; + +impl CuteCopyG2SOp { + /// Build a detached cooperative copy with validated-layout metadata and explicit pitch/index carriers. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + operands: [Value; 9], + atom_bytes: u32, + thread_layout: Layout, + value_layout: Layout, + tile_layout: Layout, + smem_layout: ComposedLayout, + leading_dim_divisor: u64, + elem: TypeHandle, + ) -> Self { + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + operands.to_vec(), + vec![], + 0, + ); + let copy = Self { op }; + copy.set_attr_atom_bytes(ctx, CuteCopyAtomAttr(atom_bytes)); + copy.set_attr_thread_layout(ctx, CuteLayoutAttr(thread_layout)); + copy.set_attr_value_layout(ctx, CuteLayoutAttr(value_layout)); + copy.set_attr_tile_layout(ctx, CuteLayoutAttr(tile_layout)); + copy.set_attr_smem_layout(ctx, CuteComposedLayoutAttr(smem_layout)); + copy.set_attr_leading_dim_divisor(ctx, CuteDivisibilityAttr(leading_dim_divisor)); + copy.set_attr_copy_g2s_elem(ctx, TypeAttr::new(elem)); + copy + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +/// Return the supported scalar element width used by cooperative-copy validation. +pub(super) fn cooperative_elem_bytes(ctx: &Context, elem: TypeHandle) -> Option { + let ty = elem.deref(ctx); + if let Some(integer) = ty.downcast_ref::() { + let width = i64::from(integer.width()); + return (width > 0 && width % 8 == 0).then_some(width / 8); + } + if ty + .downcast_ref::() + .is_some() + { + return Some(4); + } + if ty + .downcast_ref::() + .is_some() + { + return Some(8); + } + ty.downcast_ref::().is_some().then_some(2) +} + +impl Verify for CuteCopyG2SOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + if op.get_num_operands() != 9 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.copy_g2s needs exactly 9 operands and 0 results, got {} and {}", + op.get_num_operands(), + op.get_num_results() + ); + } + let Some(atom_bytes) = self.get_attr_atom_bytes(ctx).map(|attr| attr.0) else { + return verify_err!(op.loc(), "cute.copy_g2s must have atom_bytes"); + }; + if !matches!(atom_bytes, 4 | 8 | 16) { + return verify_err!(op.loc(), "cute.copy_g2s atom must be 4, 8, or 16 bytes"); + } + let Some(thread_layout) = self.get_attr_thread_layout(ctx).map(|attr| attr.0.clone()) + else { + return verify_err!(op.loc(), "cute.copy_g2s must have a thread layout"); + }; + let Some(value_layout) = self.get_attr_value_layout(ctx).map(|attr| attr.0.clone()) else { + return verify_err!(op.loc(), "cute.copy_g2s must have a value layout"); + }; + let Some(tile_layout) = self.get_attr_tile_layout(ctx).map(|attr| attr.0.clone()) else { + return verify_err!(op.loc(), "cute.copy_g2s must have a tile layout"); + }; + let Some(smem_layout) = self.get_attr_smem_layout(ctx).map(|attr| attr.0.clone()) else { + return verify_err!(op.loc(), "cute.copy_g2s must have a shared-memory layout"); + }; + let Some(leading_dim_divisor) = self.get_attr_leading_dim_divisor(ctx).map(|attr| attr.0) + else { + return verify_err!( + op.loc(), + "cute.copy_g2s must have a leading-dimension divisor" + ); + }; + if leading_dim_divisor == 0 { + return verify_err!( + op.loc(), + "cute.copy_g2s leading-dimension divisor must be positive" + ); + } + let Some(elem): Option = self + .get_attr_copy_g2s_elem(ctx) + .map(|attr| attr.clone().into()) + else { + return verify_err!(op.loc(), "cute.copy_g2s must have an element type"); + }; + let Some(elem_bytes) = cooperative_elem_bytes(ctx, elem) else { + return verify_err!(op.loc(), "unsupported cute.copy_g2s element type"); + }; + let Some(leading_dim_alignment) = leading_dim_divisor.checked_mul(elem_bytes as u64) else { + return verify_err!( + op.loc(), + "cute.copy_g2s leading-dimension alignment overflows u64" + ); + }; + if leading_dim_alignment % u64::from(atom_bytes) != 0 { + return verify_err!( + op.loc(), + "cute.copy_g2s leading-dimension promise is weaker than its copy atom" + ); + } + // One shared pure validator feeds both importer recognition and this + // op. M3 can consume the returned TV map and capacity without + // re-deriving any layout fact. + let _plan = validate_cooperative_copy_plan( + atom_bytes, + &thread_layout, + &value_layout, + &tile_layout, + &smem_layout, + elem_bytes, + ) + .map_err(|error| { + pliron::input_error!(op.loc(), "invalid cooperative copy plan: {error}") + })?; + + for (index, role) in [(0usize, "global source"), (6usize, "shared destination")] { + let ty = op.get_operand(index).get_type(ctx); + let ty_ref = ty.deref(ctx); + let Some(pointer) = ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "{role} must be a MIR pointer"); + }; + if pointer.pointee != elem { + return verify_err!(op.loc(), "{role} pointee must match elem"); + } + if index == 0 + && ![address_space::GENERIC, address_space::GLOBAL].contains(&pointer.address_space) + { + return verify_err!(op.loc(), "global source has the wrong address space"); + } + if index == 6 + && (!pointer.is_mutable + || ![address_space::GENERIC, address_space::SHARED] + .contains(&pointer.address_space)) + { + return verify_err!( + op.loc(), + "shared destination must be mutable shared/generic memory" + ); + } + } + + for (index, role) in [ + (1usize, "row count"), + (2usize, "column count"), + (3usize, "leading dimension"), + (4usize, "tile row"), + (5usize, "tile column"), + (7usize, "shared-memory capacity"), + ] { + let ty = op.get_operand(index).get_type(ctx); + if ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 64 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "{role} must be an unsigned 64-bit integer"); + } + } + let tidx_ty = op.get_operand(8).get_type(ctx); + if tidx_ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 32 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "thread index must be an unsigned 32-bit integer"); + } + Ok(()) + } +} + +/// CuTe tile copy: copy the tile described by `layout` (element type `elem`) +/// from the memory at `src` to the memory at `dst`. +/// +/// # Operands +/// +/// ```text +/// | Name | Type | Description | +/// |-------|------------|------------------------------------| +/// | `src` | MirPtrType | Source pointer (pointee == elem) | +/// | `dst` | MirPtrType | Destination pointer (pointee == elem) | +/// ``` +/// +/// # Attributes +/// +/// ```text +/// | Name | Type | Description | +/// |----------|----------------|-----------------------------------| +/// | `layout` | CuteLayoutAttr | Tile layout, e.g. (4):(1) | +/// | `elem` | TypeAttr | Element type | +/// ``` +/// +/// # Contract +/// +/// Both pointers are naturally aligned for the whole tile (like CUDA's +/// `float4`): for an identity-map layout of `n` elements of size `s`, the +/// alignment is `n * s` bytes. The device API documents this; backend +/// lowering relies on it to emit a single vectorized access. +#[pliron_op( + name = "cute.copy", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<0>], + attributes = (layout: CuteLayoutAttr, elem: TypeAttr) +)] +pub struct CuteCopyOp; + +impl CuteCopyOp { + /// Create a `cute.copy`, computing nothing: the op carries the layout + /// structurally for the selected backend continuation. + pub fn new( + ctx: &mut Context, + src: Value, + dst: Value, + layout: Layout, + elem: TypeHandle, + ) -> Self { + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![src, dst], + vec![], + 0, + ); + let copy = CuteCopyOp { op }; + copy.set_attr_layout(ctx, CuteLayoutAttr(layout)); + copy.set_attr_elem(ctx, TypeAttr::new(elem)); + copy + } + + /// Wrap an existing `cute.copy` operation (caller must have matched the + /// OpId; used by backend continuations). + pub fn wrap(op: Ptr) -> Self { + CuteCopyOp { op } + } + + /// The tile layout (cloned out of the attribute). + pub fn layout(&self, ctx: &Context) -> Option { + self.get_attr_layout(ctx).map(|a| a.0.clone()) + } + + /// The element type. + pub fn elem(&self, ctx: &Context) -> Option { + self.get_attr_elem(ctx).map(|a| a.clone().into()) + } +} + +impl Verify for CuteCopyOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.copy needs exactly 2 operands and 0 results, got {} and {}", + op.get_num_operands(), + op.get_num_results() + ); + } + let elem: TypeHandle = match self.get_attr_elem(ctx) { + Some(a) => a.clone().into(), + None => return verify_err!(op.loc(), "cute.copy must have an elem attribute"), + }; + if self.get_attr_layout(ctx).is_none() { + return verify_err!(op.loc(), "cute.copy must have a layout attribute"); + } + if self + .get_attr_layout(ctx) + .is_none_or(|attr| attr.0.checked_size().is_none()) + { + return verify_err!( + op.loc(), + "cute.copy layout needs positive extents and a non-overflowing size" + ); + } + for (i, name) in [(0usize, "src"), (1usize, "dst")] { + let ty = op.get_operand(i).get_type(ctx); + let ty_ref = ty.deref(ctx); + let Some(ptr) = ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.copy {} operand must be a mir.ptr", name); + }; + if ptr.pointee != elem { + return verify_err!( + op.loc(), + "cute.copy {} operand pointee must match the elem attribute", + name + ); + } + if i == 1 && !ptr.is_mutable { + return verify_err!(op.loc(), "cute.copy destination pointer must be mutable"); + } + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/tma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/tma.rs new file mode 100644 index 0000000000..faa0bd9712 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/copy/tma.rs @@ -0,0 +1,276 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Descriptor-backed raw transport views and their direct copy contracts. + +use cute_layout::ComposedLayout; +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::{ + builtin::{ + attributes::TypeAttr, + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + value::Value, + verify_err, +}; +use pliron_derive::pliron_op; + +use super::cooperative_elem_bytes; +use crate::attributes::CuteComposedLayoutAttr; + +/// Frozen zero-result ABI for one hardware (TMA) tile copy into CTA-local +/// shared memory. +/// +/// ```text +/// 0 smem base (T) 2 tensor map 4 tile column index +/// 1 mbarrier 3 tile row index +/// ``` +/// +/// The shared layout stays in a typed attribute so the compiler can prove +/// TMA encodability at build time; the descriptor itself carries the same +/// layout in its Rust TYPE, so pairing errors never reach this op. +/// Barrier lifecycle (init, arrive_expect_tx, wait) stays caller-owned. +/// Remote or multicast cluster-shared copies require a distinct operation; +/// accepting a generic destination here would erase the proof needed to +/// select PTX's CTA-local TMA form. +#[pliron_op( + name = "cute.copy_tma_2d", + format, + interfaces = [NOpdsInterface<5>, NResultsInterface<0>], + attributes = ( + tma_smem_layout: CuteComposedLayoutAttr, + tma_elem: TypeAttr + ) +)] +pub struct CuteTmaLoad2dOp; + +impl CuteTmaLoad2dOp { + /// Build a detached CTA-local TMA load retaining its descriptor, barrier, and tile coordinates. + pub fn new( + ctx: &mut Context, + operands: [Value; 5], + smem_layout: ComposedLayout, + elem: TypeHandle, + ) -> Self { + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + operands.to_vec(), + vec![], + 0, + ); + let load = Self { op }; + load.set_attr_tma_smem_layout(ctx, CuteComposedLayoutAttr(smem_layout)); + load.set_attr_tma_elem(ctx, TypeAttr::new(elem)); + load + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTmaLoad2dOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + if op.get_num_operands() != 5 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.copy_tma_2d expects exactly 5 operands and 0 results" + ); + } + let Some(smem_layout) = self + .get_attr_tma_smem_layout(ctx) + .map(|attr| attr.0.clone()) + else { + return verify_err!( + op.loc(), + "cute.copy_tma_2d must have a shared-memory layout" + ); + }; + let Some(elem): Option = + self.get_attr_tma_elem(ctx).map(|attr| attr.clone().into()) + else { + return verify_err!(op.loc(), "cute.copy_tma_2d must have an element type"); + }; + let Some(elem_bytes) = cooperative_elem_bytes(ctx, elem) else { + return verify_err!(op.loc(), "unsupported cute.copy_tma_2d element type"); + }; + if let Err(error) = cute_layout::validate_tma_encodable(&smem_layout, elem_bytes) { + return verify_err!( + op.loc(), + "cute.copy_tma_2d layout is not encodable: {error}" + ); + } + let smem_ty = op.get_operand(0).get_type(ctx); + let smem_ref = smem_ty.deref(ctx); + let Some(pointer) = smem_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.copy_tma_2d destination must be a MIR pointer" + ); + }; + if pointer.pointee != elem + || !pointer.is_mutable + || pointer.address_space != address_space::SHARED + { + return verify_err!( + op.loc(), + "cute.copy_tma_2d destination must be mutable CTA-shared elem memory" + ); + } + for index in [1usize, 2] { + let ty = op.get_operand(index).get_type(ctx); + if ty.deref(ctx).downcast_ref::().is_none() { + return verify_err!( + op.loc(), + "cute.copy_tma_2d operand {index} must be a pointer" + ); + } + } + for (index, what) in [(3usize, "tile row"), (4usize, "tile column")] { + let ty = op.get_operand(index).get_type(ctx); + if ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 64 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "{what} must be an unsigned 64-bit integer"); + } + } + Ok(()) + } +} + +/// Frozen zero-result ABI for one hardware (TMA) tile copy from CTA-local +/// shared memory into global memory. +/// +/// ```text +/// 0 smem source (T) 2 tile row index +/// 1 tensor map 3 tile column index +/// ``` +/// +/// The shared layout and element type remain typed attributes. The descriptor +/// carries the same witnesses at the Rust boundary, while this operation +/// preserves them for the selected backend continuation. +#[pliron_op( + name = "cute.copy_tma_s2g_2d", + format, + interfaces = [NOpdsInterface<4>, NResultsInterface<0>], + attributes = ( + tma_store_smem_layout: CuteComposedLayoutAttr, + tma_store_elem: TypeAttr + ) +)] +pub struct CuteTmaStore2dOp; + +impl CuteTmaStore2dOp { + /// Build a detached CTA-local TMA store retaining its descriptor and tile coordinates. + pub fn new( + ctx: &mut Context, + operands: [Value; 4], + smem_layout: ComposedLayout, + elem: TypeHandle, + ) -> Self { + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + operands.to_vec(), + vec![], + 0, + ); + let store = Self { op }; + store.set_attr_tma_store_smem_layout(ctx, CuteComposedLayoutAttr(smem_layout)); + store.set_attr_tma_store_elem(ctx, TypeAttr::new(elem)); + store + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTmaStore2dOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + if op.get_num_operands() != 4 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d expects exactly 4 operands and 0 results" + ); + } + let Some(smem_layout) = self + .get_attr_tma_store_smem_layout(ctx) + .map(|attr| attr.0.clone()) + else { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d must have a shared-memory layout" + ); + }; + let Some(elem): Option = self + .get_attr_tma_store_elem(ctx) + .map(|attr| attr.clone().into()) + else { + return verify_err!(op.loc(), "cute.copy_tma_s2g_2d must have an element type"); + }; + let Some(elem_bytes) = cooperative_elem_bytes(ctx, elem) else { + return verify_err!(op.loc(), "unsupported cute.copy_tma_s2g_2d element type"); + }; + if let Err(error) = cute_layout::validate_tma_encodable(&smem_layout, elem_bytes) { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d layout is not encodable: {error}" + ); + } + let smem_ty = op.get_operand(0).get_type(ctx); + let smem_ref = smem_ty.deref(ctx); + let Some(pointer) = smem_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d source must be a MIR pointer" + ); + }; + if pointer.pointee != elem || pointer.address_space != address_space::SHARED { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d source must be CTA-shared elem memory" + ); + } + let tensor_map_ty = op.get_operand(1).get_type(ctx); + if tensor_map_ty + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!( + op.loc(), + "cute.copy_tma_s2g_2d tensor map must be a pointer" + ); + } + for (index, what) in [(2usize, "tile row"), (3usize, "tile column")] { + let ty = op.get_operand(index).get_type(ctx); + if ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 64 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "{what} must be an unsigned 64-bit integer"); + } + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/divisibility.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/divisibility.rs new file mode 100644 index 0000000000..b4ec35df1c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/divisibility.rs @@ -0,0 +1,113 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SSA identity operations preserving explicit integer divisibility promises. + +use pliron::{ + builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface, OneOpdInterface, OneResultInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + result::Error, + r#type::Typed, + value::Value, + verify_err, +}; +use pliron_derive::pliron_op; + +use crate::attributes::CuteDivisibilityAttr; + +/// Preserve the promise `value % divisor == 0` in SSA form. +/// +/// ```text +/// Rust cute IR lowering +/// ───── ─────── ──────── +/// y = assume::<16>(x) -> y = cute.assume_div x {16} -> y uses x +/// │ +/// └─ direct SSA consumers read the fact +/// ``` +/// +/// The operation does not change the runtime value. Its result exists so the +/// fact follows that exact value through ordinary SSA uses; a detached +/// zero-result annotation would require guessing which later value it meant. +#[pliron_op( + name = "cute.assume_div", + format, + interfaces = [ + NOpdsInterface<1>, + NResultsInterface<1>, + OneOpdInterface, + OneResultInterface + ], + attributes = (divisor: CuteDivisibilityAttr) +)] +pub struct CuteAssumeDivOp; + +impl CuteAssumeDivOp { + /// Build a detached identity operation carrying the caller's divisibility promise. + pub fn new(ctx: &mut Context, value: Value, divisor: u64) -> Self { + let result_ty = value.get_type(ctx); + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![value], + vec![], + 0, + ); + let assume = Self { op }; + assume.set_attr_divisor(ctx, CuteDivisibilityAttr(divisor)); + assume + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the promised positive integer divisor, if the attribute is present. + pub fn divisor(&self, ctx: &Context) -> Option { + self.get_attr_divisor(ctx).map(|attr| attr.0) + } +} + +impl Verify for CuteAssumeDivOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.assume_div needs exactly 1 operand and 1 result, got {} and {}", + op.get_num_operands(), + op.get_num_results() + ); + } + let Some(divisor) = self.divisor(ctx) else { + return verify_err!(op.loc(), "cute.assume_div must have a divisor attribute"); + }; + if divisor == 0 { + return verify_err!( + op.loc(), + "cute.assume_div divisor must be greater than zero" + ); + } + + let input_ty = op.get_operand(0).get_type(ctx); + if input_ty.deref(ctx).downcast_ref::().is_none() { + return verify_err!(op.loc(), "cute.assume_div input must be an integer"); + } + if op.get_result(0).get_type(ctx) != input_ty { + return verify_err!( + op.loc(), + "cute.assume_div result must have the same type as its input" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/mod.rs new file mode 100644 index 0000000000..46e1754b95 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/mod.rs @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared placement and asynchronous output for the SM120 GEMM epilogue. +//! +//! Runtime values stay ordinary MIR values. Static attributes explain the +//! shared tile and the store protocol: +//! +//! ```text +//! 64xf32 accumulator +//! │ +//! ▼ +//! epilogue_store_fragment ──► two 128x64 shared f16 halves +//! │ │ +//! ▼ ▼ +//! proxy fence + counted sync two tma_store_2d +//! │ +//! ▼ +//! one committed group +//! ``` +//! +//! The TMA store queue is hardware state owned by the issuing thread. Its +//! acquire, commit, and tail operations therefore have no fake SSA handle or +//! loop-carried counter. Their order in the CFG is the state. + +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::builtin::{ + attributes::TypeAttr, + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteCountedCtaBarrierAttr, CuteEpilogueHalfAttr, CuteEpilogueSyncPhaseAttr, + CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTmaStorePipelineAttr, +}; +use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; +use crate::types::{CuteEpilogueTileType, CuteTmaViewType, MmaCarrierShape, mma_carrier_shape}; + +/// Check that the SSA value has an unsigned 32-bit integer type. +fn is_u32(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.is_unsigned()) +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Downcast a type handle to the static shared output-tile description. +fn epilogue_tile_of_type(ctx: &Context, ty: TypeHandle) -> Option { + ty.deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Read a TMA transport view from an SSA operand's type. +fn tma_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Check mutable shared storage with the required element type. +fn shared_pointer_matches(ctx: &Context, value: Value, storage: TypeHandle) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|pointer| { + pointer.pointee == storage + && pointer.is_mutable + && pointer.address_space == address_space::SHARED + }) +} + +/// Require a valid epilogue tile attribute before checking carrier operands. +fn checked_tile( + tile: Option, + ctx: &Context, +) -> Result { + let tile = tile.ok_or_else(|| "must carry a cute.epilogue_tile TypeAttr".to_owned())?; + tile.verify(ctx) + .map_err(|error| format!("carries an invalid epilogue tile: {error}"))?; + Ok(tile) +} + +/// Build the supported 128-by-64 swizzled FP16 half-tile placement. +fn current_half_layout() -> ComposedLayout { + let outer: Layout = "(128,64):(64,1)" + .parse() + .expect("fixed epilogue half layout is valid"); + ComposedLayout::new(Swizzle::new(3, 3, 3), 0, outer, OffsetUnit::Elements) + .expect("fixed epilogue half swizzle is valid") +} + +mod store; +mod tile; + +pub use store::{ + CuteTmaStore2dSemanticOp, CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, +}; +pub use tile::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, CuteEpilogueSyncOp, + CuteEpilogueWarpSliceOp, +}; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/store.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/store.rs new file mode 100644 index 0000000000..6f79e42e5e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/store.rs @@ -0,0 +1,292 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Asynchronous TMA output stores and their hardware queue lifecycle. + +use super::*; + +/// Require a positive store-ring stage count on a zero-result, zero-operand operation. +fn check_store_pipeline( + pipeline: Option, + ctx: &Context, +) -> Result<(), String> { + let pipeline = pipeline.ok_or_else(|| "must carry a store-pipeline config".to_owned())?; + pipeline + .verify(ctx) + .map_err(|error| format!("has an invalid store-pipeline config: {error}"))?; + if pipeline.stages != 1 { + return Err("v0 supports the one-buffer store pipeline".to_owned()); + } + Ok(()) +} + +/// Wait until the shared result buffer may be overwritten. +#[pliron_op( + name = "cute.tma_store_acquire", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], + attributes = (tma_store_acquire_config: CuteTmaStorePipelineAttr) +)] +pub struct CuteTmaStoreAcquireOp; + +impl CuteTmaStoreAcquireOp { + /// Build a detached wait operation using the output ring's maximum pending group count. + pub fn new(ctx: &mut Context, pipeline: CuteTmaStorePipelineAttr) -> Self { + let operation = Self { + op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0), + }; + operation.set_attr_tma_store_acquire_config(ctx, pipeline); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the static number of reusable output buffers, if present. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Option { + self.get_attr_tma_store_acquire_config(ctx) + .map(|pipeline| *pipeline) + } +} + +impl Verify for CuteTmaStoreAcquireOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 0 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tma_store_acquire needs 0 operands and 0 results" + ); + } + if let Err(message) = check_store_pipeline(self.pipeline(ctx), ctx) { + return verify_err!(op.loc(), "cute.tma_store_acquire {message}"); + } + Ok(()) + } +} + +/// Close this thread's pending TMA stores as one hardware copy group. +#[pliron_op( + name = "cute.tma_store_commit", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], + attributes = (tma_store_commit_config: CuteTmaStorePipelineAttr) +)] +pub struct CuteTmaStoreCommitOp; + +impl CuteTmaStoreCommitOp { + /// Build a detached operation committing the issuer's pending TMA stores as one group. + pub fn new(ctx: &mut Context, pipeline: CuteTmaStorePipelineAttr) -> Self { + let operation = Self { + op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0), + }; + operation.set_attr_tma_store_commit_config(ctx, pipeline); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the static number of reusable output buffers, if present. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Option { + self.get_attr_tma_store_commit_config(ctx) + .map(|pipeline| *pipeline) + } +} + +impl Verify for CuteTmaStoreCommitOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 0 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tma_store_commit needs 0 operands and 0 results" + ); + } + if let Err(message) = check_store_pipeline(self.pipeline(ctx), ctx) { + return verify_err!(op.loc(), "cute.tma_store_commit {message}"); + } + Ok(()) + } +} + +/// Before exit, wait until no TMA store still reads shared memory. +#[pliron_op( + name = "cute.tma_store_tail", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<0>], + attributes = (tma_store_tail_config: CuteTmaStorePipelineAttr) +)] +pub struct CuteTmaStoreTailOp; + +impl CuteTmaStoreTailOp { + /// Build a detached operation waiting for all committed stores to stop reading shared storage. + pub fn new(ctx: &mut Context, pipeline: CuteTmaStorePipelineAttr) -> Self { + let operation = Self { + op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], vec![], vec![], 0), + }; + operation.set_attr_tma_store_tail_config(ctx, pipeline); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the static number of reusable output buffers, if present. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Option { + self.get_attr_tma_store_tail_config(ctx) + .map(|pipeline| *pipeline) + } +} + +impl Verify for CuteTmaStoreTailOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 0 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tma_store_tail needs 0 operands and 0 results" + ); + } + if let Err(message) = check_store_pipeline(self.pipeline(ctx), ctx) { + return verify_err!(op.loc(), "cute.tma_store_tail {message}"); + } + Ok(()) + } +} + +/// Copy one typed shared tile to a descriptor-backed global destination. +/// +/// Both operands are short-lived direct TMA-view ghosts. The source view owns +/// the ordinary shared pointer/capacity pair; the destination view owns the +/// immutable descriptor pointer. +#[pliron_op( + name = "cute.tma_store_2d", + format, + interfaces = [NOpdsInterface<4>, NResultsInterface<0>] +)] +pub struct CuteTmaStore2dSemanticOp; + +impl CuteTmaStore2dSemanticOp { + /// Build a detached copy from a typed shared half to its descriptor-backed global tile. + pub fn new( + ctx: &mut Context, + source: Value, + destination: Value, + tile_row: Value, + tile_column: Value, + ) -> Self { + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![source, destination, tile_row, tile_column], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the source view operand of the asynchronous transfer. + #[must_use] + pub fn source(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the destination view operand of the asynchronous transfer. + #[must_use] + pub fn destination(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the runtime row index in units of complete TMA tiles. + #[must_use] + pub fn tile_row(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + /// Return the runtime column index in units of complete TMA tiles. + #[must_use] + pub fn tile_column(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } +} + +impl Verify for CuteTmaStore2dSemanticOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 4 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "cute.tma_store_2d needs 4 operands and 0 results"); + } + let Some(source) = tma_view_of(ctx, self.source(ctx)) else { + return verify_err!(op.loc(), "cute.tma_store_2d source must be a TMA view"); + }; + let Some(destination) = tma_view_of(ctx, self.destination(ctx)) else { + return verify_err!(op.loc(), "cute.tma_store_2d destination must be a TMA view"); + }; + source.verify(ctx)?; + destination.verify(ctx)?; + let source_tensor = source.tensor_view(ctx).expect("verified source tensor"); + let destination_tensor = destination + .tensor_view(ctx) + .expect("verified destination tensor"); + if source_tensor.space != CuteTensorAddressSpaceAttr::Smem + || source_tensor.access != CuteTensorAccessAttr::ReadWrite + || destination_tensor.space != CuteTensorAddressSpaceAttr::Gmem + || destination_tensor.access != CuteTensorAccessAttr::ReadWrite + { + return verify_err!( + op.loc(), + "cute.tma_store_2d direction must be writable Smem to writable Gmem" + ); + } + if source_tensor.storage != destination_tensor.storage + || source.smem_layout != destination.smem_layout + { + return verify_err!( + op.loc(), + "cute.tma_store_2d source and destination need the same carrier and tile layout" + ); + } + if source_tensor + .storage + .deref(ctx) + .downcast_ref::() + .is_none() + || source.smem_layout.0 != current_half_layout() + || source.tile_elements() != Some(8192) + || source.tile_bytes(ctx) != Some(16 * 1024) + || source_tensor.alignment.0 != 1024 + || destination_tensor.alignment.0 != 2 + { + return verify_err!( + op.loc(), + "cute.tma_store_2d v0 needs one 1024-byte-aligned B128 128x64 f16 half" + ); + } + if !is_u64(ctx, self.tile_row(ctx)) || !is_u64(ctx, self.tile_column(ctx)) { + return verify_err!( + op.loc(), + "cute.tma_store_2d row and column must be unsigned 64-bit integers" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tests.rs new file mode 100644 index 0000000000..50e00af82e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tests.rs @@ -0,0 +1,319 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use crate::attributes::{CuteEpiloguePlanAttr, CuteTiledMmaPlanAttr}; +use crate::ops::gemm_tma::{CuteTmaGmemViewOp, CuteTmaSmemViewOp}; +use crate::types::CuteEpilogueTileType; +use dialect_mir::ops::MirUndefOp; +use dialect_mir::types::{MirArrayType, MirStructType}; +use pliron::builtin::types::FP32Type; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Build the canonical shared placement consumed by the output-tile fixture. +fn data_layout() -> ComposedLayout { + let outer: Layout = "(128,32):(32,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, outer, OffsetUnit::Elements).unwrap() +} + +/// Ordinary runtime carriers and static types for a two-half epilogue test. +struct Fixture { + /// Shared storage pointer used by this fixture. + base: Value, + /// Warp index used by the output-tile fixture. + warp: Value, + /// Lane index within one 32-thread warp. + lane: Value, + /// FP32 accumulator carrier used by the compute or output operation. + accumulator: Value, + /// Storage capacity expressed as an unsigned element count. + capacity: Value, + /// Runtime tile-row coordinate supplied to the copy. + row: Value, + /// Runtime tile-column coordinate supplied to the copy. + column: Value, + /// Immutable global tensor-map pointer used by the fixture. + descriptor: Value, + /// FP16 shared-storage type. + f16_ty: TypeHandle, + /// Static type describing the supported shared output tile. + tile_ty: TypeHandle, +} + +/// Build shared output, accumulator, coordinates, and static tile metadata for output tests. +fn setup() -> (Context, Fixture) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(&ctx).into(); + let base_ty: TypeHandle = MirPtrType::get(&mut ctx, f16_ty, true, address_space::SHARED).into(); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let marker: TypeHandle = + MirStructType::get(&mut ctx, "PhantomData".into(), vec![], vec![]).into(); + let registers: TypeHandle = MirArrayType::get(&mut ctx, f32_ty, 64).into(); + let accumulator_ty: TypeHandle = MirStructType::get( + &mut ctx, + "Mxf4AccumulatorTile2x8".into(), + vec![], + vec![registers, marker], + ) + .into(); + let mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(data_layout()); + let plan = CuteEpiloguePlanAttr::sm120_mxf4_128x128(mma); + let tile_ty: TypeHandle = CuteEpilogueTileType::get(&mut ctx, f16_ty, plan).into(); + + let base = undef(&mut ctx, base_ty); + let warp = undef(&mut ctx, u64_ty); + let lane = undef(&mut ctx, u32_ty); + let accumulator = undef(&mut ctx, accumulator_ty); + let capacity = undef(&mut ctx, u64_ty); + let row = undef(&mut ctx, u64_ty); + let column = undef(&mut ctx, u64_ty); + let descriptor = undef(&mut ctx, descriptor_ty); + ( + ctx, + Fixture { + base, + warp, + lane, + accumulator, + capacity, + row, + column, + descriptor, + f16_ty, + tile_ty, + }, + ) +} + +/// Output overlay, warp selection, synchronization, and half selection preserve ordinary runtime carriers. +#[test] +fn visual_epilogue_chain_keeps_only_ordinary_runtime_carriers() { + let (mut ctx, fixture) = setup(); + let overlay = CuteEpilogueSmemOverlayOp::new(&mut ctx, fixture.base, fixture.tile_ty); + assert!(overlay.verify(&ctx).is_ok()); + let overlay_base = overlay.base(&ctx); + let slice = CuteEpilogueWarpSliceOp::new( + &mut ctx, + overlay_base, + fixture.warp, + fixture.lane, + fixture.tile_ty, + ); + assert!(slice.verify(&ctx).is_ok()); + let slice_base = slice.base(&ctx); + let slice_warp = slice.warp_id(&ctx); + let slice_lane = slice.lane(&ctx); + let store = CuteEpilogueStoreFragmentOp::new( + &mut ctx, + slice_base, + slice_warp, + slice_lane, + fixture.accumulator, + fixture.tile_ty, + ); + assert!(store.verify(&ctx).is_ok()); + + let barrier = CuteCountedCtaBarrierAttr::new(2, 0, 8, 9, 32); + let reusable = CuteEpilogueSyncOp::new( + &mut ctx, + overlay_base, + fixture.tile_ty, + barrier, + CuteEpilogueSyncPhaseAttr::Reusable, + ); + let ready = CuteEpilogueSyncOp::new( + &mut ctx, + overlay_base, + fixture.tile_ty, + barrier, + CuteEpilogueSyncPhaseAttr::ReadyForTma, + ); + assert!(reusable.verify(&ctx).is_ok()); + assert!(ready.verify(&ctx).is_ok()); + assert_eq!(barrier.participant_threads(), Some(256)); + assert_eq!(barrier.excluded_warps(), Some(1)); + + let left = CuteEpilogueHalfOp::new( + &mut ctx, + overlay_base, + fixture.tile_ty, + CuteEpilogueHalfAttr(0), + ); + let right = CuteEpilogueHalfOp::new( + &mut ctx, + overlay_base, + fixture.tile_ty, + CuteEpilogueHalfAttr(1), + ); + assert!(left.verify(&ctx).is_ok()); + assert!(right.verify(&ctx).is_ok()); + assert_eq!( + left.capacity(&ctx).get_type(&ctx), + fixture.capacity.get_type(&ctx) + ); +} + +/// Acquire/commit/tail represent issuer-local hardware state with no fabricated SSA queue object. +#[test] +fn store_pipeline_has_no_runtime_handle() { + let (mut ctx, _) = setup(); + let pipeline = CuteTmaStorePipelineAttr::new(1); + let acquire = CuteTmaStoreAcquireOp::new(&mut ctx, pipeline); + let commit = CuteTmaStoreCommitOp::new(&mut ctx, pipeline); + let tail = CuteTmaStoreTailOp::new(&mut ctx, pipeline); + assert!(acquire.verify(&ctx).is_ok()); + assert!(commit.verify(&ctx).is_ok()); + assert!(tail.verify(&ctx).is_ok()); + for operation in [ + acquire.get_operation(), + commit.get_operation(), + tail.get_operation(), + ] { + assert_eq!(operation.deref(&ctx).get_num_operands(), 0); + assert_eq!(operation.deref(&ctx).get_num_results(), 0); + } + + assert!( + CuteTmaStoreAcquireOp::new(&mut ctx, CuteTmaStorePipelineAttr::new(2)) + .verify(&ctx) + .is_err() + ); +} + +/// Shared output halves and writable global transport views must agree on storage and layout. +#[test] +fn semantic_s2g_store_requires_matching_writable_views() { + let (mut ctx, fixture) = setup(); + let source = CuteTmaSmemViewOp::new( + &mut ctx, + fixture.base, + fixture.capacity, + fixture.f16_ty, + current_half_layout(), + 1024, + ); + let destination = CuteTmaGmemViewOp::new_destination( + &mut ctx, + fixture.descriptor, + fixture.f16_ty, + current_half_layout(), + ); + assert!(source.verify(&ctx).is_ok()); + assert!(destination.verify(&ctx).is_ok()); + let source_view = source.get_operation().deref(&ctx).get_result(0); + let destination_view = destination.get_operation().deref(&ctx).get_result(0); + let store = CuteTmaStore2dSemanticOp::new( + &mut ctx, + source_view, + destination_view, + fixture.row, + fixture.column, + ); + assert!(store.verify(&ctx).is_ok()); + + let read_only = CuteTmaGmemViewOp::new( + &mut ctx, + fixture.descriptor, + fixture.f16_ty, + current_half_layout(), + ); + let read_only_view = read_only.get_operation().deref(&ctx).get_result(0); + let wrong_direction = CuteTmaStore2dSemanticOp::new( + &mut ctx, + source_view, + read_only_view, + fixture.row, + fixture.column, + ); + assert!(wrong_direction.verify(&ctx).is_err()); +} + +/// Epilogue checks reject mismatched barrier participation, accumulator shape, and out-of-range half selection. +#[test] +fn verifiers_reject_wrong_barrier_fragment_and_half() { + let (mut ctx, fixture) = setup(); + let wrong_barrier = CuteEpilogueSyncOp::new( + &mut ctx, + fixture.base, + fixture.tile_ty, + CuteCountedCtaBarrierAttr::new(2, 0, 9, 9, 32), + CuteEpilogueSyncPhaseAttr::Reusable, + ); + assert!(wrong_barrier.verify(&ctx).is_err()); + + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let short_accumulator_ty: TypeHandle = MirArrayType::get(&mut ctx, f32_ty, 63).into(); + let short_accumulator = undef(&mut ctx, short_accumulator_ty); + let wrong_fragment = CuteEpilogueStoreFragmentOp::new( + &mut ctx, + fixture.base, + fixture.warp, + fixture.lane, + short_accumulator, + fixture.tile_ty, + ); + assert!(wrong_fragment.verify(&ctx).is_err()); + + let wrong_half = CuteEpilogueHalfOp::new( + &mut ctx, + fixture.base, + fixture.tile_ty, + CuteEpilogueHalfAttr(2), + ); + assert!(wrong_half.verify(&ctx).is_err()); +} + +/// Unsupported output-plan geometry or storage interpretation is rejected independently of backend choice. +#[test] +fn static_plan_and_tile_fail_closed_before_backend_selection() { + let (mut ctx, fixture) = setup(); + let mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(data_layout()); + let valid = CuteEpiloguePlanAttr::sm120_mxf4_128x128(mma.clone()); + assert!(valid.verify(&ctx).is_ok()); + + let mut one_half = valid.clone(); + one_half.halves = 1; + assert!(one_half.verify(&ctx).is_err()); + + let alternate_data = + ComposedLayout::from_layout("(128,32):(32,1)".parse().unwrap(), OffsetUnit::Elements); + let alternate_mma = CuteTiledMmaPlanAttr::mxf4_128x128x128(alternate_data); + assert!( + CuteEpiloguePlanAttr::sm120_mxf4_128x128(alternate_mma) + .verify(&ctx) + .is_err() + ); + + let wrong_storage: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let wrong_tile = CuteEpilogueTileType::get( + &mut ctx, + wrong_storage, + CuteEpiloguePlanAttr::sm120_mxf4_128x128(mma), + ); + assert!(wrong_tile.deref(&ctx).verify(&ctx).is_err()); + assert!(fixture.tile_ty.deref(&ctx).verify(&ctx).is_ok()); + + assert!(CuteTmaStorePipelineAttr::new(0).verify(&ctx).is_err()); + assert!( + CuteCountedCtaBarrierAttr::new(16, 0, 8, 9, 32) + .verify(&ctx) + .is_err() + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tile.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tile.rs new file mode 100644 index 0000000000..9235928796 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/epilogue/tile.rs @@ -0,0 +1,547 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared output tiles, warp placement, fragment stores, and synchronization. + +use super::*; + +/// Attach result-tile meaning to one shared pointer without changing it. +#[pliron_op( + name = "cute.epilogue_smem_overlay", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], + attributes = (epilogue_overlay_tile: TypeAttr) +)] +pub struct CuteEpilogueSmemOverlayOp; + +impl CuteEpilogueSmemOverlayOp { + /// Build a detached overlay retaining the shared pointer and attaching the output-tile type. + pub fn new(ctx: &mut Context, base: Value, tile: TypeHandle) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![base.get_type(ctx)], + vec![base], + vec![], + 0, + ), + }; + operation.set_attr_epilogue_overlay_tile(ctx, TypeAttr::new(tile)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the original shared pointer supplied to the semantic overlay. + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Read the static shared output-tile type attached to the operation. + #[must_use] + pub fn tile(&self, ctx: &Context) -> Option { + let attr = self.get_attr_epilogue_overlay_tile(ctx)?; + epilogue_tile_of_type(ctx, attr.get_type(ctx)) + } +} + +impl Verify for CuteEpilogueSmemOverlayOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.epilogue_smem_overlay needs 1 operand and 1 result" + ); + } + let tile = match checked_tile(self.tile(ctx), ctx) { + Ok(tile) => tile, + Err(message) => return verify_err!(op.loc(), "cute.epilogue_smem_overlay {message}"), + }; + if !shared_pointer_matches(ctx, self.input_base(ctx), tile.storage) { + return verify_err!( + op.loc(), + "cute.epilogue_smem_overlay base must be a mutable CTA-shared f16 pointer" + ); + } + if self.base(ctx).get_type(ctx) != self.input_base(ctx).get_type(ctx) { + return verify_err!( + op.loc(), + "cute.epilogue_smem_overlay must preserve the base pointer type" + ); + } + Ok(()) + } +} + +/// Keep the pointer, warp, and lane selected by `get_slice` visible. +/// +/// All three results are ordinary values and may be forwarded through the +/// Rust `Sm120EpilogueWarp128x128` aggregate. +#[pliron_op( + name = "cute.epilogue_warp_slice", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<3>], + attributes = (epilogue_warp_slice_tile: TypeAttr) +)] +pub struct CuteEpilogueWarpSliceOp; + +impl CuteEpilogueWarpSliceOp { + /// Build a detached output-slice operation forwarding the pointer, warp, and lane carriers. + pub fn new( + ctx: &mut Context, + base: Value, + warp_id: Value, + lane: Value, + tile: TypeHandle, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ + base.get_type(ctx), + warp_id.get_type(ctx), + lane.get_type(ctx), + ], + vec![base, warp_id, lane], + vec![], + 0, + ), + }; + operation.set_attr_epilogue_warp_slice_tile(ctx, TypeAttr::new(tile)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the original shared pointer supplied to the semantic overlay. + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the warp identifier supplied when selecting the output slice. + #[must_use] + pub fn input_warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the lane identifier supplied when selecting the output slice. + #[must_use] + pub fn input_lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the ordinary warp-index carrier for output placement. + #[must_use] + pub fn warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Return the lane-index carrier identifying register ownership within a warp. + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + /// Read the static shared output-tile type attached to the operation. + #[must_use] + pub fn tile(&self, ctx: &Context) -> Option { + let attr = self.get_attr_epilogue_warp_slice_tile(ctx)?; + epilogue_tile_of_type(ctx, attr.get_type(ctx)) + } +} + +impl Verify for CuteEpilogueWarpSliceOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 3 { + return verify_err!( + op.loc(), + "cute.epilogue_warp_slice needs 3 operands and 3 results" + ); + } + let tile = match checked_tile(self.tile(ctx), ctx) { + Ok(tile) => tile, + Err(message) => return verify_err!(op.loc(), "cute.epilogue_warp_slice {message}"), + }; + if !shared_pointer_matches(ctx, self.input_base(ctx), tile.storage) + || !is_u64(ctx, self.input_warp_id(ctx)) + || !is_u32(ctx, self.input_lane(ctx)) + { + return verify_err!( + op.loc(), + "cute.epilogue_warp_slice needs a shared f16 base, u64 warp ID, and u32 lane" + ); + } + if (0..3) + .any(|index| op.get_result(index).get_type(ctx) != op.get_operand(index).get_type(ctx)) + { + return verify_err!( + op.loc(), + "cute.epilogue_warp_slice must preserve all three carrier types" + ); + } + Ok(()) + } +} + +/// Convert and place one warp's complete 2x8 accumulator tile in shared memory. +#[pliron_op( + name = "cute.epilogue_store_fragment", + format, + interfaces = [NOpdsInterface<4>, NResultsInterface<0>], + attributes = (epilogue_store_fragment_tile: TypeAttr) +)] +pub struct CuteEpilogueStoreFragmentOp; + +impl CuteEpilogueStoreFragmentOp { + /// Build a detached operation placing one warp's complete accumulator into shared output storage. + pub fn new( + ctx: &mut Context, + base: Value, + warp_id: Value, + lane: Value, + accumulator: Value, + tile: TypeHandle, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![base, warp_id, lane, accumulator], + vec![], + 0, + ), + }; + operation.set_attr_epilogue_store_fragment_tile(ctx, TypeAttr::new(tile)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the ordinary warp-index carrier for output placement. + #[must_use] + pub fn warp_id(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the lane-index carrier identifying register ownership within a warp. + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + /// Return the ordinary FP32 accumulator carrier. + #[must_use] + pub fn accumulator(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } + + /// Read the static shared output-tile type attached to the operation. + #[must_use] + pub fn tile(&self, ctx: &Context) -> Option { + let attr = self.get_attr_epilogue_store_fragment_tile(ctx)?; + epilogue_tile_of_type(ctx, attr.get_type(ctx)) + } +} + +impl Verify for CuteEpilogueStoreFragmentOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 4 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.epilogue_store_fragment needs 4 operands and 0 results" + ); + } + let tile = match checked_tile(self.tile(ctx), ctx) { + Ok(tile) => tile, + Err(message) => { + return verify_err!(op.loc(), "cute.epilogue_store_fragment {message}"); + } + }; + if !shared_pointer_matches(ctx, self.base(ctx), tile.storage) + || !is_u64(ctx, self.warp_id(ctx)) + || !is_u32(ctx, self.lane(ctx)) + { + return verify_err!( + op.loc(), + "cute.epilogue_store_fragment needs a shared f16 base, u64 warp ID, and u32 lane" + ); + } + let expected = MmaCarrierShape::f32( + tile.plan + .tiled_mma + .accumulator_registers_per_lane() + .expect("verified epilogue plan has an accumulator size"), + ); + if mma_carrier_shape(ctx, self.accumulator(ctx).get_type(ctx)) != Some(expected) { + return verify_err!( + op.loc(), + "cute.epilogue_store_fragment accumulator must be an ordinary 64xf32 MIR carrier" + ); + } + Ok(()) + } +} + +/// Meet before shared reuse or after every writer has published its result. +#[pliron_op( + name = "cute.epilogue_sync", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<0>], + attributes = ( + epilogue_sync_tile: TypeAttr, + epilogue_sync_barrier: CuteCountedCtaBarrierAttr, + epilogue_sync_phase: CuteEpilogueSyncPhaseAttr + ) +)] +pub struct CuteEpilogueSyncOp; + +impl CuteEpilogueSyncOp { + /// Build an epilogue synchronization boundary. + /// + /// `ReadyForTma` semantically includes publication from the generic + /// shared-memory proxy before the counted hand-off. Backends must lower + /// that phase to both their proxy fence and barrier sequence. + pub fn new( + ctx: &mut Context, + base: Value, + tile: TypeHandle, + barrier: CuteCountedCtaBarrierAttr, + phase: CuteEpilogueSyncPhaseAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![base], + vec![], + 0, + ), + }; + operation.set_attr_epilogue_sync_tile(ctx, TypeAttr::new(tile)); + operation.set_attr_epilogue_sync_barrier(ctx, barrier); + operation.set_attr_epilogue_sync_phase(ctx, phase); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Read the static shared output-tile type attached to the operation. + #[must_use] + pub fn tile(&self, ctx: &Context) -> Option { + let attr = self.get_attr_epilogue_sync_tile(ctx)?; + epilogue_tile_of_type(ctx, attr.get_type(ctx)) + } + + /// Clone the counted CTA-barrier participation contract, if present. + #[must_use] + pub fn barrier(&self, ctx: &Context) -> Option { + self.get_attr_epilogue_sync_barrier(ctx) + .map(|barrier| *barrier) + } + + /// Return the reusable/ready hand-off phase encoded by the synchronization operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Option { + self.get_attr_epilogue_sync_phase(ctx).map(|phase| *phase) + } +} + +impl Verify for CuteEpilogueSyncOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "cute.epilogue_sync needs 1 operand and 0 results"); + } + let tile = match checked_tile(self.tile(ctx), ctx) { + Ok(tile) => tile, + Err(message) => return verify_err!(op.loc(), "cute.epilogue_sync {message}"), + }; + if !shared_pointer_matches(ctx, self.base(ctx), tile.storage) { + return verify_err!( + op.loc(), + "cute.epilogue_sync base must be a mutable CTA-shared f16 pointer" + ); + } + let Some(barrier) = self.barrier(ctx) else { + return verify_err!(op.loc(), "cute.epilogue_sync must carry a counted barrier"); + }; + barrier.verify(ctx)?; + let compute_warps = u32::try_from( + tile.plan + .tiled_mma + .compute_warps() + .expect("verified epilogue plan has a warp count"), + ) + .expect("v0 compute-warp count fits u32"); + if barrier.barrier_id != 2 + || barrier.first_warp != 0 + || barrier.warp_count != compute_warps + || barrier.cta_warps != compute_warps + 1 + || barrier.lanes_per_warp != 32 + || barrier.participant_threads() != Some(256) + || barrier.excluded_warps() != Some(1) + { + return verify_err!( + op.loc(), + "cute.epilogue_sync v0 needs barrier 2 for 8 compute warps, with one producer warp excluded" + ); + } + if self.phase(ctx).is_none() { + return verify_err!(op.loc(), "cute.epilogue_sync must name its protocol phase"); + } + Ok(()) + } +} + +/// Select one ordinary pointer/capacity pair from the full shared result tile. +#[pliron_op( + name = "cute.epilogue_half", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<2>], + attributes = ( + epilogue_half_tile: TypeAttr, + epilogue_half_index: CuteEpilogueHalfAttr + ) +)] +pub struct CuteEpilogueHalfOp; + +impl CuteEpilogueHalfOp { + /// Build a detached selector producing one half's shared pointer and element capacity. + pub fn new( + ctx: &mut Context, + base: Value, + tile: TypeHandle, + half: CuteEpilogueHalfAttr, + ) -> Self { + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![base.get_type(ctx), u64_ty], + vec![base], + vec![], + 0, + ), + }; + operation.set_attr_epilogue_half_tile(ctx, TypeAttr::new(tile)); + operation.set_attr_epilogue_half_index(ctx, half); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the shared base pointer of the complete output tile. + #[must_use] + pub fn full_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the shared pointer result selecting one output half. + #[must_use] + pub fn half_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the ordinary unsigned storage-capacity carrier, in elements. + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Read the static shared output-tile type attached to the operation. + #[must_use] + pub fn tile(&self, ctx: &Context) -> Option { + let attr = self.get_attr_epilogue_half_tile(ctx)?; + epilogue_tile_of_type(ctx, attr.get_type(ctx)) + } + + /// Return the static left/right output-half selection. + #[must_use] + pub fn half(&self, ctx: &Context) -> Option { + self.get_attr_epilogue_half_index(ctx).map(|half| *half) + } +} + +impl Verify for CuteEpilogueHalfOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 2 { + return verify_err!(op.loc(), "cute.epilogue_half needs 1 operand and 2 results"); + } + let tile = match checked_tile(self.tile(ctx), ctx) { + Ok(tile) => tile, + Err(message) => return verify_err!(op.loc(), "cute.epilogue_half {message}"), + }; + if !shared_pointer_matches(ctx, self.full_base(ctx), tile.storage) + || self.half_base(ctx).get_type(ctx) != self.full_base(ctx).get_type(ctx) + || !is_u64(ctx, self.capacity(ctx)) + { + return verify_err!( + op.loc(), + "cute.epilogue_half must return a shared f16 pointer and u64 capacity" + ); + } + let Some(half) = self.half(ctx) else { + return verify_err!(op.loc(), "cute.epilogue_half must name half 0 or 1"); + }; + half.verify(ctx)?; + if half.0 >= tile.plan.halves { + return verify_err!( + op.loc(), + "cute.epilogue_half index is outside the full tile" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/mod.rs new file mode 100644 index 0000000000..29a6793e3b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/mod.rs @@ -0,0 +1,510 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! High-level TMA transport used by the first GEMM slice. +//! +//! This layer keeps one simple story visible: +//! +//! ```text +//! descriptor ─► gmem carrier view ─┐ +//! ├─ tma_copy_2d(row, col, barrier) +//! shared base ─► smem carrier view ┘ +//! ``` +//! +//! TMA moves raw storage boxes. A packed-FP4 tile therefore travels as `u8`, +//! and a canonical scale tile travels as `u16`. These transport views are +//! deliberately `Plain` and `Generic`: E2M1, UE8M0, Mkl, and Nkl become +//! meaningful only when later shared-tensor and MMA operations interpret the +//! copied bytes. +//! +//! A backend continuation can lower each semantic copy without changing its +//! ordinary runtime ABI: +//! +//! ```text +//! tma_copy_2d(gmem, smem, row, col, barrier) +//! │ resolve the two direct view producers +//! ▼ +//! copy_tma_2d(smem.base, barrier, gmem.descriptor, row, col) +//! ``` + +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface, OneResultInterface}, + types::{FP32Type, IntegerType}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, + CuteTensorRoleAttr, +}; +use crate::layout::ComposedLayout; +use crate::types::{CuteTensorViewType, CuteTmaViewType}; + +/// Read a TMA transport view from an SSA operand's type. +fn tma_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Return the supported scalar carrier width in bytes, excluding decoded packed formats. +fn transport_element_bytes(ctx: &Context, element: TypeHandle) -> Option { + let element = element.deref(ctx); + if element.downcast_ref::().is_some() { + Some(2) + } else if element.downcast_ref::().is_some() { + Some(4) + } else { + element.downcast_ref::().and_then(|integer| { + match (integer.width(), integer.is_unsigned()) { + (8, true) => Some(1), + (16, true) => Some(2), + _ => None, + } + }) + } +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Build the raw-carrier tensor/TMA type with the requested address space and access. +fn transport_view_type( + ctx: &mut Context, + element: TypeHandle, + space: CuteTensorAddressSpaceAttr, + access: CuteTensorAccessAttr, + alignment_bytes: u64, + smem_layout: ComposedLayout, +) -> TypeHandle { + let tensor: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + element, + element, + space, + access, + alignment_bytes, + CuteTensorFormatAttr::Plain, + CuteTensorRoleAttr::Generic, + CuteTensorLayoutAttr::Tma2D, + ) + .into(); + CuteTmaViewType::get(ctx, tensor, smem_layout).into() +} + +/// Wrap a TMA descriptor as a global raw-carrier tensor view. +/// +/// The descriptor owns the real global pointer, dimensions, and strides. The +/// result type keeps the carrier element and the destination tile placement +/// visible without pretending the descriptor pointer is the data pointer. +/// The inner tensor's natural carrier alignment is therefore not a claim +/// about the hidden global data address; descriptor creation and the unsafe +/// copy contract own that separate requirement. +#[pliron_op( + name = "cute.tma_gmem_view", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteTmaGmemViewOp; + +impl CuteTmaGmemViewOp { + /// Build a descriptor-backed transport view with explicit read or write access. + fn new_with_access( + ctx: &mut Context, + descriptor: Value, + element: TypeHandle, + smem_layout: ComposedLayout, + access: CuteTensorAccessAttr, + ) -> Self { + let alignment = transport_element_bytes(ctx, element) + .expect("cute.tma_gmem_view builder needs a supported carrier element"); + let result = transport_view_type( + ctx, + element, + CuteTensorAddressSpaceAttr::Gmem, + access, + alignment, + smem_layout, + ); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![descriptor], + vec![], + 0, + ), + } + } + + /// Wrap a descriptor used as a global source. + pub fn new( + ctx: &mut Context, + descriptor: Value, + element: TypeHandle, + smem_layout: ComposedLayout, + ) -> Self { + Self::new_with_access( + ctx, + descriptor, + element, + smem_layout, + CuteTensorAccessAttr::ReadOnly, + ) + } + + /// Wrap a descriptor used as a global destination. + /// + /// The descriptor pointer is still immutable. `ReadWrite` describes the + /// data reached through the descriptor, not the descriptor bytes. + pub fn new_destination( + ctx: &mut Context, + descriptor: Value, + element: TypeHandle, + smem_layout: ComposedLayout, + ) -> Self { + Self::new_with_access( + ctx, + descriptor, + element, + smem_layout, + CuteTensorAccessAttr::ReadWrite, + ) + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the descriptor handle hidden behind the semantic view. + #[must_use] + pub fn descriptor(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } +} + +impl Verify for CuteTmaGmemViewOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tma_gmem_view needs 1 operand and 1 result"); + } + let descriptor_ty = self.descriptor(ctx).get_type(ctx); + let descriptor_ty = descriptor_ty.deref(ctx); + let Some(descriptor) = descriptor_ty.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tma_gmem_view descriptor must be a pointer"); + }; + if descriptor.is_mutable || descriptor.address_space != address_space::GENERIC { + return verify_err!( + op.loc(), + "cute.tma_gmem_view descriptor must be an immutable generic pointer" + ); + } + let Some(view) = tma_view_of(ctx, op.get_result(0)) else { + return verify_err!(op.loc(), "cute.tma_gmem_view result must be a TMA view"); + }; + view.verify(ctx)?; + let tensor = view.tensor_view(ctx).expect("verified TMA tensor view"); + let element_bytes = tensor.storage_bytes(ctx).expect("verified carrier width"); + if tensor.space != CuteTensorAddressSpaceAttr::Gmem + || !matches!( + tensor.access, + CuteTensorAccessAttr::ReadOnly | CuteTensorAccessAttr::ReadWrite + ) + || tensor.alignment.0 != element_bytes + { + return verify_err!( + op.loc(), + "cute.tma_gmem_view must be a global carrier with natural alignment" + ); + } + Ok(()) + } +} + +/// View one existing CTA-shared allocation as a TMA destination tile. +/// +/// `capacity` counts carrier elements, not bytes or logical FP4 values. +/// `alignment_bytes` is the unsafe caller's promise about this selected stage +/// base; it may be stronger than the carrier's natural alignment. +#[pliron_op( + name = "cute.tma_smem_view", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteTmaSmemViewOp; + +impl CuteTmaSmemViewOp { + /// Build a detached transport view over shared storage and its element capacity. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + base: Value, + capacity: Value, + element: TypeHandle, + smem_layout: ComposedLayout, + alignment_bytes: u64, + ) -> Self { + let result = transport_view_type( + ctx, + element, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadWrite, + alignment_bytes, + smem_layout, + ); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![base, capacity], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the selected CTA-shared stage base. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the available carrier-element count from that base. + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } +} + +impl Verify for CuteTmaSmemViewOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tma_smem_view needs 2 operands and 1 result"); + } + let Some(view) = tma_view_of(ctx, op.get_result(0)) else { + return verify_err!(op.loc(), "cute.tma_smem_view result must be a TMA view"); + }; + view.verify(ctx)?; + let tensor = view.tensor_view(ctx).expect("verified TMA tensor view"); + if tensor.space != CuteTensorAddressSpaceAttr::Smem + || tensor.access != CuteTensorAccessAttr::ReadWrite + { + return verify_err!( + op.loc(), + "cute.tma_smem_view must be a writable shared carrier" + ); + } + let element_bytes = i64::try_from( + tensor + .storage_bytes(ctx) + .expect("verified TMA carrier has a byte width"), + ) + .expect("supported TMA carrier width fits i64"); + let required_alignment = + crate::layout::tma_phase_alignment_bytes(&view.smem_layout.0, element_bytes) + .expect("verified TMA layout has a phase alignment"); + let promised_alignment = tensor.alignment.0; + if promised_alignment < required_alignment as u64 { + return verify_err!( + op.loc(), + "cute.tma_smem_view promises {promised_alignment}-byte alignment, but this layout needs {required_alignment} bytes" + ); + } + let base_ty = self.base(ctx).get_type(ctx); + let base_ty = base_ty.deref(ctx); + let Some(base) = base_ty.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tma_smem_view base must be a MIR pointer"); + }; + if base.pointee != tensor.storage + || !base.is_mutable + || base.address_space != address_space::SHARED + { + return verify_err!( + op.loc(), + "cute.tma_smem_view base must be a mutable CTA-shared carrier pointer" + ); + } + if !is_u64(ctx, self.capacity(ctx)) { + return verify_err!( + op.loc(), + "cute.tma_smem_view capacity must be an unsigned 64-bit integer" + ); + } + Ok(()) + } +} + +/// Start one asynchronous two-dimensional copy from a descriptor-backed +/// global view into one selected shared stage. +/// +/// The completion barrier stays explicit because the later pipeline slice +/// owns its initialization, expected byte count, wait, and release protocol. +#[pliron_op( + name = "cute.tma_copy_2d", + format, + interfaces = [NOpdsInterface<5>, NResultsInterface<0>] +)] +pub struct CuteTmaCopy2dOp; + +impl CuteTmaCopy2dOp { + /// Build a detached TMA copy preserving the source/destination views, tile indices, and barrier. + pub fn new( + ctx: &mut Context, + source: Value, + destination: Value, + tile_row: Value, + tile_column: Value, + completion_barrier: Value, + ) -> Self { + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![ + source, + destination, + tile_row, + tile_column, + completion_barrier, + ], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the source view operand of the asynchronous transfer. + #[must_use] + pub fn source(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the destination view operand of the asynchronous transfer. + #[must_use] + pub fn destination(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the runtime row index in units of complete TMA tiles. + #[must_use] + pub fn tile_row(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + /// Return the runtime column index in units of complete TMA tiles. + #[must_use] + pub fn tile_column(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } + + /// Return the shared barrier pointer tracking completion of the transfer. + #[must_use] + pub fn completion_barrier(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(4) + } +} + +impl Verify for CuteTmaCopy2dOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 5 || op.get_num_results() != 0 { + return verify_err!(op.loc(), "cute.tma_copy_2d needs 5 operands and 0 results"); + } + let Some(source) = tma_view_of(ctx, self.source(ctx)) else { + return verify_err!(op.loc(), "cute.tma_copy_2d source must be a TMA view"); + }; + let Some(destination) = tma_view_of(ctx, self.destination(ctx)) else { + return verify_err!(op.loc(), "cute.tma_copy_2d destination must be a TMA view"); + }; + source.verify(ctx)?; + destination.verify(ctx)?; + let source_tensor = source.tensor_view(ctx).expect("verified source tensor"); + let destination_tensor = destination + .tensor_view(ctx) + .expect("verified destination tensor"); + if source_tensor.space != CuteTensorAddressSpaceAttr::Gmem + || source_tensor.access != CuteTensorAccessAttr::ReadOnly + || destination_tensor.space != CuteTensorAddressSpaceAttr::Smem + || destination_tensor.access != CuteTensorAccessAttr::ReadWrite + { + return verify_err!( + op.loc(), + "cute.tma_copy_2d direction must be read-only Gmem to writable Smem" + ); + } + if source_tensor.storage != destination_tensor.storage + || source.smem_layout != destination.smem_layout + { + return verify_err!( + op.loc(), + "cute.tma_copy_2d source and destination need the same carrier and tile layout" + ); + } + if !is_u64(ctx, self.tile_row(ctx)) || !is_u64(ctx, self.tile_column(ctx)) { + return verify_err!( + op.loc(), + "cute.tma_copy_2d tile row and column must be unsigned 64-bit integers" + ); + } + let barrier_ty = self.completion_barrier(ctx).get_type(ctx); + let barrier_ty = barrier_ty.deref(ctx); + let Some(barrier) = barrier_ty.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.tma_copy_2d completion barrier must be a MIR pointer" + ); + }; + let barrier_is_u64 = barrier + .pointee + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()); + if !barrier.is_mutable || barrier.address_space != address_space::SHARED || !barrier_is_u64 + { + return verify_err!( + op.loc(), + "cute.tma_copy_2d completion barrier must be a mutable CTA-shared pointer to canonical unsigned u64 Barrier storage" + ); + } + Ok(()) + } +} + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/tests.rs new file mode 100644 index 0000000000..c1a41e8f40 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemm_tma/tests.rs @@ -0,0 +1,311 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use crate::layout::{Layout, OffsetUnit, Swizzle}; +use dialect_mir::ops::MirUndefOp; +use pliron::builtin::types::Signedness; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Return result zero from the fixture operation for connecting a semantic chain. +fn result(ctx: &Context, op: Ptr) -> Value { + op.deref(ctx).get_result(0) +} + +/// Build the canonical packed-value transport tile used by the four-copy stage fixture. +fn byte_layout() -> ComposedLayout { + let outer: Layout = "(128,64):(64,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 4, 3), 0, outer, OffsetUnit::Elements).unwrap() +} + +/// Build the canonical scale transport tile used by the four-copy stage fixture. +fn scale_layout() -> ComposedLayout { + let outer: Layout = "(1,256):(256,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::IDENTITY, 0, outer, OffsetUnit::Elements).unwrap() +} + +/// Typed placeholder operands shared by the operation-family verifier tests. +struct Inputs { + /// Immutable global tensor-map pointer used by the fixture. + descriptor: Value, + /// Shared byte pointer carrying packed FP4 input data. + byte_base: Value, + /// Shared pointer carrying scale transport elements. + scale_base: Value, + /// Unsigned runtime extent/index value used by the fixture. + size: Value, + /// Shared completion-barrier pointer used by the copy fixture. + barrier: Value, + /// Unsigned byte storage type. + u8_ty: TypeHandle, + /// Unsigned 16-bit transport-carrier type. + u16_ty: TypeHandle, +} + +/// Create descriptor, shared byte/scale pointers, extents, and barriers for transport tests. +fn setup() -> (Context, Inputs) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u16_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let byte_base_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let scale_base_ty: TypeHandle = + MirPtrType::get(&mut ctx, u16_ty, true, address_space::SHARED).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let descriptor = undef(&mut ctx, descriptor_ty); + let byte_base = undef(&mut ctx, byte_base_ty); + let scale_base = undef(&mut ctx, scale_base_ty); + let size = undef(&mut ctx, u64_ty); + let barrier = undef(&mut ctx, barrier_ty); + ( + ctx, + Inputs { + descriptor, + byte_base, + scale_base, + size, + barrier, + u8_ty, + u16_ty, + }, + ) +} + +/// Build a detached transport copy from locally verified source and destination views. +fn copy( + ctx: &mut Context, + descriptor: Value, + base: Value, + size: Value, + barrier: Value, + view: (TypeHandle, ComposedLayout, u64), +) -> CuteTmaCopy2dOp { + let (element, layout, alignment) = view; + let source = CuteTmaGmemViewOp::new(ctx, descriptor, element, layout.clone()); + assert!(source.verify(ctx).is_ok()); + let source = result(ctx, source.get_operation()); + let destination = CuteTmaSmemViewOp::new(ctx, base, size, element, layout, alignment); + assert!(destination.verify(ctx).is_ok()); + let destination = result(ctx, destination.get_operation()); + CuteTmaCopy2dOp::new(ctx, source, destination, size, size, barrier) +} + +/// A/B values and A/B scales remain four explicit copies with individually typed view pairs. +#[test] +fn four_stage_copies_stay_four_composable_operations() { + let (mut ctx, inputs) = setup(); + let a = copy( + &mut ctx, + inputs.descriptor, + inputs.byte_base, + inputs.size, + inputs.barrier, + (inputs.u8_ty, byte_layout(), 1024), + ); + let b = copy( + &mut ctx, + inputs.descriptor, + inputs.byte_base, + inputs.size, + inputs.barrier, + (inputs.u8_ty, byte_layout(), 1024), + ); + let sfa = copy( + &mut ctx, + inputs.descriptor, + inputs.scale_base, + inputs.size, + inputs.barrier, + (inputs.u16_ty, scale_layout(), 512), + ); + let sfb = copy( + &mut ctx, + inputs.descriptor, + inputs.scale_base, + inputs.size, + inputs.barrier, + (inputs.u16_ty, scale_layout(), 512), + ); + + for operation in [&a, &b, &sfa, &sfb] { + assert!(operation.verify(&ctx).is_ok()); + } +} + +/// TMA transports storage carriers without assigning packed-number decoding or matrix roles. +#[test] +fn transport_view_keeps_bytes_plain_and_role_free() { + let (mut ctx, inputs) = setup(); + let source = CuteTmaGmemViewOp::new(&mut ctx, inputs.descriptor, inputs.u8_ty, byte_layout()); + assert!(source.verify(&ctx).is_ok()); + let source = result(&ctx, source.get_operation()); + let view = tma_view_of(&ctx, source).unwrap(); + let tensor = view.tensor_view(&ctx).unwrap(); + + assert_eq!(tensor.format, CuteTensorFormatAttr::Plain); + assert_eq!(tensor.role, CuteTensorRoleAttr::Generic); + assert_eq!(tensor.space, CuteTensorAddressSpaceAttr::Gmem); + assert_eq!(view.tile_elements(), Some(128 * 64)); + assert_eq!(view.tile_bytes(&ctx), Some(128 * 64)); +} + +/// Source and destination transport views must describe the same physical shared tile. +#[test] +fn copy_rejects_a_different_destination_layout() { + let (mut ctx, inputs) = setup(); + let source = CuteTmaGmemViewOp::new(&mut ctx, inputs.descriptor, inputs.u8_ty, byte_layout()); + let source = result(&ctx, source.get_operation()); + let destination = CuteTmaSmemViewOp::new( + &mut ctx, + inputs.byte_base, + inputs.size, + inputs.u8_ty, + ComposedLayout::new( + Swizzle::IDENTITY, + 0, + "(128,64):(64,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + 1024, + ); + let destination = result(&ctx, destination.get_operation()); + let copy = CuteTmaCopy2dOp::new( + &mut ctx, + source, + destination, + inputs.size, + inputs.size, + inputs.barrier, + ); + + assert!( + copy.verify(&ctx) + .unwrap_err() + .to_string() + .contains("same carrier and tile layout") + ); +} + +/// Shared transport metadata cannot manufacture an address-space proof for a generic pointer. +#[test] +fn shared_view_rejects_a_generic_pointer() { + let (mut ctx, inputs) = setup(); + let generic_base_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, inputs.u8_ty, true).into(); + let generic_base = undef(&mut ctx, generic_base_ty); + let destination = CuteTmaSmemViewOp::new( + &mut ctx, + generic_base, + inputs.size, + inputs.u8_ty, + byte_layout(), + 1024, + ); + + assert!( + destination + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("CTA-shared") + ); +} + +/// The destination base must satisfy the full swizzle-period alignment promise. +#[test] +fn shared_view_rejects_alignment_smaller_than_the_swizzle_phase() { + let (mut ctx, inputs) = setup(); + let destination = CuteTmaSmemViewOp::new( + &mut ctx, + inputs.byte_base, + inputs.size, + inputs.u8_ty, + byte_layout(), + 16, + ); + + assert!( + destination + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("needs 512 bytes") + ); +} + +/// Descriptor-backed transport preserves the immutable tensor-map pointer contract. +#[test] +fn global_view_rejects_a_mutable_descriptor_pointer() { + let (mut ctx, inputs) = setup(); + let mutable_descriptor_ty: TypeHandle = + MirPtrType::get_generic(&mut ctx, inputs.u8_ty, true).into(); + let mutable_descriptor = undef(&mut ctx, mutable_descriptor_ty); + let source = CuteTmaGmemViewOp::new(&mut ctx, mutable_descriptor, inputs.u8_ty, byte_layout()); + + assert!( + source + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("immutable generic pointer") + ); +} + +/// Completion requires mutable shared u64 barrier storage, not merely a pointer-shaped carrier. +#[test] +fn copy_rejects_generic_or_wrong_pointee_barriers() { + let (mut ctx, inputs) = setup(); + let u64_ty = inputs.size.get_type(&ctx); + let generic_barrier_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, true).into(); + let generic_barrier = undef(&mut ctx, generic_barrier_ty); + let generic = copy( + &mut ctx, + inputs.descriptor, + inputs.byte_base, + inputs.size, + generic_barrier, + (inputs.u8_ty, byte_layout(), 1024), + ); + assert!( + generic + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("canonical unsigned u64 Barrier") + ); + + let wrong_barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, inputs.u8_ty, true, address_space::SHARED).into(); + let wrong_barrier = undef(&mut ctx, wrong_barrier_ty); + let wrong = copy( + &mut ctx, + inputs.descriptor, + inputs.byte_base, + inputs.size, + wrong_barrier, + (inputs.u8_ty, byte_layout(), 1024), + ); + assert!( + wrong + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("canonical unsigned u64 Barrier") + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/mod.rs new file mode 100644 index 0000000000..6fca5a60ca --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/mod.rs @@ -0,0 +1,818 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! High-level block-scaled tensor operations used by GEMV. +//! +//! The dialect keeps the same small steps as the Rust kernel: +//! +//! ```text +//! packed E2M1 tensor ──┐ +//! ├─ scaled_view ─ row ─ KTile<64> ─ load ─ fragment ─┐ +//! UE8M0 scale tensor ──┘ │ +//! ├─ dot(acc) +//! packed E2M1 tensor ──┐ │ +//! ├─ scaled_view ─ row ─ KTile<64> ─ load ─ fragment ─┘ +//! UE8M0 scale tensor ──┘ +//! ``` +//! +//! Every arrow is its own operation. Nothing here means “run GEMV”. The +//! selected backend can lower the load and dot either to native pointer and +//! arithmetic operations or to the corresponding CuTe MLIR steps. + +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface, OneOpdInterface, OneResultInterface}, + types::{FP32Type, IntegerType, Signedness}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteAlignmentAttr, CuteScaledLayoutAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, +}; +use crate::types::{CuteFragmentType, CuteScaledViewType, CuteTensorViewType}; + +/// The first preserved GEMV slice loads 64 logical K values at a time. +pub const GEMV_K_TILE_WIDTH: u64 = 64; + +/// Downcast an SSA operand's type to its packed tensor-view facts. +fn tensor_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Downcast an SSA operand's type to its bound value/scale view. +fn scaled_view_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Read the source-view facts retained by a loaded block-scaled fragment. +fn fragment_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Check for the FP32 scalar type required by accumulator carriers. +fn is_f32(ctx: &Context, ty: TypeHandle) -> bool { + ty.deref(ctx).downcast_ref::().is_some() +} + +/// Compare the bound value/scale types and role, independently of the selected subview. +fn same_scaled_members(left: &CuteScaledViewType, right: &CuteScaledViewType) -> bool { + left.values == right.values && left.scales == right.scales && left.role == right.role +} + +/// Compare logical/storage types, format, and layout while allowing other view facts to differ. +fn same_tensor_facts_except_role(left: &CuteTensorViewType, right: &CuteTensorViewType) -> bool { + left.logical == right.logical + && left.storage == right.storage + && left.format == right.format + && left.layout == right.layout +} + +/// Require compatible matrix/vector fragments for the supported block-scaled FP32 dot. +fn fragments_can_dot(ctx: &Context, matrix: &CuteFragmentType, vector: &CuteFragmentType) -> bool { + let Some(matrix_source) = matrix.source_view(ctx) else { + return false; + }; + let Some(vector_source) = vector.source_view(ctx) else { + return false; + }; + if matrix_source.role != CuteTensorRoleAttr::Mkl + || vector_source.role != CuteTensorRoleAttr::Nkl + || matrix_source.layout != vector_source.layout + { + return false; + } + + let matrix_values_ty = matrix_source.values.deref(ctx); + let Some(matrix_values) = matrix_values_ty.downcast_ref::() else { + return false; + }; + let vector_values_ty = vector_source.values.deref(ctx); + let Some(vector_values) = vector_values_ty.downcast_ref::() else { + return false; + }; + let matrix_scales_ty = matrix_source.scales.deref(ctx); + let Some(matrix_scales) = matrix_scales_ty.downcast_ref::() else { + return false; + }; + let vector_scales_ty = vector_source.scales.deref(ctx); + let Some(vector_scales) = vector_scales_ty.downcast_ref::() else { + return false; + }; + + same_tensor_facts_except_role(matrix_values, vector_values) + && same_tensor_facts_except_role(matrix_scales, vector_scales) +} + +/// Recover rows and K extents only from a direct two-dimensional tensor constructor. +fn tensor_make_2d_shape(ctx: &Context, value: Value) -> Option<(Value, Value)> { + let producer = value.defining_op()?; + if Operation::get_opid(producer, ctx) != CuteTensorMake2DOp::get_opid_static() { + return None; + } + let producer = producer.deref(ctx); + Some((producer.get_operand(2), producer.get_operand(3))) +} + +/// Describe one packed or scale tensor using its two logical dimensions. +/// +/// `len` counts storage elements. `rows` and `k` count logical values. The +/// operation creates a view only; it does not read the pointer. +#[pliron_op( + name = "cute.tensor_make_2d", + format, + interfaces = [NOpdsInterface<4>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteTensorMake2DOp; + +impl CuteTensorMake2DOp { + /// Create a read-only global tensor with explicit format, role, and layout. + /// + /// `alignment_bytes` describes the storage pointer at view construction. + /// Pass `1` when the safe source API has no stronger guarantee. An unsafe + /// load records its selected-address promise on the load operation. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + data: Value, + len: Value, + rows: Value, + k: Value, + logical: TypeHandle, + storage: TypeHandle, + format: CuteTensorFormatAttr, + role: CuteTensorRoleAttr, + alignment_bytes: u64, + layout: CuteTensorLayoutAttr, + ) -> Self { + let result: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + alignment_bytes, + format, + role, + layout, + ) + .into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![data, len, rows, k], + vec![], + 0, + ), + } + } + + /// Create the packed E2M1 value view used by the GEMV matrix or vector. + #[allow(clippy::too_many_arguments)] + pub fn new_e2m1( + ctx: &mut Context, + data: Value, + len: Value, + rows: Value, + k: Value, + role: CuteTensorRoleAttr, + alignment_bytes: u64, + ) -> Self { + let logical: TypeHandle = FP32Type::get(ctx).into(); + let storage: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + Self::new( + ctx, + data, + len, + rows, + k, + logical, + storage, + CuteTensorFormatAttr::E2M1, + role, + alignment_bytes, + CuteTensorLayoutAttr::KMajor, + ) + } + + /// Create the UE8M0 scale view paired with an E2M1 tensor. + #[allow(clippy::too_many_arguments)] + pub fn new_ue8m0( + ctx: &mut Context, + data: Value, + len: Value, + rows: Value, + k: Value, + role: CuteTensorRoleAttr, + alignment_bytes: u64, + values_per_scale: u64, + ) -> Self { + let logical: TypeHandle = FP32Type::get(ctx).into(); + let storage: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + Self::new( + ctx, + data, + len, + rows, + k, + logical, + storage, + CuteTensorFormatAttr::UE8M0, + role, + alignment_bytes, + CuteTensorLayoutAttr::BlockScaleKMajor(values_per_scale), + ) + } + + /// Create the E4M3FN scale view paired with an E2M1 tensor (NVFP4). + #[allow(clippy::too_many_arguments)] + pub fn new_e4m3( + ctx: &mut Context, + data: Value, + len: Value, + rows: Value, + k: Value, + role: CuteTensorRoleAttr, + alignment_bytes: u64, + values_per_scale: u64, + ) -> Self { + let logical: TypeHandle = FP32Type::get(ctx).into(); + let storage: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + Self::new( + ctx, + data, + len, + rows, + k, + logical, + storage, + CuteTensorFormatAttr::E4M3, + role, + alignment_bytes, + CuteTensorLayoutAttr::BlockScaleKMajor(values_per_scale), + ) + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the storage-element count from the source slice. + #[must_use] + pub fn storage_len(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the logical row count. Packed storage does not change it. + #[must_use] + pub fn logical_rows(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } + + /// Return the logical K count. For E2M1, storage uses `K / 2` bytes. + #[must_use] + pub fn logical_k(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(3) + } +} + +impl Verify for CuteTensorMake2DOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 4 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.tensor_make_2d needs 4 operands and 1 result" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(view) = result_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tensor_make_2d result must be a tensor view"); + }; + view.verify(ctx)?; + if view.space != CuteTensorAddressSpaceAttr::Gmem + || view.access != CuteTensorAccessAttr::ReadOnly + { + return verify_err!( + op.loc(), + "cute.tensor_make_2d must produce a read-only global tensor view" + ); + } + if !matches!( + (view.format, view.layout), + (CuteTensorFormatAttr::E2M1, CuteTensorLayoutAttr::KMajor) + | ( + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3, + CuteTensorLayoutAttr::BlockScaleKMajor(_) + ) + ) { + return verify_err!( + op.loc(), + "cute.tensor_make_2d needs an E2M1 value or UE8M0/E4M3 scale layout" + ); + } + + let data_ty = op.get_operand(0).get_type(ctx); + let data_ty_ref = data_ty.deref(ctx); + let Some(data) = data_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tensor_make_2d data must be a MIR pointer"); + }; + if data.pointee != view.storage { + return verify_err!( + op.loc(), + "cute.tensor_make_2d data pointee must match tensor storage" + ); + } + if ![address_space::GENERIC, address_space::GLOBAL].contains(&data.address_space) { + return verify_err!( + op.loc(), + "cute.tensor_make_2d data must point at global or generic memory" + ); + } + for (index, name) in [(1usize, "length"), (2, "rows"), (3, "K")] { + if !is_u64(ctx, op.get_operand(index)) { + return verify_err!( + op.loc(), + "cute.tensor_make_2d {name} must be an unsigned 64-bit integer" + ); + } + } + Ok(()) + } +} + +/// Bind a packed value tensor to the scale tensor that explains its range. +#[pliron_op( + name = "cute.scaled_view_make", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteScaledViewMakeOp; + +impl CuteScaledViewMakeOp { + /// Bind value and scale tensor ghosts into one detached scaled-view operation. + pub fn new(ctx: &mut Context, values: Value, scales: Value) -> Self { + let values_view = tensor_view_of(ctx, values) + .expect("cute.scaled_view_make builder needs a tensor-view values input"); + let result: TypeHandle = CuteScaledViewType::get( + ctx, + values.get_type(ctx), + scales.get_type(ctx), + values_view.role, + CuteScaledLayoutAttr::Full, + ) + .into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![values, scales], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteScaledViewMakeOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scaled_view_make needs 2 operands and 1 result" + ); + } + if tensor_view_of(ctx, op.get_operand(0)).is_none() + || tensor_view_of(ctx, op.get_operand(1)).is_none() + { + return verify_err!( + op.loc(), + "cute.scaled_view_make inputs must be tensor views" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.scaled_view_make result must be a scaled view" + ); + }; + if result.values != op.get_operand(0).get_type(ctx) + || result.scales != op.get_operand(1).get_type(ctx) + || result.layout != CuteScaledLayoutAttr::Full + { + return verify_err!( + op.loc(), + "cute.scaled_view_make result must bind its two inputs at Full layout" + ); + } + let Some((value_rows, value_k)) = tensor_make_2d_shape(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.scaled_view_make values must come directly from cute.tensor_make_2d" + ); + }; + let Some((scale_rows, scale_k)) = tensor_make_2d_shape(ctx, op.get_operand(1)) else { + return verify_err!( + op.loc(), + "cute.scaled_view_make scales must come directly from cute.tensor_make_2d" + ); + }; + if value_rows != scale_rows || value_k != scale_k { + return verify_err!( + op.loc(), + "cute.scaled_view_make values and scales must share the same logical rows and K" + ); + } + result.verify(ctx) + } +} + +/// Select one logical row without reading values or scales. +#[pliron_op( + name = "cute.scaled_view_row", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteScaledViewRowOp; + +impl CuteScaledViewRowOp { + /// Build a detached selection of one batch and row from a scaled tensor view. + pub fn new(ctx: &mut Context, tensor: Value, batch: Value, row: Value) -> Self { + let source = scaled_view_of(ctx, tensor) + .expect("cute.scaled_view_row builder needs a scaled-view input"); + let result: TypeHandle = source.with_layout(ctx, CuteScaledLayoutAttr::Row).into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![tensor, batch, row], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteScaledViewRowOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scaled_view_row needs 3 operands and 1 result" + ); + } + let Some(input) = scaled_view_of(ctx, op.get_operand(0)) else { + return verify_err!(op.loc(), "cute.scaled_view_row input must be a scaled view"); + }; + if input.layout != CuteScaledLayoutAttr::Full { + return verify_err!(op.loc(), "cute.scaled_view_row input must use Full layout"); + } + if !is_u64(ctx, op.get_operand(1)) || !is_u64(ctx, op.get_operand(2)) { + return verify_err!( + op.loc(), + "cute.scaled_view_row batch and row must be unsigned 64-bit integers" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.scaled_view_row result must be a scaled view" + ); + }; + if !same_scaled_members(&input, result) || result.layout != CuteScaledLayoutAttr::Row { + return verify_err!( + op.loc(), + "cute.scaled_view_row may change only Full layout to Row" + ); + } + result.verify(ctx) + } +} + +/// Select 64 neighboring logical K values and their four scales. +#[pliron_op( + name = "cute.scaled_view_k_tile", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteScaledViewKTileOp; + +impl CuteScaledViewKTileOp { + /// Build a detached K-tile selection with the source row's packed-format facts. + pub fn new(ctx: &mut Context, row: Value, tile_index: Value) -> Self { + let source = scaled_view_of(ctx, row) + .expect("cute.scaled_view_k_tile builder needs a scaled row input"); + let result: TypeHandle = source + .with_layout(ctx, CuteScaledLayoutAttr::KTile(GEMV_K_TILE_WIDTH)) + .into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![row, tile_index], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteScaledViewKTileOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile needs 2 operands and 1 result" + ); + } + let Some(input) = scaled_view_of(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile input must be a scaled view" + ); + }; + if input.layout != CuteScaledLayoutAttr::Row { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile input must use Row layout" + ); + } + if !is_u64(ctx, op.get_operand(1)) { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile index must be an unsigned 64-bit integer" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile result must be a scaled view" + ); + }; + if !same_scaled_members(&input, result) + || result.layout != CuteScaledLayoutAttr::KTile(GEMV_K_TILE_WIDTH) + { + return verify_err!( + op.loc(), + "cute.scaled_view_k_tile may change only Row layout to KTile<64>" + ); + } + result.verify(ctx) + } +} + +/// Load one selected K tile into a register fragment. +/// +/// The source view is safe to construct from byte slices and therefore makes +/// no strong alignment claim. This unsafe transfer carries the stronger +/// selected-address promises needed by its two physical reads: +/// +/// ```text +/// 32 packed value bytes -> address aligned to at least 16 bytes +/// 4 scale bytes -> address aligned to at least 4 bytes +/// ``` +#[pliron_op( + name = "cute.scaled_view_load", + format, + interfaces = [ + NOpdsInterface<1>, + NResultsInterface<1>, + OneOpdInterface, + OneResultInterface + ], + attributes = ( + value_alignment_bytes: CuteAlignmentAttr, + scale_alignment_bytes: CuteAlignmentAttr + ) +)] +pub struct CuteScaledViewLoadOp; + +impl CuteScaledViewLoadOp { + /// Build a detached block-scaled load with explicit selected-address alignment promises. + pub fn new( + ctx: &mut Context, + tile: Value, + value_alignment_bytes: u64, + scale_alignment_bytes: u64, + ) -> Self { + let result: TypeHandle = CuteFragmentType::get(ctx, tile.get_type(ctx)).into(); + let load = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result], + vec![tile], + vec![], + 0, + ), + }; + load.set_attr_value_alignment_bytes(ctx, CuteAlignmentAttr(value_alignment_bytes)); + load.set_attr_scale_alignment_bytes(ctx, CuteAlignmentAttr(scale_alignment_bytes)); + load + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the promised alignment of the selected packed-value address. + #[must_use] + pub fn promised_value_alignment(&self, ctx: &Context) -> Option { + self.get_attr_value_alignment_bytes(ctx) + .map(|alignment| alignment.0) + } + + /// Return the promised alignment of the selected scale address. + #[must_use] + pub fn promised_scale_alignment(&self, ctx: &Context) -> Option { + self.get_attr_scale_alignment_bytes(ctx) + .map(|alignment| alignment.0) + } +} + +impl Verify for CuteScaledViewLoadOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scaled_view_load needs 1 operand and 1 result" + ); + } + let Some(input) = scaled_view_of(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.scaled_view_load input must be a scaled view" + ); + }; + if input.layout != CuteScaledLayoutAttr::KTile(GEMV_K_TILE_WIDTH) { + return verify_err!(op.loc(), "cute.scaled_view_load input must be KTile<64>"); + } + let Some(value_alignment) = self.promised_value_alignment(ctx) else { + return verify_err!( + op.loc(), + "cute.scaled_view_load must carry a packed-value alignment promise" + ); + }; + let Some(scale_alignment) = self.promised_scale_alignment(ctx) else { + return verify_err!( + op.loc(), + "cute.scaled_view_load must carry a scale alignment promise" + ); + }; + if !value_alignment.is_power_of_two() || value_alignment < 16 { + return verify_err!( + op.loc(), + "cute.scaled_view_load needs packed-value alignment of at least 16 bytes" + ); + } + if !scale_alignment.is_power_of_two() || scale_alignment < 4 { + return verify_err!( + op.loc(), + "cute.scaled_view_load needs scale alignment of at least 4 bytes" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.scaled_view_load result must be a fragment"); + }; + if result.source != op.get_operand(0).get_type(ctx) { + return verify_err!( + op.loc(), + "cute.scaled_view_load fragment must remember its input tile" + ); + } + result.verify(ctx) + } +} + +/// Add one M-row fragment × one N-row fragment dot product to `acc`. +/// +/// This operation owns no loop or output store. It represents one K=64 +/// reduction step and returns the next `f32` accumulator. +#[pliron_op( + name = "cute.dot", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteDotOp; + +impl CuteDotOp { + /// Build a detached FP32 dot update from matching matrix/vector K fragments. + pub fn new(ctx: &mut Context, matrix: Value, vector: Value, acc: Value) -> Self { + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![acc.get_type(ctx)], + vec![matrix, vector, acc], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteDotOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.dot needs 3 operands and 1 result"); + } + let Some(matrix) = fragment_of(ctx, op.get_operand(0)) else { + return verify_err!(op.loc(), "cute.dot left input must be a fragment"); + }; + let Some(vector) = fragment_of(ctx, op.get_operand(1)) else { + return verify_err!(op.loc(), "cute.dot right input must be a fragment"); + }; + matrix.verify(ctx)?; + vector.verify(ctx)?; + if !fragments_can_dot(ctx, &matrix, &vector) { + return verify_err!(op.loc(), "cute.dot needs compatible Mkl and Nkl fragments"); + } + let acc_ty = op.get_operand(2).get_type(ctx); + if !is_f32(ctx, acc_ty) { + return verify_err!(op.loc(), "cute.dot accumulator must be f32"); + } + if op.get_result(0).get_type(ctx) != acc_ty { + return verify_err!( + op.loc(), + "cute.dot result must match its f32 accumulator type" + ); + } + Ok(()) + } +} + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/tests.rs new file mode 100644 index 0000000000..61748bfbea --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/gemv/tests.rs @@ -0,0 +1,347 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use dialect_mir::ops::MirUndefOp; +use dialect_mir::types::MirFP16Type; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Typed placeholder operands shared by the operation-family verifier tests. +struct Inputs { + /// Global pointer to packed matrix values. + matrix_values: Value, + /// Global pointer to matrix block scales. + matrix_scales: Value, + /// Global pointer to packed vector values. + vector_values: Value, + /// Global pointer to vector block scales. + vector_scales: Value, + /// Unsigned runtime extent/index value used by the fixture. + size: Value, + /// FP32 seed accumulator used by the dot-product fixture. + acc: Value, +} + +/// Create matrix/vector value and scale pointers plus the FP32 dot seed. +fn setup() -> (Context, Inputs) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let global_u8_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, false, address_space::GLOBAL).into(); + let matrix_values = undef(&mut ctx, global_u8_ty); + let matrix_scales = undef(&mut ctx, global_u8_ty); + let vector_values = undef(&mut ctx, global_u8_ty); + let vector_scales = undef(&mut ctx, global_u8_ty); + let size = undef(&mut ctx, u64_ty); + let acc = undef(&mut ctx, f32_ty); + ( + ctx, + Inputs { + matrix_values, + matrix_scales, + vector_values, + vector_scales, + size, + acc, + }, + ) +} + +/// Return result zero from the fixture operation for connecting a semantic chain. +fn result(ctx: &Context, op: Ptr) -> Value { + op.deref(ctx).get_result(0) +} + +/// Bind compatible fixture value and scale tensors for a requested matrix/vector coordinate role. +fn make_scaled( + ctx: &mut Context, + values: Value, + scales: Value, + size: Value, + role: CuteTensorRoleAttr, +) -> (CuteScaledViewMakeOp, Value) { + let values = CuteTensorMake2DOp::new_e2m1(ctx, values, size, size, size, role, 1); + assert!(values.verify(ctx).is_ok()); + let values = result(ctx, values.get_operation()); + let scales = CuteTensorMake2DOp::new_ue8m0(ctx, scales, size, size, size, role, 1, 16); + assert!(scales.verify(ctx).is_ok()); + let scales = result(ctx, scales.get_operation()); + let scaled = CuteScaledViewMakeOp::new(ctx, values, scales); + let scaled_value = result(ctx, scaled.get_operation()); + (scaled, scaled_value) +} + +/// Select a fixture row and K tile, then load its block-scaled register fragment. +fn load_fragment(ctx: &mut Context, scaled: Value, index: Value) -> (CuteScaledViewLoadOp, Value) { + let row = CuteScaledViewRowOp::new(ctx, scaled, index, index); + assert!(row.verify(ctx).is_ok()); + let row = result(ctx, row.get_operation()); + let tile = CuteScaledViewKTileOp::new(ctx, row, index); + assert!(tile.verify(ctx).is_ok()); + let tile = result(ctx, tile.get_operation()); + let load = CuteScaledViewLoadOp::new(ctx, tile, 16, 4); + let fragment = result(ctx, load.get_operation()); + (load, fragment) +} + +/// The packed-value chain remains separate tensor/view/select/load/dot operations rather than a fused kernel call. +#[test] +fn gemv_story_stays_visible_as_six_small_operations() { + let (mut ctx, inputs) = setup(); + let (matrix, matrix_value) = make_scaled( + &mut ctx, + inputs.matrix_values, + inputs.matrix_scales, + inputs.size, + CuteTensorRoleAttr::Mkl, + ); + assert!(matrix.verify(&ctx).is_ok()); + let (vector, vector_value) = make_scaled( + &mut ctx, + inputs.vector_values, + inputs.vector_scales, + inputs.size, + CuteTensorRoleAttr::Nkl, + ); + assert!(vector.verify(&ctx).is_ok()); + + let (matrix_load, matrix_fragment) = load_fragment(&mut ctx, matrix_value, inputs.size); + let (vector_load, vector_fragment) = load_fragment(&mut ctx, vector_value, inputs.size); + assert!(matrix_load.verify(&ctx).is_ok()); + assert!(vector_load.verify(&ctx).is_ok()); + + let dot = CuteDotOp::new(&mut ctx, matrix_fragment, vector_fragment, inputs.acc); + assert!(dot.verify(&ctx).is_ok()); + assert_eq!( + dot.get_operation().deref(&ctx).get_result(0).get_type(&ctx), + inputs.acc.get_type(&ctx) + ); +} + +/// Values and their scales must address the same matrix or vector row coordinate. +#[test] +fn scaled_view_rejects_a_scale_tensor_with_the_wrong_role() { + let (mut ctx, inputs) = setup(); + let values = CuteTensorMake2DOp::new_e2m1( + &mut ctx, + inputs.matrix_values, + inputs.size, + inputs.size, + inputs.size, + CuteTensorRoleAttr::Mkl, + 1, + ); + let values = result(&ctx, values.get_operation()); + let scales = CuteTensorMake2DOp::new_ue8m0( + &mut ctx, + inputs.matrix_scales, + inputs.size, + inputs.size, + inputs.size, + CuteTensorRoleAttr::Nkl, + 1, + 16, + ); + let scales = result(&ctx, scales.get_operation()); + let scaled = CuteScaledViewMakeOp::new(&mut ctx, values, scales); + + assert!( + scaled + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("share one Mkl or Nkl role") + ); +} + +/// A dot pairs matrix M and vector N fragments; two matrix-role inputs are invalid. +#[test] +fn dot_rejects_two_mkl_fragments() { + let (mut ctx, inputs) = setup(); + let (_, left) = make_scaled( + &mut ctx, + inputs.matrix_values, + inputs.matrix_scales, + inputs.size, + CuteTensorRoleAttr::Mkl, + ); + let (_, right) = make_scaled( + &mut ctx, + inputs.vector_values, + inputs.vector_scales, + inputs.size, + CuteTensorRoleAttr::Mkl, + ); + let (_, left) = load_fragment(&mut ctx, left, inputs.size); + let (_, right) = load_fragment(&mut ctx, right, inputs.size); + let dot = CuteDotOp::new(&mut ctx, left, right, inputs.acc); + + assert!( + dot.verify(&ctx) + .unwrap_err() + .to_string() + .contains("Mkl and Nkl") + ); +} + +/// The supported GEMV fragment requires the expected number of logical values per scale. +#[test] +fn scaled_view_rejects_the_wrong_scale_group_width() { + let (mut ctx, inputs) = setup(); + let values = CuteTensorMake2DOp::new_e2m1( + &mut ctx, + inputs.matrix_values, + inputs.size, + inputs.size, + inputs.size, + CuteTensorRoleAttr::Mkl, + 1, + ); + let values = result(&ctx, values.get_operation()); + let scales = CuteTensorMake2DOp::new_ue8m0( + &mut ctx, + inputs.matrix_scales, + inputs.size, + inputs.size, + inputs.size, + CuteTensorRoleAttr::Mkl, + 1, + 32, + ); + let scales = result(&ctx, scales.get_operation()); + let scaled = CuteScaledViewMakeOp::new(&mut ctx, values, scales); + + assert!( + scaled + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("BlockScaleKMajor<16>") + ); +} + +/// Equal static tensor types do not excuse differing runtime value and scale extents. +#[test] +fn scaled_view_rejects_different_runtime_shapes() { + let (mut ctx, inputs) = setup(); + let u64_ty = inputs.size.get_type(&ctx); + let other_rows = undef(&mut ctx, u64_ty); + let values = CuteTensorMake2DOp::new_e2m1( + &mut ctx, + inputs.matrix_values, + inputs.size, + inputs.size, + inputs.size, + CuteTensorRoleAttr::Mkl, + 1, + ); + let values = result(&ctx, values.get_operation()); + let scales = CuteTensorMake2DOp::new_ue8m0( + &mut ctx, + inputs.matrix_scales, + inputs.size, + other_rows, + inputs.size, + CuteTensorRoleAttr::Mkl, + 1, + 16, + ); + let scales = result(&ctx, scales.get_operation()); + let scaled = CuteScaledViewMakeOp::new(&mut ctx, values, scales); + + assert!( + scaled + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("same logical rows and K") + ); +} + +/// A selected tile load may carry stronger alignment than the original tensor base promises. +#[test] +fn load_owns_the_selected_address_alignment_promises() { + let (mut ctx, inputs) = setup(); + let (_, scaled) = make_scaled( + &mut ctx, + inputs.matrix_values, + inputs.matrix_scales, + inputs.size, + CuteTensorRoleAttr::Mkl, + ); + let row = CuteScaledViewRowOp::new(&mut ctx, scaled, inputs.size, inputs.size); + let row = result(&ctx, row.get_operation()); + let tile = CuteScaledViewKTileOp::new(&mut ctx, row, inputs.size); + let tile = result(&ctx, tile.get_operation()); + + let weak_values = CuteScaledViewLoadOp::new(&mut ctx, tile, 8, 4); + assert!( + weak_values + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("at least 16 bytes") + ); + let weak_scales = CuteScaledViewLoadOp::new(&mut ctx, tile, 16, 2); + assert!( + weak_scales + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("at least 4 bytes") + ); + let valid = CuteScaledViewLoadOp::new(&mut ctx, tile, 16, 4); + assert_eq!(valid.promised_value_alignment(&ctx), Some(16)); + assert_eq!(valid.promised_scale_alignment(&ctx), Some(4)); + assert!(valid.verify(&ctx).is_ok()); +} + +/// A tensor constructor cannot claim shared storage for an ordinary global source pointer. +#[test] +fn tensor_make_2d_rejects_shared_result_facts_over_a_global_pointer() { + let (mut ctx, inputs) = setup(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(&ctx).into(); + let forged: TypeHandle = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + f16_ty, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + 2, + CuteTensorFormatAttr::E2M1, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::KMajor, + ) + .into(); + let raw = Operation::new( + &mut ctx, + CuteTensorMake2DOp::get_concrete_op_info(), + vec![forged], + vec![inputs.matrix_values, inputs.size, inputs.size, inputs.size], + vec![], + 0, + ); + let make = CuteTensorMake2DOp::wrap(raw); + + let error = make.verify(&ctx).unwrap_err(); + assert!( + error.to_string().contains("read-only global"), + "unexpected verifier error: {error}" + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/fragment.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/fragment.rs new file mode 100644 index 0000000000..209d478dde --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/fragment.rs @@ -0,0 +1,191 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Accumulator initialization and scale-fragment K selection. + +use super::*; + +/// Fill one physical register carrier while retaining its fragment meaning. +/// +/// V0 uses this for the 64-FP32 accumulator carrier only. +#[pliron_op( + name = "cute.fragment_fill", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>, OneResultInterface], + attributes = ( + fragment_fill_plan: CuteTiledMmaPlanAttr, + fragment_fill_kind: CuteMmaCarrierKindAttr + ) +)] +pub struct CuteFragmentFillOp; + +impl CuteFragmentFillOp { + /// Build a detached ordinary register aggregate initialized from the supplied FP32 scalar. + pub fn new( + ctx: &mut Context, + fill: Value, + result_type: TypeHandle, + plan: CuteTiledMmaPlanAttr, + kind: CuteMmaCarrierKindAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_type], + vec![fill], + vec![], + 0, + ), + }; + operation.set_attr_fragment_fill_plan(ctx, plan); + operation.set_attr_fragment_fill_kind(ctx, kind); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the FP32 value used to initialize every accumulator register. + #[must_use] + pub fn fill(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the ordinary register aggregate produced by this operation. + #[must_use] + pub fn fragment(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_fragment_fill_plan(ctx)) + } + + /// Return the semantic role assigned to the ordinary register aggregate. + #[must_use] + pub fn kind(&self, ctx: &Context) -> Option { + self.get_attr_fragment_fill_kind(ctx).map(|kind| *kind) + } +} + +impl Verify for CuteFragmentFillOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.fragment_fill needs 1 operand and 1 result"); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.fragment_fill must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.fragment_fill {message}"); + } + if self.kind(ctx) != Some(CuteMmaCarrierKindAttr::Accumulator) + || !is_f32(ctx, self.fill(ctx)) + { + return verify_err!( + op.loc(), + "cute.fragment_fill v0 fills an accumulator from one f32 value" + ); + } + let expected = MmaCarrierShape::f32( + plan.accumulator_registers_per_lane() + .expect("verified plan has an accumulator size"), + ); + if let Err(message) = check_carrier(ctx, self.fragment(ctx), expected) { + return verify_err!(op.loc(), "cute.fragment_fill result {message}"); + } + Ok(()) + } +} +/// Select the low or high K=64 scale pair from every packed K=128 word. +#[pliron_op( + name = "cute.fragment_slice_k", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface], + attributes = (fragment_slice_k_plan: CuteTiledMmaPlanAttr) +)] +pub struct CuteFragmentSliceKOp; + +impl CuteFragmentSliceKOp { + /// Build a detached K=64 selection from a loaded stage-scale aggregate. + pub fn new( + ctx: &mut Context, + stage_scales: Value, + k_half: Value, + result_type: TypeHandle, + plan: CuteTiledMmaPlanAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_type], + vec![stage_scales, k_half], + vec![], + 0, + ), + }; + operation.set_attr_fragment_slice_k_plan(ctx, plan); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_fragment_slice_k_plan(ctx)) + } + + /// Return the ordinary scale-register aggregate for this K selection. + #[must_use] + pub fn scales(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } +} + +impl Verify for CuteFragmentSliceKOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.fragment_slice_k needs 2 operands and 1 result" + ); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.fragment_slice_k must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.fragment_slice_k {message}"); + } + if plan.cta_k != plan.atom.k * 2 || !is_u64(ctx, op.get_operand(1)) { + return verify_err!( + op.loc(), + "cute.fragment_slice_k v0 selects one u64-indexed K=64 half of K=128" + ); + } + let expected = MmaCarrierShape::u32( + plan.scale_words_per_lane() + .expect("verified plan has a scale carrier size"), + ); + if let Err(message) = check_carrier(ctx, op.get_operand(0), expected) { + return verify_err!(op.loc(), "cute.fragment_slice_k input {message}"); + } + if let Err(message) = check_carrier(ctx, self.scales(ctx), expected) { + return verify_err!(op.loc(), "cute.fragment_slice_k result {message}"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/gemm.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/gemm.rs new file mode 100644 index 0000000000..d4c153bad4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/gemm.rs @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tiled GEMM over ordinary register and shared-memory carriers. + +use super::*; + +/// Multiply one A K=64 bundle by one lazy B selection and add it to C. +/// +/// B arrives as four scalar carriers, never as eight eager fragments. Backend +/// A must honor `plan.b_load_group`: load two B fragments, issue their four +/// MMA atoms, then move to the next pair. +#[pliron_op( + name = "cute.tiled_gemm", + format, + interfaces = [NOpdsInterface<8>, NResultsInterface<1>, OneResultInterface], + attributes = ( + tiled_gemm_plan: CuteTiledMmaPlanAttr, + tiled_gemm_b_view: TypeAttr + ) +)] +pub struct CuteTiledGemmOp; + +impl CuteTiledGemmOp { + /// Build a detached accumulator update from typed A/scales and lazy shared B carriers. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + lane: Value, + a: Value, + b_base: Value, + b_capacity: Value, + b_warp_n: Value, + b_k_half: Value, + scales: Value, + accumulator: Value, + b_view: TypeHandle, + plan: CuteTiledMmaPlanAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![accumulator.get_type(ctx)], + vec![ + lane, + a, + b_base, + b_capacity, + b_warp_n, + b_k_half, + scales, + accumulator, + ], + vec![], + 0, + ), + }; + operation.set_attr_tiled_gemm_plan(ctx, plan); + operation.set_attr_tiled_gemm_b_view(ctx, TypeAttr::new(b_view)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_tiled_gemm_plan(ctx)) + } + + /// Static N-role shared tensor facts carried by the four lazy B scalars. + #[must_use] + pub fn b_view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_tiled_gemm_b_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } + + /// Return the ordinary FP32 accumulator carrier. + #[must_use] + pub fn accumulator(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } +} + +impl Verify for CuteTiledGemmOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 8 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tiled_gemm needs 8 operands and 1 result"); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.tiled_gemm must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.tiled_gemm {message}"); + } + if !is_u32(ctx, op.get_operand(0)) { + return verify_err!(op.loc(), "cute.tiled_gemm lane must be an unsigned u32"); + } + let expected_a = MmaCarrierShape::u32( + plan.a_registers_per_lane() + .expect("verified plan has an A carrier size"), + ); + if let Err(message) = check_carrier(ctx, op.get_operand(1), expected_a) { + return verify_err!(op.loc(), "cute.tiled_gemm A input {message}"); + } + let Some(b_view) = self.b_view(ctx) else { + return verify_err!( + op.loc(), + "cute.tiled_gemm must carry an N-role shared view for lazy B" + ); + }; + if let Err(message) = check_data_view(ctx, &b_view, &plan, CuteTensorRoleAttr::Nkl) { + return verify_err!(op.loc(), "cute.tiled_gemm B input {message}"); + } + if !shared_carriers_match_view(ctx, op.get_operand(2), op.get_operand(3), &b_view) + || !is_u64(ctx, op.get_operand(4)) + || !is_u64(ctx, op.get_operand(5)) + { + return verify_err!( + op.loc(), + "cute.tiled_gemm lazy B must be pointer, capacity, warp-N, and K-half scalars matching its typed view" + ); + } + let expected_scales = MmaCarrierShape::u32( + plan.scale_words_per_lane() + .expect("verified plan has a scale carrier size"), + ); + if let Err(message) = check_carrier(ctx, op.get_operand(6), expected_scales) { + return verify_err!(op.loc(), "cute.tiled_gemm scale input {message}"); + } + let expected_accumulator = MmaCarrierShape::f32( + plan.accumulator_registers_per_lane() + .expect("verified plan has an accumulator size"), + ); + if let Err(message) = check_carrier(ctx, op.get_operand(7), expected_accumulator) { + return verify_err!(op.loc(), "cute.tiled_gemm accumulator input {message}"); + } + if self.accumulator(ctx).get_type(ctx) != op.get_operand(7).get_type(ctx) { + return verify_err!( + op.loc(), + "cute.tiled_gemm result must preserve the accumulator carrier type" + ); + } + if let Err(message) = check_carrier(ctx, self.accumulator(ctx), expected_accumulator) { + return verify_err!(op.loc(), "cute.tiled_gemm result {message}"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/ldmatrix.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/ldmatrix.rs new file mode 100644 index 0000000000..4000014a79 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/ldmatrix.rs @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Warp-cooperative matrix loads with typed operand roles and shared placement. + +use cute_layout::ComposedLayout; +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::{ + builtin::{ + attributes::TypeAttr, + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::IntegerType, + }, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + result::Error, + r#type::{TypeHandle, Typed}, + value::Value, + verify_err, +}; +use pliron_derive::pliron_op; + +use crate::attributes::{CuteComposedLayoutAttr, CuteMatrixRoleAttr}; + +/// Frozen zero-result ABI for one warp-cooperative `ldmatrix` fragment load. +/// +/// ```text +/// 0 smem base (f16) 2 warp tile column 4 fragment slot (u32 array) +/// 1 warp tile row 3 lane id +/// ``` +/// +/// The shared layout (with any swizzle) and the operand role stay in typed +/// attributes. Backend lowering computes each lane's row address through the +/// composed byte map and selects the target load, storing the returned values +/// into the fragment slot. Operands 1..=2 are `u64`; operand 3 is `u32`. +/// As in the unsafe `cute-rs` matrix-load API, the lane must be the caller's +/// hardware lane (0..32) and the whole selected window must lie in the layout. +/// Lowering may use these bounds to simplify tile-relative address arithmetic; +/// this operation does not insert dynamic window or lane checks. +#[pliron_op( + name = "cute.ldmatrix", + format, + interfaces = [NOpdsInterface<5>, NResultsInterface<0>], + attributes = ( + matrix_role: CuteMatrixRoleAttr, + ldmatrix_smem_layout: CuteComposedLayoutAttr, + ldmatrix_elem: TypeAttr + ) +)] +pub struct CuteLdmatrixOp; + +impl CuteLdmatrixOp { + /// Build a detached warp matrix load with its typed shared-layout and operand-role contract. + pub fn new( + ctx: &mut Context, + operands: [Value; 5], + role: CuteMatrixRoleAttr, + smem_layout: ComposedLayout, + elem: TypeHandle, + ) -> Self { + let op = Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + operands.to_vec(), + vec![], + 0, + ); + let load = Self { op }; + load.set_attr_matrix_role(ctx, role); + load.set_attr_ldmatrix_smem_layout(ctx, CuteComposedLayoutAttr(smem_layout)); + load.set_attr_ldmatrix_elem(ctx, TypeAttr::new(elem)); + load + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteLdmatrixOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = &*self.get_operation().deref(ctx); + let Some(role) = self.get_attr_matrix_role(ctx).map(|attr| *attr) else { + return verify_err!(op.loc(), "cute.ldmatrix must have a matrix role"); + }; + let Some(smem_layout) = self + .get_attr_ldmatrix_smem_layout(ctx) + .map(|attr| attr.0.clone()) + else { + return verify_err!(op.loc(), "cute.ldmatrix must have a shared-memory layout"); + }; + let Some(elem): Option = self + .get_attr_ldmatrix_elem(ctx) + .map(|attr| attr.clone().into()) + else { + return verify_err!(op.loc(), "cute.ldmatrix must have an element type"); + }; + if elem.deref(ctx).downcast_ref::().is_none() { + return verify_err!(op.loc(), "cute.ldmatrix v0 supports f16 elements only"); + } + let _ = role; + let modes = smem_layout.outer().modes(); + if modes.len() != 2 { + return verify_err!(op.loc(), "cute.ldmatrix shared layout needs two modes"); + } + let (Some(rows), Some(columns)) = (modes[0].checked_size(), modes[1].checked_size()) else { + return verify_err!(op.loc(), "cute.ldmatrix shared layout extents are invalid"); + }; + let byte_layout = match smem_layout.to_byte_offsets(2) { + Ok(layout) => layout, + Err(error) => { + return verify_err!( + op.loc(), + "cute.ldmatrix shared layout has no byte form: {error}" + ); + } + }; + if let Err(error) = cute_layout::validate_ldmatrix_source(&byte_layout, rows, columns) { + return verify_err!(op.loc(), "cute.ldmatrix source is not loadable: {error}"); + } + + let smem_ty = op.get_operand(0).get_type(ctx); + let smem_ref = smem_ty.deref(ctx); + let Some(pointer) = smem_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.ldmatrix source must be a MIR pointer"); + }; + if pointer.pointee != elem + || ![address_space::GENERIC, address_space::SHARED].contains(&pointer.address_space) + { + return verify_err!( + op.loc(), + "cute.ldmatrix source must point at shared/generic f16" + ); + } + for (index, role_name) in [(1usize, "warp tile row"), (2usize, "warp tile column")] { + let ty = op.get_operand(index).get_type(ctx); + if ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 64 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "{role_name} must be an unsigned 64-bit integer"); + } + } + let lane_ty = op.get_operand(3).get_type(ctx); + if lane_ty + .deref(ctx) + .downcast_ref::() + .is_none_or(|integer| integer.width() != 32 || !integer.is_unsigned()) + { + return verify_err!(op.loc(), "lane id must be an unsigned 32-bit integer"); + } + let slot_ty = op.get_operand(4).get_type(ctx); + let slot_ref = slot_ty.deref(ctx); + let Some(slot) = slot_ref.downcast_ref::() else { + return verify_err!(op.loc(), "fragment slot must be a MIR pointer"); + }; + if !slot.is_mutable { + return verify_err!(op.loc(), "fragment slot must be mutable"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/load.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/load.rs new file mode 100644 index 0000000000..20e581f3ff --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/load.rs @@ -0,0 +1,404 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared scale and A loads with lazy B selection. + +use super::*; + +/// Load the ten packed scale words used by one warp for a K=128 stage. +#[pliron_op( + name = "cute.mma_load_scales", + format, + interfaces = [NOpdsInterface<7>, NResultsInterface<1>, OneResultInterface], + attributes = ( + mma_load_scales_plan: CuteTiledMmaPlanAttr, + mma_load_scales_a_view: TypeAttr, + mma_load_scales_b_view: TypeAttr + ) +)] +pub struct CuteMmaLoadScalesOp; + +impl CuteMmaLoadScalesOp { + /// Build a detached stage-scale load retaining both operand views and lane selection. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + lane: Value, + scale_a_base: Value, + scale_a_capacity: Value, + scale_b_base: Value, + scale_b_capacity: Value, + warp_m: Value, + warp_n: Value, + scale_a_view: TypeHandle, + scale_b_view: TypeHandle, + result_type: TypeHandle, + plan: CuteTiledMmaPlanAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_type], + vec![ + lane, + scale_a_base, + scale_a_capacity, + scale_b_base, + scale_b_capacity, + warp_m, + warp_n, + ], + vec![], + 0, + ), + }; + operation.set_attr_mma_load_scales_plan(ctx, plan); + operation.set_attr_mma_load_scales_a_view(ctx, TypeAttr::new(scale_a_view)); + operation.set_attr_mma_load_scales_b_view(ctx, TypeAttr::new(scale_b_view)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_mma_load_scales_plan(ctx)) + } + + /// Static M-role scale tensor facts for operands 1 and 2. + #[must_use] + pub fn scale_a_view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_mma_load_scales_a_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } + + /// Static N-role scale tensor facts for operands 3 and 4. + #[must_use] + pub fn scale_b_view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_mma_load_scales_b_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } + + /// Return the ordinary scale-register aggregate for this K selection. + #[must_use] + pub fn scales(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } +} + +impl Verify for CuteMmaLoadScalesOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 7 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.mma_load_scales needs 7 operands and 1 result" + ); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.mma_load_scales must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.mma_load_scales {message}"); + } + if !is_u32(ctx, op.get_operand(0)) + || !is_u64(ctx, op.get_operand(5)) + || !is_u64(ctx, op.get_operand(6)) + { + return verify_err!( + op.loc(), + "cute.mma_load_scales needs a u32 lane and u64 warp-M/warp-N positions" + ); + } + let Some(scale_a) = self.scale_a_view(ctx) else { + return verify_err!( + op.loc(), + "cute.mma_load_scales must carry an M-role shared scale view" + ); + }; + let Some(scale_b) = self.scale_b_view(ctx) else { + return verify_err!( + op.loc(), + "cute.mma_load_scales must carry an N-role shared scale view" + ); + }; + if let Err(message) = check_scale_view(ctx, &scale_a, &plan, CuteTensorRoleAttr::Mkl) { + return verify_err!(op.loc(), "cute.mma_load_scales A {message}"); + } + if let Err(message) = check_scale_view(ctx, &scale_b, &plan, CuteTensorRoleAttr::Nkl) { + return verify_err!(op.loc(), "cute.mma_load_scales B {message}"); + } + if !shared_carriers_match_view(ctx, op.get_operand(1), op.get_operand(2), &scale_a) + || !shared_carriers_match_view(ctx, op.get_operand(3), op.get_operand(4), &scale_b) + { + return verify_err!( + op.loc(), + "cute.mma_load_scales pointer/capacity carriers must match their typed scale views" + ); + } + let expected = MmaCarrierShape::u32( + plan.scale_words_per_lane() + .expect("verified plan has a scale carrier size"), + ); + if let Err(message) = check_carrier(ctx, self.scales(ctx), expected) { + return verify_err!(op.loc(), "cute.mma_load_scales result {message}"); + } + Ok(()) + } +} +/// Load the two A fragments owned by one warp for one K=64 step. +#[pliron_op( + name = "cute.mma_load_a", + format, + interfaces = [NOpdsInterface<5>, NResultsInterface<1>, OneResultInterface], + attributes = ( + mma_load_a_plan: CuteTiledMmaPlanAttr, + mma_load_a_view: TypeAttr + ) +)] +pub struct CuteMmaLoadAOp; + +impl CuteMmaLoadAOp { + /// Build a detached packed A-register load for the selected lane, warp-M, and K position. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + lane: Value, + base: Value, + capacity: Value, + warp_m: Value, + k_half: Value, + view: TypeHandle, + result_type: TypeHandle, + plan: CuteTiledMmaPlanAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_type], + vec![lane, base, capacity, warp_m, k_half], + vec![], + 0, + ), + }; + operation.set_attr_mma_load_a_plan(ctx, plan); + operation.set_attr_mma_load_a_view(ctx, TypeAttr::new(view)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_mma_load_a_plan(ctx)) + } + + /// Static M-role shared tensor facts for the pointer/capacity operands. + #[must_use] + pub fn view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_mma_load_a_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } + + /// Return the ordinary register aggregate produced by this operation. + #[must_use] + pub fn fragment(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } +} + +impl Verify for CuteMmaLoadAOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 5 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.mma_load_a needs 5 operands and 1 result"); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.mma_load_a must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.mma_load_a {message}"); + } + if !is_u32(ctx, op.get_operand(0)) + || !is_u64(ctx, op.get_operand(3)) + || !is_u64(ctx, op.get_operand(4)) + { + return verify_err!( + op.loc(), + "cute.mma_load_a needs a u32 lane and u64 warp-M/K-half positions" + ); + } + let Some(view) = self.view(ctx) else { + return verify_err!( + op.loc(), + "cute.mma_load_a must carry an M-role shared value view" + ); + }; + if let Err(message) = check_data_view(ctx, &view, &plan, CuteTensorRoleAttr::Mkl) { + return verify_err!(op.loc(), "cute.mma_load_a input {message}"); + } + if !shared_carriers_match_view(ctx, op.get_operand(1), op.get_operand(2), &view) { + return verify_err!( + op.loc(), + "cute.mma_load_a pointer/capacity carriers must match its typed A view" + ); + } + let expected = MmaCarrierShape::u32( + plan.a_registers_per_lane() + .expect("verified plan has an A carrier size"), + ); + if let Err(message) = check_carrier(ctx, self.fragment(ctx), expected) { + return verify_err!(op.loc(), "cute.mma_load_a result {message}"); + } + Ok(()) + } +} + +/// Keep B as a lazy shared-memory selection. +/// +/// All four results preserve ordinary operand types. In particular, this op +/// has no register-fragment result and performs no `ldmatrix` load. +#[pliron_op( + name = "cute.mma_partition_b", + format, + interfaces = [NOpdsInterface<4>, NResultsInterface<4>], + attributes = ( + mma_partition_b_plan: CuteTiledMmaPlanAttr, + mma_partition_b_view: TypeAttr + ) +)] +pub struct CuteMmaPartitionBOp; + +impl CuteMmaPartitionBOp { + /// Build a detached lazy B view retaining its base, capacity, warp-N, and K-half carriers. + pub fn new( + ctx: &mut Context, + base: Value, + capacity: Value, + warp_n: Value, + k_half: Value, + view: TypeHandle, + plan: CuteTiledMmaPlanAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ + base.get_type(ctx), + capacity.get_type(ctx), + warp_n.get_type(ctx), + k_half.get_type(ctx), + ], + vec![base, capacity, warp_n, k_half], + vec![], + 0, + ), + }; + operation.set_attr_mma_partition_b_plan(ctx, plan); + operation.set_attr_mma_partition_b_view(ctx, TypeAttr::new(view)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_mma_partition_b_plan(ctx)) + } + + /// Static N-role shared tensor facts for the lazy B selection. + #[must_use] + pub fn view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_mma_partition_b_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the ordinary unsigned storage-capacity carrier, in elements. + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Return the lazy B partition's warp-N coordinate carrier. + #[must_use] + pub fn warp_n(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + /// Return the lazy B partition's K-half coordinate carrier. + #[must_use] + pub fn k_half(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(3) + } +} + +impl Verify for CuteMmaPartitionBOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 4 || op.get_num_results() != 4 { + return verify_err!( + op.loc(), + "cute.mma_partition_b needs 4 operands and exactly 4 scalar results" + ); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.mma_partition_b must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.mma_partition_b {message}"); + } + let Some(view) = self.view(ctx) else { + return verify_err!( + op.loc(), + "cute.mma_partition_b must carry an N-role shared value view" + ); + }; + if let Err(message) = check_data_view(ctx, &view, &plan, CuteTensorRoleAttr::Nkl) { + return verify_err!(op.loc(), "cute.mma_partition_b input {message}"); + } + if !shared_carriers_match_view(ctx, op.get_operand(0), op.get_operand(1), &view) { + return verify_err!( + op.loc(), + "cute.mma_partition_b pointer/capacity carriers must match its typed B view" + ); + } + if !is_u64(ctx, op.get_operand(2)) || !is_u64(ctx, op.get_operand(3)) { + return verify_err!( + op.loc(), + "cute.mma_partition_b warp-N and K-half must be unsigned u64 values" + ); + } + if (0..4) + .any(|index| op.get_result(index).get_type(ctx) != op.get_operand(index).get_type(ctx)) + { + return verify_err!( + op.loc(), + "cute.mma_partition_b must preserve pointer, capacity, warp-N, and K-half types" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/mod.rs new file mode 100644 index 0000000000..579ead2dd0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/mod.rs @@ -0,0 +1,252 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared-tensor and tiled-MMA meaning for the GEMM compute loop. +//! +//! Runtime values stay exactly as they are in MIR today. Typed attributes say +//! what those values mean: +//! +//! ```text +//! shared pointer + capacity ── overlay ─────────────────────────────┐ +//! │ +//! lane ── tiled_mma_slice ──────────────────────────────────────┐ │ +//! │ │ +//! scales A/B ── load_scales ── slice K=64 ──────────────────┐ │ │ +//! A shared tile ─────────────── load_a ───────────────────┐ │ │ │ +//! B shared tile ─────────────── partition_b ───────────┐ │ │ │ │ +//! zero / previous C ───────────────────────────────────┴──┴──┴───┴───┘ +//! │ +//! ▼ +//! tiled_gemm +//! ``` +//! +//! `partition_b` is intentionally lazy. It returns only B's shared pointer, +//! capacity, warp-N position, and K half. `tiled_gemm` can therefore load two +//! B fragments, use them immediately, and then load the next pair. No eager +//! eight-fragment B aggregate exists in this schema. +//! +//! Rust may temporarily store a shared view inside an ordinary aggregate. +//! Consumers therefore repeat its static `!cute.smem_tensor` TypeAttr instead +//! of requiring their pointer/capacity operands to come directly from the +//! overlay operation. Normal and no-inline MIR can flatten to the same ABI. +//! +//! Backend continuations consume the same ordinary carriers: they can erase +//! forwarding operations into native loads and MMA sites or group the values +//! into CuTe MLIR tensors and fragments. + +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::builtin::{ + attributes::TypeAttr, + op_interfaces::{NOpdsInterface, NResultsInterface, OneResultInterface}, + types::{FP32Type, IntegerType}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteMmaCarrierKindAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, + CuteTiledMmaPlanAttr, SM1XX_BLOCK_SCALE_ATOM_BYTES, +}; +use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; +use crate::types::{CuteSmemTensorType, MmaCarrierShape, mma_carrier_shape}; + +/// Check that the SSA value has an unsigned 32-bit integer type. +fn is_u32(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.is_unsigned()) +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Check for the FP32 scalar type required by accumulator carriers. +fn is_f32(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some() +} + +/// Downcast static metadata to the shared tensor's format and placement. +fn smem_tensor_of_type(ctx: &Context, ty: TypeHandle) -> Option { + ty.deref(ctx).downcast_ref::().cloned() +} + +/// Require mutable shared storage with the view's exact storage element type. +fn shared_pointer_matches(ctx: &Context, value: Value, storage: TypeHandle) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|pointer| { + pointer.pointee == storage + && pointer.is_mutable + && pointer.address_space == address_space::SHARED + }) +} + +/// Check the shared pointer and unsigned capacity carrier against the static tensor view. +fn shared_carriers_match_view( + ctx: &Context, + base: Value, + capacity: Value, + view: &CuteSmemTensorType, +) -> bool { + view.tensor_view(ctx).is_some_and(|tensor| { + shared_pointer_matches(ctx, base, tensor.storage) && is_u64(ctx, capacity) + }) +} + +/// Build the supported dense placement for one stage of block-scale carriers. +fn canonical_scale_placement() -> ComposedLayout { + let outer: Layout = "(32,4):(4,1)" + .parse() + .expect("fixed SM120 scale placement is valid"); + ComposedLayout::new(Swizzle::IDENTITY, 0, outer, OffsetUnit::Elements) + .expect("identity scale placement is valid") +} + +/// Build the supported swizzled placement for one stage of packed FP4 values. +fn canonical_data_placement() -> ComposedLayout { + let outer: Layout = "(128,32):(32,1)" + .parse() + .expect("fixed SM120 data placement is valid"); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, outer, OffsetUnit::Elements) + .expect("fixed SM120 data swizzle is valid") +} + +/// Clone an optional borrowed tiled-MMA attribute for a consumer. +fn plan_of_attr( + plan: Option>, +) -> Option { + plan.map(|plan| plan.clone()) +} + +/// Restrict tiled MMA to the supported SM120 shape, atom, and ownership layouts. +fn check_plan(plan: &CuteTiledMmaPlanAttr, ctx: &Context) -> Result<(), &'static str> { + plan.verify(ctx) + .map_err(|_| "has an invalid tiled-MMA plan")?; + let valid_atom = plan.atom == crate::attributes::CuteMmaAtomAttr::mxf4_m16n8k64() + || plan.atom == crate::attributes::CuteMmaAtomAttr::nvf4_m16n8k64(); + if !valid_atom + || plan.cta_m != 128 + || plan.cta_n != 128 + || plan.cta_k != 128 + || plan.warp_m != 4 + || plan.warp_n != 2 + || plan.b_load_group != 2 + || plan.shared_layout.0 != canonical_data_placement() + { + return Err("v0 supports the 128x128x128 SM120 MXFP4/NVFP4 plan"); + } + let canonical = CuteTiledMmaPlanAttr::mxf4_128x128x128(plan.shared_layout.0.clone()); + if plan.m_ownership != canonical.m_ownership || plan.n_ownership != canonical.n_ownership { + return Err("v0 needs the SM120 2x8 warp-cell ownership maps"); + } + Ok(()) +} + +/// Require the packed-value format, operand role, and shared placement needed by tiled MMA. +fn check_data_view( + ctx: &Context, + view: &CuteSmemTensorType, + plan: &CuteTiledMmaPlanAttr, + role: CuteTensorRoleAttr, +) -> Result<(), &'static str> { + view.verify(ctx) + .map_err(|_| "uses an invalid shared tensor type")?; + let tensor = view.tensor_view(ctx).ok_or("must wrap a tensor view")?; + if tensor.format != CuteTensorFormatAttr::E2M1 + || tensor.role != role + || tensor.layout != CuteTensorLayoutAttr::KMajor + || tensor.alignment.0 < 16 + || tensor + .storage + .deref(ctx) + .downcast_ref::() + .is_none() + || view.placement != plan.shared_layout + || view.storage_elements() != Some(4096) + || view.storage_bytes(ctx) != Some(8192) + { + return Err("must be the matching 128x128 packed-E2M1 shared tile"); + } + Ok(()) +} + +/// Require the scale format, operand role, and canonical stage placement needed by tiled MMA. +fn check_scale_view( + ctx: &Context, + view: &CuteSmemTensorType, + plan: &CuteTiledMmaPlanAttr, + role: CuteTensorRoleAttr, +) -> Result<(), &'static str> { + view.verify(ctx) + .map_err(|_| "uses an invalid shared tensor type")?; + let tensor = view.tensor_view(ctx).ok_or("must wrap a tensor view")?; + let storage_is_u32 = tensor + .storage + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.is_unsigned()); + if tensor.format != plan.atom.scale_format + || tensor.role != role + || tensor.layout != CuteTensorLayoutAttr::BlockScaleKMajor(plan.atom.values_per_scale) + || tensor.alignment.0 < 4 + || !storage_is_u32 + || view.placement.0 != canonical_scale_placement() + || view.storage_elements() != Some(SM1XX_BLOCK_SCALE_ATOM_BYTES / 4) + || view.storage_bytes(ctx) != Some(SM1XX_BLOCK_SCALE_ATOM_BYTES) + { + return Err("must be the matching 128-row canonical UE8M0/E4M3 shared scale atom"); + } + Ok(()) +} + +/// Require the ordinary aggregate's register counts to match its semantic fragment kind. +fn check_carrier( + ctx: &Context, + value: Value, + expected: MmaCarrierShape, +) -> Result<(), &'static str> { + if mma_carrier_shape(ctx, value.get_type(ctx)) == Some(expected) { + Ok(()) + } else { + Err("has the wrong ordinary MIR carrier shape") + } +} + +mod fragment; +mod gemm; +mod ldmatrix; +mod load; +mod tensor; + +pub use fragment::{CuteFragmentFillOp, CuteFragmentSliceKOp}; +pub use gemm::CuteTiledGemmOp; +pub use ldmatrix::CuteLdmatrixOp; +pub use load::{CuteMmaLoadAOp, CuteMmaLoadScalesOp, CuteMmaPartitionBOp}; +pub use tensor::{CuteSmemTensorOverlayOp, CuteTiledMmaSliceOp}; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tensor.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tensor.rs new file mode 100644 index 0000000000..032d3ffe18 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tensor.rs @@ -0,0 +1,196 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared tensor overlays and per-lane tiled MMA selection. + +use super::*; + +/// Give a shared pointer/capacity pair typed tensor meaning without changing +/// either runtime value. +/// +/// The `smem_overlay_view` attribute is a `!cute.smem_tensor` type. Results +/// have exactly the operand types, so backend lowering may erase this +/// operation to two identity replacements. +#[pliron_op( + name = "cute.smem_tensor_overlay", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<2>], + attributes = (smem_overlay_view: TypeAttr) +)] +pub struct CuteSmemTensorOverlayOp; + +impl CuteSmemTensorOverlayOp { + /// Build a detached overlay annotating ordinary shared pointer/capacity carriers. + pub fn new(ctx: &mut Context, base: Value, capacity: Value, view: TypeHandle) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![base.get_type(ctx), capacity.get_type(ctx)], + vec![base, capacity], + vec![], + 0, + ), + }; + operation.set_attr_smem_overlay_view(ctx, TypeAttr::new(view)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the original shared pointer supplied to the semantic overlay. + #[must_use] + pub fn input_base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the original shared element-capacity operand. + #[must_use] + pub fn input_capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the shared storage pointer retained by this semantic operation. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the ordinary unsigned storage-capacity carrier, in elements. + #[must_use] + pub fn capacity(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Static logical/storage/role/layout facts attached to this pair. + #[must_use] + pub fn view(&self, ctx: &Context) -> Option { + let attr = self.get_attr_smem_overlay_view(ctx)?; + smem_tensor_of_type(ctx, attr.get_type(ctx)) + } +} + +impl Verify for CuteSmemTensorOverlayOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 2 { + return verify_err!( + op.loc(), + "cute.smem_tensor_overlay needs 2 operands and 2 results" + ); + } + let Some(view) = self.view(ctx) else { + return verify_err!( + op.loc(), + "cute.smem_tensor_overlay must carry a cute.smem_tensor type" + ); + }; + view.verify(ctx)?; + let tensor = view.tensor_view(ctx).expect("verified shared tensor"); + if !shared_pointer_matches(ctx, self.input_base(ctx), tensor.storage) { + return verify_err!( + op.loc(), + "cute.smem_tensor_overlay base must be a mutable CTA-shared storage pointer" + ); + } + if !is_u64(ctx, self.input_capacity(ctx)) { + return verify_err!( + op.loc(), + "cute.smem_tensor_overlay capacity must be an unsigned 64-bit integer" + ); + } + if self.base(ctx).get_type(ctx) != self.input_base(ctx).get_type(ctx) + || self.capacity(ctx).get_type(ctx) != self.input_capacity(ctx).get_type(ctx) + { + return verify_err!( + op.loc(), + "cute.smem_tensor_overlay results must preserve the pointer and capacity types" + ); + } + Ok(()) + } +} + +/// Attach one tiled-MMA plan to the calling lane number. +/// +/// The result remains an ordinary `u32`; it may safely cross loop edges. +#[pliron_op( + name = "cute.tiled_mma_slice", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>, OneResultInterface], + attributes = (tiled_mma_slice_plan: CuteTiledMmaPlanAttr) +)] +pub struct CuteTiledMmaSliceOp; + +impl CuteTiledMmaSliceOp { + /// Build a detached lane selection carrying the static tiled-MMA ownership plan. + pub fn new(ctx: &mut Context, lane: Value, plan: CuteTiledMmaPlanAttr) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![lane.get_type(ctx)], + vec![lane], + vec![], + 0, + ), + }; + operation.set_attr_tiled_mma_slice_plan(ctx, plan); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the lane-index carrier identifying register ownership within a warp. + #[must_use] + pub fn lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the forwarded lane result annotated by the tiled-MMA plan. + #[must_use] + pub fn sliced_lane(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Clone the operation's typed tiled-MMA plan, if present. + #[must_use] + pub fn plan(&self, ctx: &Context) -> Option { + plan_of_attr(self.get_attr_tiled_mma_slice_plan(ctx)) + } +} + +impl Verify for CuteTiledMmaSliceOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.tiled_mma_slice needs 1 operand and 1 result" + ); + } + if !is_u32(ctx, self.lane(ctx)) + || self.sliced_lane(ctx).get_type(ctx) != self.lane(ctx).get_type(ctx) + { + return verify_err!( + op.loc(), + "cute.tiled_mma_slice must preserve one unsigned u32 lane" + ); + } + let Some(plan) = self.plan(ctx) else { + return verify_err!(op.loc(), "cute.tiled_mma_slice must carry a plan"); + }; + if let Err(message) = check_plan(&plan, ctx) { + return verify_err!(op.loc(), "cute.tiled_mma_slice {message}"); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tests.rs new file mode 100644 index 0000000000..fa370a6f51 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mma/tests.rs @@ -0,0 +1,507 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use crate::attributes::{CuteTensorAccessAttr, CuteTensorAddressSpaceAttr}; +use crate::types::CuteTensorViewType; +use dialect_mir::ops::MirUndefOp; +use pliron::builtin::types::Signedness; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Build the canonical swizzled packed-data shared layout for the MMA fixture. +fn data_placement() -> ComposedLayout { + let outer: Layout = "(128,32):(32,1)".parse().unwrap(); + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, outer, OffsetUnit::Elements).unwrap() +} + +/// Storage interpretation and placement used to construct shared-view test types. +struct SmemTypeFacts { + /// Logical interpretation of the bits in each physical storage element. + format: CuteTensorFormatAttr, + /// Logical M/N coordinate role assigned to the tensor. + role: CuteTensorRoleAttr, + /// Promised storage-base alignment, in bytes. + alignment: u64, + /// Static coordinate layout used to interpret the storage. + layout: CuteTensorLayoutAttr, + /// Mapping from logical coordinates to physical shared storage offsets. + placement: ComposedLayout, +} + +/// Construct a typed shared view from independently configurable storage facts for positive and negative cases. +fn smem_type( + ctx: &mut Context, + logical: TypeHandle, + storage: TypeHandle, + facts: SmemTypeFacts, +) -> TypeHandle { + let tensor: TypeHandle = CuteTensorViewType::get_with_facts( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + facts.alignment, + facts.format, + facts.role, + facts.layout, + ) + .into(); + CuteSmemTensorType::get(ctx, tensor, facts.placement).into() +} + +/// A mutually compatible set of shared views, register carriers, and one MMA plan. +struct TestTypes { + /// Unsigned 32-bit index/register type. + u32_ty: TypeHandle, + /// Unsigned 64-bit extent/barrier type. + u64_ty: TypeHandle, + /// FP32 accumulator element type. + f32_ty: TypeHandle, + /// Pointer type for A's shared packed-value storage. + a_ptr_ty: TypeHandle, + /// Pointer type for B's shared packed-value storage. + b_ptr_ty: TypeHandle, + /// Pointer type for A's shared scale carriers. + scale_a_ptr_ty: TypeHandle, + /// Pointer type for B's shared scale carriers. + scale_b_ptr_ty: TypeHandle, + /// Semantic shared-view type for the A operand. + a_view_ty: TypeHandle, + /// Semantic shared-view type for the B operand. + b_view_ty: TypeHandle, + /// Semantic shared-view type for A scales. + scale_a_view_ty: TypeHandle, + /// Semantic shared-view type for B scales. + scale_b_view_ty: TypeHandle, + /// Ordinary aggregate type holding packed A registers. + a_carrier_ty: TypeHandle, + /// Ordinary aggregate type holding a stage of scale registers. + scale_carrier_ty: TypeHandle, + /// Ordinary aggregate type holding the FP32 accumulator registers. + accumulator_ty: TypeHandle, + /// Static layout and instruction contract used by this fixture or type. + plan: CuteTiledMmaPlanAttr, +} + +/// Build the compatible shared view types, register shapes, and tiled plan for compute tests. +fn setup() -> (Context, TestTypes) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(&ctx).into(); + let a_ptr_ty: TypeHandle = + MirPtrType::get(&mut ctx, f16_ty, true, address_space::SHARED).into(); + let b_ptr_ty = a_ptr_ty; + let scale_a_ptr_ty: TypeHandle = + MirPtrType::get(&mut ctx, u32_ty, true, address_space::SHARED).into(); + let scale_b_ptr_ty = scale_a_ptr_ty; + let plan = CuteTiledMmaPlanAttr::mxf4_128x128x128(data_placement()); + let a_view_ty = smem_type( + &mut ctx, + f32_ty, + f16_ty, + SmemTypeFacts { + format: CuteTensorFormatAttr::E2M1, + role: CuteTensorRoleAttr::Mkl, + alignment: 16, + layout: CuteTensorLayoutAttr::KMajor, + placement: data_placement(), + }, + ); + let b_view_ty = smem_type( + &mut ctx, + f32_ty, + f16_ty, + SmemTypeFacts { + format: CuteTensorFormatAttr::E2M1, + role: CuteTensorRoleAttr::Nkl, + alignment: 16, + layout: CuteTensorLayoutAttr::KMajor, + placement: data_placement(), + }, + ); + let scale_a_view_ty = smem_type( + &mut ctx, + f32_ty, + u32_ty, + SmemTypeFacts { + format: CuteTensorFormatAttr::UE8M0, + role: CuteTensorRoleAttr::Mkl, + alignment: 4, + layout: CuteTensorLayoutAttr::BlockScaleKMajor(32), + placement: canonical_scale_placement(), + }, + ); + let scale_b_view_ty = smem_type( + &mut ctx, + f32_ty, + u32_ty, + SmemTypeFacts { + format: CuteTensorFormatAttr::UE8M0, + role: CuteTensorRoleAttr::Nkl, + alignment: 4, + layout: CuteTensorLayoutAttr::BlockScaleKMajor(32), + placement: canonical_scale_placement(), + }, + ); + let a_carrier_ty: TypeHandle = + dialect_mir::types::MirArrayType::get(&mut ctx, u32_ty, 8).into(); + let scale_carrier_ty: TypeHandle = + dialect_mir::types::MirArrayType::get(&mut ctx, u32_ty, 10).into(); + let accumulator_ty: TypeHandle = + dialect_mir::types::MirArrayType::get(&mut ctx, f32_ty, 64).into(); + ( + ctx, + TestTypes { + u32_ty, + u64_ty, + f32_ty, + a_ptr_ty, + b_ptr_ty, + scale_a_ptr_ty, + scale_b_ptr_ty, + a_view_ty, + b_view_ty, + scale_a_view_ty, + scale_b_view_ty, + a_carrier_ty, + scale_carrier_ty, + accumulator_ty, + plan, + }, + ) +} + +/// Attach semantic shared-tensor metadata to fixture pointer/capacity carriers. +fn overlay( + ctx: &mut Context, + base_ty: TypeHandle, + capacity_ty: TypeHandle, + view_ty: TypeHandle, +) -> CuteSmemTensorOverlayOp { + let base = undef(ctx, base_ty); + let capacity = undef(ctx, capacity_ty); + CuteSmemTensorOverlayOp::new(ctx, base, capacity, view_ty) +} + +/// Compute operations retain Rust-shaped carriers and defer B loads until the tiled GEMM consumes them. +#[test] +fn visual_chain_keeps_carriers_ordinary_and_b_lazy() { + let (mut ctx, types) = setup(); + assert!(types.plan.verify(&ctx).is_ok()); + assert_eq!(types.plan.compute_warps(), Some(8)); + assert_eq!(types.plan.scale_words_per_lane(), Some(10)); + assert_eq!(types.plan.a_registers_per_lane(), Some(8)); + assert_eq!(types.plan.accumulator_registers_per_lane(), Some(64)); + + let a = overlay(&mut ctx, types.a_ptr_ty, types.u64_ty, types.a_view_ty); + let b = overlay(&mut ctx, types.b_ptr_ty, types.u64_ty, types.b_view_ty); + let scale_a = overlay( + &mut ctx, + types.scale_a_ptr_ty, + types.u64_ty, + types.scale_a_view_ty, + ); + let scale_b = overlay( + &mut ctx, + types.scale_b_ptr_ty, + types.u64_ty, + types.scale_b_view_ty, + ); + for operation in [&a, &b, &scale_a, &scale_b] { + assert!(operation.verify(&ctx).is_ok()); + } + + let lane = undef(&mut ctx, types.u32_ty); + let warp_m = undef(&mut ctx, types.u64_ty); + let warp_n = undef(&mut ctx, types.u64_ty); + let k_half = undef(&mut ctx, types.u64_ty); + let mma = CuteTiledMmaSliceOp::new(&mut ctx, lane, types.plan.clone()); + assert!(mma.verify(&ctx).is_ok()); + + let zero = undef(&mut ctx, types.f32_ty); + let fill = CuteFragmentFillOp::new( + &mut ctx, + zero, + types.accumulator_ty, + types.plan.clone(), + CuteMmaCarrierKindAttr::Accumulator, + ); + assert!(fill.verify(&ctx).is_ok()); + + let sliced_lane = mma.sliced_lane(&ctx); + let scale_a_base = scale_a.base(&ctx); + let scale_a_capacity = scale_a.capacity(&ctx); + let scale_b_base = scale_b.base(&ctx); + let scale_b_capacity = scale_b.capacity(&ctx); + let load_scales = CuteMmaLoadScalesOp::new( + &mut ctx, + sliced_lane, + scale_a_base, + scale_a_capacity, + scale_b_base, + scale_b_capacity, + warp_m, + warp_n, + types.scale_a_view_ty, + types.scale_b_view_ty, + types.scale_carrier_ty, + types.plan.clone(), + ); + assert!(load_scales.verify(&ctx).is_ok()); + let stage_scales = load_scales.scales(&ctx); + let slice_scales = CuteFragmentSliceKOp::new( + &mut ctx, + stage_scales, + k_half, + types.scale_carrier_ty, + types.plan.clone(), + ); + assert!(slice_scales.verify(&ctx).is_ok()); + + let a_base = a.base(&ctx); + let a_capacity = a.capacity(&ctx); + let load_a = CuteMmaLoadAOp::new( + &mut ctx, + sliced_lane, + a_base, + a_capacity, + warp_m, + k_half, + types.a_view_ty, + types.a_carrier_ty, + types.plan.clone(), + ); + assert!(load_a.verify(&ctx).is_ok()); + + let b_base = b.base(&ctx); + let b_capacity = b.capacity(&ctx); + let partition_b = CuteMmaPartitionBOp::new( + &mut ctx, + b_base, + b_capacity, + warp_n, + k_half, + types.b_view_ty, + types.plan.clone(), + ); + assert!(partition_b.verify(&ctx).is_ok()); + { + let partition_operation = partition_b.get_operation().deref(&ctx); + assert_eq!(partition_operation.get_num_results(), 4); + assert!((0..4).all(|index| { + let value = partition_operation.get_result(index); + mma_carrier_shape(&ctx, value.get_type(&ctx)).is_none() + })); + } + + let a_fragment = load_a.fragment(&ctx); + let partition_base = partition_b.base(&ctx); + let partition_capacity = partition_b.capacity(&ctx); + let partition_warp_n = partition_b.warp_n(&ctx); + let partition_k_half = partition_b.k_half(&ctx); + let selected_scales = slice_scales.scales(&ctx); + let initial_accumulator = fill.fragment(&ctx); + let gemm = CuteTiledGemmOp::new( + &mut ctx, + sliced_lane, + a_fragment, + partition_base, + partition_capacity, + partition_warp_n, + partition_k_half, + selected_scales, + initial_accumulator, + types.b_view_ty, + types.plan, + ); + assert!(gemm.verify(&ctx).is_ok()); + assert_eq!( + mma_carrier_shape(&ctx, gemm.accumulator(&ctx).get_type(&ctx)), + Some(MmaCarrierShape::f32(64)) + ); +} + +/// The compute ABI requires a typed lazy B partition, preventing an eager register aggregate from replacing it. +#[test] +fn tiled_gemm_rejects_an_eager_or_untyped_b_fragment() { + let (mut ctx, types) = setup(); + let lane = undef(&mut ctx, types.u32_ty); + let a = undef(&mut ctx, types.a_carrier_ty); + let eager_b = undef(&mut ctx, types.a_carrier_ty); + let capacity = undef(&mut ctx, types.u64_ty); + let warp_n = undef(&mut ctx, types.u64_ty); + let k_half = undef(&mut ctx, types.u64_ty); + let scales = undef(&mut ctx, types.scale_carrier_ty); + let accumulator = undef(&mut ctx, types.accumulator_ty); + let gemm = CuteTiledGemmOp::new( + &mut ctx, + lane, + a, + eager_b, + capacity, + warp_n, + k_half, + scales, + accumulator, + types.b_view_ty, + types.plan, + ); + assert!( + gemm.verify(&ctx) + .unwrap_err() + .to_string() + .contains("lazy B") + ); +} + +/// Static view metadata lets flattened pointer/capacity carriers retain their meaning after ordinary aggregate forwarding. +#[test] +fn typed_consumers_accept_flattened_carriers_after_aggregate_forwarding() { + let (mut ctx, types) = setup(); + let lane = undef(&mut ctx, types.u32_ty); + let warp_m = undef(&mut ctx, types.u64_ty); + let warp_n = undef(&mut ctx, types.u64_ty); + let k_half = undef(&mut ctx, types.u64_ty); + let a_base = undef(&mut ctx, types.a_ptr_ty); + let a_capacity = undef(&mut ctx, types.u64_ty); + let load_a = CuteMmaLoadAOp::new( + &mut ctx, + lane, + a_base, + a_capacity, + warp_m, + k_half, + types.a_view_ty, + types.a_carrier_ty, + types.plan.clone(), + ); + assert!(load_a.verify(&ctx).is_ok()); + + let a_fragment = load_a.fragment(&ctx); + let b_base = undef(&mut ctx, types.b_ptr_ty); + let b_capacity = undef(&mut ctx, types.u64_ty); + let scales = undef(&mut ctx, types.scale_carrier_ty); + let accumulator = undef(&mut ctx, types.accumulator_ty); + let gemm = CuteTiledGemmOp::new( + &mut ctx, + lane, + a_fragment, + b_base, + b_capacity, + warp_n, + k_half, + scales, + accumulator, + types.b_view_ty, + types.plan, + ); + assert!(gemm.verify(&ctx).is_ok()); +} + +/// A/B scale views cannot exchange coordinate roles even when their storage shapes match. +#[test] +fn scale_load_rejects_swapped_m_and_n_roles() { + let (mut ctx, types) = setup(); + let scale_a_wrong = overlay( + &mut ctx, + types.scale_a_ptr_ty, + types.u64_ty, + types.scale_b_view_ty, + ); + let scale_b = overlay( + &mut ctx, + types.scale_b_ptr_ty, + types.u64_ty, + types.scale_b_view_ty, + ); + let lane = undef(&mut ctx, types.u32_ty); + let warp_m = undef(&mut ctx, types.u64_ty); + let warp_n = undef(&mut ctx, types.u64_ty); + let scale_a_base = scale_a_wrong.base(&ctx); + let scale_a_capacity = scale_a_wrong.capacity(&ctx); + let scale_b_base = scale_b.base(&ctx); + let scale_b_capacity = scale_b.capacity(&ctx); + let load = CuteMmaLoadScalesOp::new( + &mut ctx, + lane, + scale_a_base, + scale_a_capacity, + scale_b_base, + scale_b_capacity, + warp_m, + warp_n, + types.scale_b_view_ty, + types.scale_b_view_ty, + types.scale_carrier_ty, + types.plan, + ); + assert!( + load.verify(&ctx) + .unwrap_err() + .to_string() + .contains("A must be") + ); +} + +/// Unsupported tiled plans and incorrect register aggregate sizes fail before backend selection. +#[test] +fn plan_and_carrier_verifiers_fail_closed() { + let (mut ctx, types) = setup(); + let mut bad_plan = types.plan.clone(); + bad_plan.b_load_group = 3; + assert!( + bad_plan + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("B load group") + ); + + let mut wrong_ownership = types.plan.clone(); + wrong_ownership.m_ownership = crate::attributes::CuteLayoutAttr( + "(4,2):(2,1)".parse().expect("dense alternate map is valid"), + ); + assert!(wrong_ownership.verify(&ctx).is_ok()); + let lane = undef(&mut ctx, types.u32_ty); + let slice = CuteTiledMmaSliceOp::new(&mut ctx, lane, wrong_ownership); + assert!( + slice + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("ownership maps") + ); + + let wrong_accumulator_ty: TypeHandle = + dialect_mir::types::MirArrayType::get(&mut ctx, types.f32_ty, 63).into(); + let zero = undef(&mut ctx, types.f32_ty); + let fill = CuteFragmentFillOp::new( + &mut ctx, + zero, + wrong_accumulator_ty, + types.plan, + CuteMmaCarrierKindAttr::Accumulator, + ); + assert!( + fill.verify(&ctx) + .unwrap_err() + .to_string() + .contains("carrier shape") + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mod.rs new file mode 100644 index 0000000000..a622dc1525 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/mod.rs @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CuTe operation families grouped by their semantic responsibility. +//! +//! Use the family modules for tensor, copy, GEMV, MMA, pipeline, scheduler, +//! and epilogue operations. The direct copy, divisibility, matrix-load, and +//! TMA operation names remain available here for existing callers. + +pub mod copy; +mod divisibility; +pub mod epilogue; +pub mod gemm_tma; +pub mod gemv; +pub mod mma; +pub mod pipeline; +pub mod scheduler; +pub mod sm100; +pub mod sm100_attention; +pub mod tensor; + +pub use copy::{CuteCopyG2SOp, CuteCopyOp, CuteTmaLoad2dOp, CuteTmaStore2dOp}; +pub use divisibility::CuteAssumeDivOp; +pub use mma::CuteLdmatrixOp; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/construction.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/construction.rs new file mode 100644 index 0000000000..c467c053fc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/construction.rs @@ -0,0 +1,196 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Load-pipeline handles and barrier-ring initialization. + +use super::*; + +/// Bind one shared-memory barrier allocation to a TMA load-pipeline config. +/// +/// The base points at this exact physical order: +/// +/// ```text +/// [full 0 ... full S-1][empty 0 ... empty S-1] +/// ``` +/// +/// MIR represents CUDA's 8-byte `Barrier` storage as one unsigned `u64`. +#[pliron_op( + name = "cute.tma_load_pipeline_make", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>, OneResultInterface], + attributes = (pipeline_base_alignment_bytes: CuteAlignmentAttr) +)] +pub struct CuteTmaLoadPipelineMakeOp; + +impl CuteTmaLoadPipelineMakeOp { + /// Create a compiler-only view over an existing shared barrier ring. + pub fn new( + ctx: &mut Context, + base: Value, + stages: u64, + consumer_warps: u32, + transaction_bytes: u32, + base_alignment_bytes: u64, + ) -> Self { + let pipeline: TypeHandle = + CuteTmaLoadPipelineType::get(ctx, stages, consumer_warps, transaction_bytes).into(); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![pipeline], + vec![base], + vec![], + 0, + ), + }; + operation + .set_attr_pipeline_base_alignment_bytes(ctx, CuteAlignmentAttr(base_alignment_bytes)); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Shared base of the first full barrier. + #[must_use] + pub fn base(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Compiler-only pipeline handle. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Alignment promised for the start of the whole barrier ring. + #[must_use] + pub fn promised_base_alignment(&self, ctx: &Context) -> Option { + self.get_attr_pipeline_base_alignment_bytes(ctx) + .map(|alignment| alignment.0) + } + + /// Static pipeline configuration carried by the result type. + #[must_use] + pub fn pipeline_type(&self, ctx: &Context) -> Option { + pipeline_type_of(ctx, self.pipeline(ctx)) + } +} + +impl Verify for CuteTmaLoadPipelineMakeOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_make needs 1 operand and 1 result" + ); + } + if !valid_shared_barrier_pointer(ctx, self.base(ctx)) { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_make base must be a mutable shared pointer to the canonical unsigned u64 Barrier storage" + ); + } + let Some(pipeline) = self.pipeline_type(ctx) else { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_make result must be a TMA load-pipeline handle" + ); + }; + pipeline.verify(ctx)?; + let Some(alignment) = self.promised_base_alignment(ctx) else { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_make must carry a base-alignment promise" + ); + }; + CuteAlignmentAttr(alignment).verify(ctx)?; + if alignment < BARRIER_BYTES { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_make base alignment must be at least 8 bytes" + ); + } + if self + .pipeline(ctx) + .uses(ctx) + .into_iter() + .any(|r#use| !allowed_pipeline_user(ctx, r#use.user_op())) + { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline handle may only feed load-pipeline lifecycle operations" + ); + } + Ok(()) + } +} + +/// Initialize all full/empty barriers once, then publish them to the CTA. +#[pliron_op( + name = "cute.tma_load_pipeline_init", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<0>] +)] +pub struct CuteTmaLoadPipelineInitOp; + +impl CuteTmaLoadPipelineInitOp { + /// Build a detached initialization/publication operation for both barrier rings. + pub fn new(ctx: &mut Context, pipeline: Value, init_thread: Value) -> Self { + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![pipeline, init_thread], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the direct compiler-only load-pipeline handle operand. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Thread index that initializes every barrier. All threads synchronize. + #[must_use] + pub fn init_thread(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } +} + +impl Verify for CuteTmaLoadPipelineInitOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_init needs 2 operands and 0 results" + ); + } + if let Err(message) = checked_pipeline(ctx, self.pipeline(ctx)) { + return verify_err!(op.loc(), "cute.tma_load_pipeline_init {message}"); + } + if !is_u32(ctx, self.init_thread(ctx)) { + return verify_err!( + op.loc(), + "cute.tma_load_pipeline_init thread index must be an unsigned 32-bit integer" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/lifecycle.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/lifecycle.rs new file mode 100644 index 0000000000..d195fe25c9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/lifecycle.rs @@ -0,0 +1,433 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Producer and consumer barrier-ring lifecycle operations. + +use super::*; + +/// Generate the common handle, slot, and static cursor-state accessors. +macro_rules! pipeline_state_accessors { + ($attr_getter:ident) => { + /// Return the direct compiler-only load-pipeline handle operand. + #[must_use] + pub fn pipeline(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the ordinary u32 slot selecting the current barrier stage. + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Clone the static producer/consumer role and ring size, if present. + #[must_use] + pub fn state(&self, ctx: &Context) -> Option { + self.$attr_getter(ctx).map(|state| *state) + } + }; +} + +/// Wait until the producer may reuse its current empty stage. +#[pliron_op( + name = "cute.pipeline_producer_acquire", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<0>], + attributes = (pipeline_producer_acquire_state: CutePipelineStateAttr) +)] +pub struct CutePipelineProducerAcquireOp; + +impl CutePipelineProducerAcquireOp { + /// Build a detached wait for the current empty stage and producer parity. + pub fn new( + ctx: &mut Context, + pipeline: Value, + slot: Value, + phase: Value, + state: CutePipelineStateAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![pipeline, slot, phase], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_producer_acquire_state(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_producer_acquire_state); + + /// Return the ordinary u32 barrier-parity carrier used by this operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } +} + +impl Verify for CutePipelineProducerAcquireOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.pipeline_producer_acquire needs 3 operands and 0 results" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_producer_acquire must carry producer state facts" + ); + }; + if let Err(message) = checked_pipeline_state( + ctx, + self.pipeline(ctx), + state, + CutePipelineRoleAttr::Producer, + ) { + return verify_err!(op.loc(), "cute.pipeline_producer_acquire {message}"); + } + if !is_u32(ctx, self.slot(ctx)) || !is_u32(ctx, self.phase(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_producer_acquire slot and phase must be unsigned 32-bit integers" + ); + } + Ok(()) + } +} + +/// Set the expected byte count and return this stage's full barrier pointer. +/// +/// The returned pointer is real SSA. The four existing TMA-copy operations +/// consume it directly, so their ABI does not change. +#[pliron_op( + name = "cute.pipeline_producer_expect_tx", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface], + attributes = (pipeline_producer_expect_state: CutePipelineStateAttr) +)] +pub struct CutePipelineProducerExpectTxOp; + +impl CutePipelineProducerExpectTxOp { + /// Build a detached byte-expectation operation returning the stage's full barrier. + pub fn new( + ctx: &mut Context, + pipeline: Value, + slot: Value, + state: CutePipelineStateAttr, + ) -> Self { + let make = pipeline_make_of(ctx, pipeline) + .expect("cute.pipeline_producer_expect_tx builder needs a direct pipeline make"); + let barrier_pointer = make.base(ctx).get_type(ctx); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![barrier_pointer], + vec![pipeline, slot], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_producer_expect_state(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_producer_expect_state); + + /// Mutable shared pointer consumed by the existing four TMA copies. + #[must_use] + pub fn completion_barrier(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } +} + +impl Verify for CutePipelineProducerExpectTxOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.pipeline_producer_expect_tx needs 2 operands and 1 result" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_producer_expect_tx must carry producer state facts" + ); + }; + let (_, make) = match checked_pipeline_state( + ctx, + self.pipeline(ctx), + state, + CutePipelineRoleAttr::Producer, + ) { + Ok(pair) => pair, + Err(message) => { + return verify_err!(op.loc(), "cute.pipeline_producer_expect_tx {message}"); + } + }; + if !is_u32(ctx, self.slot(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_producer_expect_tx slot must be an unsigned 32-bit integer" + ); + } + if self.completion_barrier(ctx).get_type(ctx) != make.base(ctx).get_type(ctx) { + return verify_err!( + op.loc(), + "cute.pipeline_producer_expect_tx result must match the pipeline's shared Barrier pointer type" + ); + } + Ok(()) + } +} + +/// Wait until all expected TMA bytes have reached the current full stage. +#[pliron_op( + name = "cute.pipeline_consumer_wait", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<0>], + attributes = (pipeline_consumer_wait_state: CutePipelineStateAttr) +)] +pub struct CutePipelineConsumerWaitOp; + +impl CutePipelineConsumerWaitOp { + /// Build a detached wait for all expected input bytes at the consumer cursor. + pub fn new( + ctx: &mut Context, + pipeline: Value, + slot: Value, + phase: Value, + state: CutePipelineStateAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![pipeline, slot, phase], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_consumer_wait_state(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_consumer_wait_state); + + /// Return the ordinary u32 barrier-parity carrier used by this operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } +} + +impl Verify for CutePipelineConsumerWaitOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_wait needs 3 operands and 0 results" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_wait must carry consumer state facts" + ); + }; + if let Err(message) = checked_pipeline_state( + ctx, + self.pipeline(ctx), + state, + CutePipelineRoleAttr::Consumer, + ) { + return verify_err!(op.loc(), "cute.pipeline_consumer_wait {message}"); + } + if !is_u32(ctx, self.slot(ctx)) || !is_u32(ctx, self.phase(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_wait slot and phase must be unsigned 32-bit integers" + ); + } + Ok(()) + } +} + +/// Mark the current stage empty after all consumer warps finish reading it. +#[pliron_op( + name = "cute.pipeline_consumer_release", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<0>], + attributes = (pipeline_consumer_release_state: CutePipelineStateAttr) +)] +pub struct CutePipelineConsumerReleaseOp; + +impl CutePipelineConsumerReleaseOp { + /// Build a detached empty-stage arrival after consumer reads finish. + pub fn new( + ctx: &mut Context, + pipeline: Value, + slot: Value, + state: CutePipelineStateAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![pipeline, slot], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_consumer_release_state(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_consumer_release_state); +} + +impl Verify for CutePipelineConsumerReleaseOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_release needs 2 operands and 0 results" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_release must carry consumer state facts" + ); + }; + if let Err(message) = checked_pipeline_state( + ctx, + self.pipeline(ctx), + state, + CutePipelineRoleAttr::Consumer, + ) { + return verify_err!(op.loc(), "cute.pipeline_consumer_release {message}"); + } + if !is_u32(ctx, self.slot(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_consumer_release slot must be an unsigned 32-bit integer" + ); + } + Ok(()) + } +} + +/// Before the producer exits, wait until every stage is empty again. +#[pliron_op( + name = "cute.pipeline_producer_tail", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<0>], + attributes = (pipeline_producer_tail_state: CutePipelineStateAttr) +)] +pub struct CutePipelineProducerTailOp; + +impl CutePipelineProducerTailOp { + /// Build a detached exit wait draining the producer's reusable-stage ring. + pub fn new( + ctx: &mut Context, + pipeline: Value, + slot: Value, + phase: Value, + state: CutePipelineStateAttr, + ) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![pipeline, slot, phase], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_producer_tail_state(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + pipeline_state_accessors!(get_attr_pipeline_producer_tail_state); + + /// Return the ordinary u32 barrier-parity carrier used by this operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(2) + } +} + +impl Verify for CutePipelineProducerTailOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.pipeline_producer_tail needs 3 operands and 0 results" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_producer_tail must carry producer state facts" + ); + }; + if let Err(message) = checked_pipeline_state( + ctx, + self.pipeline(ctx), + state, + CutePipelineRoleAttr::Producer, + ) { + return verify_err!(op.loc(), "cute.pipeline_producer_tail {message}"); + } + if !is_u32(ctx, self.slot(ctx)) || !is_u32(ctx, self.phase(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_producer_tail slot and phase must be unsigned 32-bit integers" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/mod.rs new file mode 100644 index 0000000000..2f8dd90568 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/mod.rs @@ -0,0 +1,193 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! A small semantic layer for the GEMM TMA load pipeline. +//! +//! The pipeline owns two hardware-barrier rings: +//! +//! ```text +//! one slot +//! producer: wait empty ──► expect bytes ──► four TMA copies +//! │ +//! consumer: release empty ◄── read tile ◄── wait full +//! ``` +//! +//! Only [`CuteTmaLoadPipelineType`] is a compiler-only handle. The loop still +//! carries the same ordinary scalar state as Rust: +//! +//! ```text +//! slot: u32, phase: u32 ── advance<3> ──► next_slot, next_phase +//! ``` +//! +//! A backend continuation maps these operations to its barrier, transaction, +//! and publication primitives. It does not change the four TMA copies or add +//! a runtime pipeline object. +//! +//! A pipeline handle must come directly from `tma_load_pipeline_make`. It may +//! only feed the lifecycle operations in this module. It cannot enter a +//! function signature, pointer, aggregate, load/store, or block argument. + +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface, OneResultInterface}, + types::{IntegerType, Signedness}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{CuteAlignmentAttr, CutePipelineRoleAttr, CutePipelineStateAttr}; +use crate::types::CuteTmaLoadPipelineType; + +/// Shared storage occupied by one hardware barrier, in bytes. +const BARRIER_BYTES: u64 = 8; + +/// Intern the unsigned 32-bit scalar type used for cursor state. +fn u32_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 32, Signedness::Unsigned).into() +} + +/// Intern the unsigned 64-bit scalar type used for extents and indices. +fn u64_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 64, Signedness::Unsigned).into() +} + +/// Check that the SSA value has an unsigned 32-bit integer type. +fn is_u32(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.is_unsigned()) +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Read the ring configuration from a compiler-only pipeline handle's type. +fn pipeline_type_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Resolve a handle only when defined directly by a load-pipeline constructor. +fn pipeline_make_of(ctx: &Context, value: Value) -> Option { + let producer = value.defining_op()?; + if Operation::get_opid(producer, ctx) != CuteTmaLoadPipelineMakeOp::get_opid_static() { + return None; + } + let make = CuteTmaLoadPipelineMakeOp::wrap(producer); + (make.pipeline(ctx) == value).then_some(make) +} + +/// Require a valid direct pipeline handle before checking a lifecycle operation. +fn checked_pipeline( + ctx: &Context, + value: Value, +) -> Result<(CuteTmaLoadPipelineType, CuteTmaLoadPipelineMakeOp), String> { + let pipeline = pipeline_type_of(ctx, value) + .ok_or_else(|| "pipeline operand must be a TMA load-pipeline handle".to_owned())?; + pipeline + .verify(ctx) + .map_err(|error| format!("pipeline type is invalid: {error}"))?; + let make = pipeline_make_of(ctx, value).ok_or_else(|| { + "pipeline handle must come directly from cute.tma_load_pipeline_make".to_owned() + })?; + Ok((pipeline, make)) +} + +/// Require a valid cursor role and stage count consistent with the pipeline. +fn checked_pipeline_state( + ctx: &Context, + pipeline_value: Value, + state: CutePipelineStateAttr, + required_role: CutePipelineRoleAttr, +) -> Result<(CuteTmaLoadPipelineType, CuteTmaLoadPipelineMakeOp), String> { + state + .verify(ctx) + .map_err(|error| format!("pipeline state is invalid: {error}"))?; + if state.role != required_role { + return Err(format!( + "pipeline state must have {required_role:?} role, got {:?}", + state.role + )); + } + let (pipeline, make) = checked_pipeline(ctx, pipeline_value)?; + if state.stages != pipeline.stages { + return Err(format!( + "pipeline state has {} stages, but its pipeline has {}", + state.stages, pipeline.stages + )); + } + Ok((pipeline, make)) +} + +/// Downcast an SSA value to its MIR pointer type for barrier checks. +fn barrier_pointer_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Require mutable shared storage containing unsigned 64-bit barrier words. +fn valid_shared_barrier_pointer(ctx: &Context, value: Value) -> bool { + let Some(pointer) = barrier_pointer_of(ctx, value) else { + return false; + }; + if !pointer.is_mutable || pointer.address_space != address_space::SHARED { + return false; + } + pointer + .pointee + .deref(ctx) + .downcast_ref::() + .is_some_and(|barrier| barrier.width() == 64 && barrier.is_unsigned()) +} + +/// Recognize lifecycle operations allowed to consume a direct pipeline handle. +fn allowed_pipeline_user(ctx: &Context, user: Ptr) -> bool { + let opid = Operation::get_opid(user, ctx); + [ + CuteTmaLoadPipelineInitOp::get_opid_static(), + CutePipelineProducerAcquireOp::get_opid_static(), + CutePipelineProducerExpectTxOp::get_opid_static(), + CutePipelineConsumerWaitOp::get_opid_static(), + CutePipelineConsumerReleaseOp::get_opid_static(), + CutePipelineProducerTailOp::get_opid_static(), + ] + .contains(&opid) +} + +mod construction; +mod lifecycle; +mod state; + +pub use construction::{CuteTmaLoadPipelineInitOp, CuteTmaLoadPipelineMakeOp}; +pub use lifecycle::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, +}; +pub use state::{CutePipelineStateAdvanceOp, CutePipelineStateNewOp, CutePipelineStateSlotOp}; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/state.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/state.rs new file mode 100644 index 0000000000..9a4aae54cb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/state.rs @@ -0,0 +1,252 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Ordinary scalar slot and phase state for a pipeline ring. + +use super::*; + +/// Create the initial ordinary `(slot, phase)` pair for one pipeline side. +#[pliron_op( + name = "cute.pipeline_state_new", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<2>], + attributes = (pipeline_state_new_config: CutePipelineStateAttr) +)] +pub struct CutePipelineStateNewOp; + +impl CutePipelineStateNewOp { + /// Build detached initial slot/phase results for the selected pipeline role. + pub fn new(ctx: &mut Context, state: CutePipelineStateAttr) -> Self { + let ty = u32_type(ctx); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ty, ty], + vec![], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_state_new_config(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the ordinary u32 ring-slot carrier used by this operation. + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the ordinary u32 barrier-parity carrier used by this operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Clone the static producer/consumer role and ring-size attribute, if present. + #[must_use] + pub fn state(&self, ctx: &Context) -> Option { + self.get_attr_pipeline_state_new_config(ctx) + .map(|state| *state) + } +} + +impl Verify for CutePipelineStateNewOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 0 || op.get_num_results() != 2 { + return verify_err!( + op.loc(), + "cute.pipeline_state_new needs 0 operands and 2 results" + ); + } + if !is_u32(ctx, self.slot(ctx)) || !is_u32(ctx, self.phase(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_state_new slot and phase must be unsigned 32-bit integers" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!(op.loc(), "cute.pipeline_state_new must carry state facts"); + }; + state.verify(ctx) + } +} + +/// Widen a `u32` slot to the `u64` index used for shared-memory addressing. +#[pliron_op( + name = "cute.pipeline_state_slot", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>, OneResultInterface], + attributes = (pipeline_state_slot_config: CutePipelineStateAttr) +)] +pub struct CutePipelineStateSlotOp; + +impl CutePipelineStateSlotOp { + /// Build a detached conversion from a u32 cursor slot to a u64 addressing index. + pub fn new(ctx: &mut Context, slot: Value, state: CutePipelineStateAttr) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![u64_type(ctx)], + vec![slot], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_state_slot_config(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the ordinary u32 ring-slot carrier used by this operation. + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Slot widened for pointer offset arithmetic. + #[must_use] + pub fn index(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Clone the static producer/consumer role and ring-size attribute, if present. + #[must_use] + pub fn state(&self, ctx: &Context) -> Option { + self.get_attr_pipeline_state_slot_config(ctx) + .map(|state| *state) + } +} + +impl Verify for CutePipelineStateSlotOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.pipeline_state_slot needs 1 operand and 1 result" + ); + } + if !is_u32(ctx, self.slot(ctx)) || !is_u64(ctx, self.index(ctx)) { + return verify_err!( + op.loc(), + "cute.pipeline_state_slot needs a u32 slot and returns a u64 index" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!(op.loc(), "cute.pipeline_state_slot must carry state facts"); + }; + state.verify(ctx) + } +} + +/// Move one ordinary `(slot, phase)` pair to the next ring position. +#[pliron_op( + name = "cute.pipeline_state_advance", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<2>], + attributes = (pipeline_state_advance_config: CutePipelineStateAttr) +)] +pub struct CutePipelineStateAdvanceOp; + +impl CutePipelineStateAdvanceOp { + /// Build a detached ring-cursor advance preserving the producer/consumer role. + pub fn new(ctx: &mut Context, slot: Value, phase: Value, state: CutePipelineStateAttr) -> Self { + let ty = u32_type(ctx); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ty, ty], + vec![slot, phase], + vec![], + 0, + ), + }; + operation.set_attr_pipeline_state_advance_config(ctx, state); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the ordinary u32 ring-slot carrier used by this operation. + #[must_use] + pub fn slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the ordinary u32 barrier-parity carrier used by this operation. + #[must_use] + pub fn phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the advanced ring-slot result, wrapping at the static stage count. + #[must_use] + pub fn next_slot(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the advanced parity result, toggled when the ring wraps. + #[must_use] + pub fn next_phase(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Clone the static producer/consumer role and ring-size attribute, if present. + #[must_use] + pub fn state(&self, ctx: &Context) -> Option { + self.get_attr_pipeline_state_advance_config(ctx) + .map(|state| *state) + } +} + +impl Verify for CutePipelineStateAdvanceOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 2 { + return verify_err!( + op.loc(), + "cute.pipeline_state_advance needs 2 operands and 2 results" + ); + } + if ![ + self.slot(ctx), + self.phase(ctx), + self.next_slot(ctx), + self.next_phase(ctx), + ] + .into_iter() + .all(|value| is_u32(ctx, value)) + { + return verify_err!( + op.loc(), + "cute.pipeline_state_advance slot and phase values must be unsigned 32-bit integers" + ); + } + let Some(state) = self.state(ctx) else { + return verify_err!( + op.loc(), + "cute.pipeline_state_advance must carry state facts" + ); + }; + state.verify(ctx) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/tests.rs new file mode 100644 index 0000000000..1530e62dff --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/pipeline/tests.rs @@ -0,0 +1,394 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use dialect_mir::ops::MirUndefOp; + +/// Stage-ring size used by the pipeline regression fixture. +const STAGES: u64 = 3; +/// Number of consuming warps in the pipeline regression fixture. +const CONSUMER_WARPS: u32 = 8; +/// Expected bytes per completed stage in the pipeline regression fixture. +const TRANSACTION_BYTES: u32 = 17_408; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Construct a mutable shared-pointer type with the fixture's chosen pointee. +fn shared_pointer(ctx: &mut Context, pointee: TypeHandle, mutable: bool) -> Value { + let pointer: TypeHandle = MirPtrType::get(ctx, pointee, mutable, address_space::SHARED).into(); + undef(ctx, pointer) +} + +/// Context and scalar/shared-pointer carriers for the load-pipeline verifier tests. +struct Setup { + /// IR context owning all handles used by this fixture. + ctx: Context, + /// Shared storage pointer used by this fixture. + base: Value, + /// Unsigned 32-bit cursor/phase fixture value. + u32_value: Value, + /// Unsigned 64-bit extent/index fixture value. + u64_value: Value, +} + +/// Create the barrier-ring pointer and correctly typed u32/u64 cursor carriers. +fn setup() -> Setup { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let barrier = u64_type(&ctx); + let base = shared_pointer(&mut ctx, barrier, true); + let u32_ty = u32_type(&ctx); + let u64_ty = u64_type(&ctx); + let u32_value = undef(&mut ctx, u32_ty); + let u64_value = undef(&mut ctx, u64_ty); + Setup { + ctx, + base, + u32_value, + u64_value, + } +} + +/// Creation, initialization, cursor updates, and producer/consumer effects retain their separate operation boundaries. +#[test] +fn load_pipeline_stays_a_composable_ring_protocol() { + let Setup { + mut ctx, + base, + u32_value, + .. + } = setup(); + let make = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + base, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + let pipeline = make.pipeline(&ctx); + + let init = CuteTmaLoadPipelineInitOp::new(&mut ctx, pipeline, u32_value); + + let producer_state = CutePipelineStateAttr::producer(STAGES); + let producer = CutePipelineStateNewOp::new(&mut ctx, producer_state); + let producer_slot = producer.slot(&ctx); + let producer_phase = producer.phase(&ctx); + let producer_index = CutePipelineStateSlotOp::new(&mut ctx, producer_slot, producer_state); + let acquire = CutePipelineProducerAcquireOp::new( + &mut ctx, + pipeline, + producer_slot, + producer_phase, + producer_state, + ); + let expect = + CutePipelineProducerExpectTxOp::new(&mut ctx, pipeline, producer_slot, producer_state); + let producer_advance = + CutePipelineStateAdvanceOp::new(&mut ctx, producer_slot, producer_phase, producer_state); + let producer_next_slot = producer_advance.next_slot(&ctx); + let producer_next_phase = producer_advance.next_phase(&ctx); + let tail = CutePipelineProducerTailOp::new( + &mut ctx, + pipeline, + producer_next_slot, + producer_next_phase, + producer_state, + ); + + let consumer_state = CutePipelineStateAttr::consumer(STAGES); + let consumer = CutePipelineStateNewOp::new(&mut ctx, consumer_state); + let consumer_slot = consumer.slot(&ctx); + let consumer_phase = consumer.phase(&ctx); + let consumer_index = CutePipelineStateSlotOp::new(&mut ctx, consumer_slot, consumer_state); + let wait = CutePipelineConsumerWaitOp::new( + &mut ctx, + pipeline, + consumer_slot, + consumer_phase, + consumer_state, + ); + let release = + CutePipelineConsumerReleaseOp::new(&mut ctx, pipeline, consumer_slot, consumer_state); + let consumer_advance = + CutePipelineStateAdvanceOp::new(&mut ctx, consumer_slot, consumer_phase, consumer_state); + + assert!(make.verify(&ctx).is_ok()); + assert!(init.verify(&ctx).is_ok()); + assert!(producer.verify(&ctx).is_ok()); + assert!(producer_index.verify(&ctx).is_ok()); + assert!(acquire.verify(&ctx).is_ok()); + assert!(expect.verify(&ctx).is_ok()); + assert!(producer_advance.verify(&ctx).is_ok()); + assert!(tail.verify(&ctx).is_ok()); + assert!(consumer.verify(&ctx).is_ok()); + assert!(consumer_index.verify(&ctx).is_ok()); + assert!(wait.verify(&ctx).is_ok()); + assert!(release.verify(&ctx).is_ok()); + assert!(consumer_advance.verify(&ctx).is_ok()); + assert_eq!(make.pipeline_type(&ctx).unwrap().storage_bytes(), Some(48)); + assert_eq!( + expect.completion_barrier(&ctx).get_type(&ctx), + base.get_type(&ctx) + ); +} + +/// Zero/overflowing stages and unsupported transaction or participation counts fail configuration checks. +#[test] +fn pipeline_config_rejects_impossible_rings() { + let ctx = Context::new(); + assert!( + CuteTmaLoadPipelineType { + stages: 0, + consumer_warps: 8, + transaction_bytes: 1, + } + .verify(&ctx) + .is_err() + ); + assert!( + CuteTmaLoadPipelineType { + stages: 3, + consumer_warps: 33, + transaction_bytes: 1, + } + .verify(&ctx) + .is_err() + ); + assert!( + CuteTmaLoadPipelineType { + stages: 3, + consumer_warps: 8, + transaction_bytes: 0, + } + .verify(&ctx) + .is_err() + ); + let overflow = CuteTmaLoadPipelineType { + stages: u64::MAX, + consumer_warps: 8, + transaction_bytes: 1, + }; + assert_eq!(overflow.storage_bytes(), None); + assert!(overflow.verify(&ctx).is_err()); + assert!(CutePipelineStateAttr::producer(0).verify(&ctx).is_err()); +} + +/// Pipeline creation requires matching shared barrier storage, capacity, and alignment. +#[test] +fn make_rejects_wrong_storage_and_weak_promises() { + let Setup { + mut ctx, + base, + u32_value, + .. + } = setup(); + let weak = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + base, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 4, + ); + assert!( + weak.verify(&ctx) + .unwrap_err() + .to_string() + .contains("at least 8") + ); + + let wrong_pointee_ty = u32_value.get_type(&ctx); + let wrong_pointee = shared_pointer(&mut ctx, wrong_pointee_ty, true); + let wrong = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + wrong_pointee, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + assert!( + wrong + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("unsigned u64 Barrier") + ); + + let barrier = u64_type(&ctx); + let immutable = shared_pointer(&mut ctx, barrier, false); + let immutable = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + immutable, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + assert!(immutable.verify(&ctx).is_err()); + + let barrier = u64_type(&ctx); + let generic_pointer_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, barrier, true).into(); + let generic_pointer = undef(&mut ctx, generic_pointer_ty); + let generic = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + generic_pointer, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + assert!(generic.verify(&ctx).is_err()); + + let signed_u64: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signed).into(); + let signed_storage = shared_pointer(&mut ctx, signed_u64, true); + let signed_storage = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + signed_storage, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + assert!(signed_storage.verify(&ctx).is_err()); +} + +/// Lifecycle effects cannot mix producer/consumer roles, ring sizes, or cursor scalar types. +#[test] +fn effects_reject_role_stage_and_scalar_drift() { + let Setup { + mut ctx, + base, + u32_value, + u64_value, + } = setup(); + let make = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + base, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + let pipeline = make.pipeline(&ctx); + + let wrong_role = CutePipelineProducerAcquireOp::new( + &mut ctx, + pipeline, + u32_value, + u32_value, + CutePipelineStateAttr::consumer(STAGES), + ); + assert!( + wrong_role + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("Producer") + ); + + let wrong_stages = CutePipelineConsumerWaitOp::new( + &mut ctx, + pipeline, + u32_value, + u32_value, + CutePipelineStateAttr::consumer(2), + ); + assert!( + wrong_stages + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("3") + ); + + let wrong_scalar = CutePipelineConsumerReleaseOp::new( + &mut ctx, + pipeline, + u64_value, + CutePipelineStateAttr::consumer(STAGES), + ); + assert!( + wrong_scalar + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("32-bit") + ); +} + +/// The expectation result must remain the shared u64 barrier pointer consumed by TMA copies. +#[test] +fn expect_tx_rejects_barrier_pointer_type_drift() { + let Setup { + mut ctx, + base, + u32_value, + .. + } = setup(); + let make = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + base, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + let pipeline = make.pipeline(&ctx); + let base_pointer_ty = base.get_type(&ctx); + let wrong_result: TypeHandle = MirPtrType::get_generic(&mut ctx, base_pointer_ty, true).into(); + let raw = Operation::new( + &mut ctx, + CutePipelineProducerExpectTxOp::get_concrete_op_info(), + vec![wrong_result], + vec![pipeline, u32_value], + vec![], + 0, + ); + let expect = CutePipelineProducerExpectTxOp::wrap(raw); + expect.set_attr_pipeline_producer_expect_state(&ctx, CutePipelineStateAttr::producer(STAGES)); + + assert!( + expect + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("must match") + ); +} + +/// The compiler-only handle cannot escape through an ordinary non-lifecycle operation. +#[test] +fn ghost_pipeline_rejects_an_unrelated_consumer() { + let Setup { mut ctx, base, .. } = setup(); + let make = CuteTmaLoadPipelineMakeOp::new( + &mut ctx, + base, + STAGES, + CONSUMER_WARPS, + TRANSACTION_BYTES, + 8, + ); + let pipeline = make.pipeline(&ctx); + let _unrelated = + CutePipelineStateSlotOp::new(&mut ctx, pipeline, CutePipelineStateAttr::producer(STAGES)); + + assert!( + make.verify(&ctx) + .unwrap_err() + .to_string() + .contains("lifecycle operations") + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/mod.rs new file mode 100644 index 0000000000..0f8e2bfabc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/mod.rs @@ -0,0 +1,501 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Static scheduling kept as five small CuTe operations. +//! +//! A physical CUDA block can process more than one logical output tile: +//! +//! ```text +//! new_1d +//! │ current = blockIdx.x, stride = gridDim.x +//! ▼ +//! current ── has_work? ── yes ──► get_current_work ──► (linear, m, n, batch) +//! ▲ │ +//! └──────── advance_to_next_work(current, stride) ◄─────────────┘ +//! ``` +//! +//! The loop carries only ordinary `u64` values. `work_tile` is a short-lived +//! compiler handle, so this layer adds no runtime scheduler object. A backend +//! continuation materializes the required special-register reads and integer +//! arithmetic. + +use pliron::builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface}, + types::{IntegerType, Signedness}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::CuteTileGridAttr; +use crate::types::CuteWorkTileType; + +/// Intern the unsigned 64-bit scalar type used for extents and indices. +fn u64_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 64, Signedness::Unsigned).into() +} + +/// Intern the one-bit integer type used by work predicates. +fn bool_type(ctx: &Context) -> TypeHandle { + IntegerType::get(ctx, 1, Signedness::Signless).into() +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Check that the SSA value has the one-bit predicate type. +fn is_bool(ctx: &Context, value: Value) -> bool { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) +} + +/// Read the selected work tile's static grid from its compiler-only type. +fn work_tile_of(ctx: &Context, value: Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .cloned() +} + +/// Check whether the current linear number is inside the logical tile grid. +/// +/// Keeping this separate from [`CuteSchedulerCurrentOp`] matters: the tile's +/// M/N/batch arithmetic remains inside the true loop path. A finished block +/// performs only the bounds check. +#[pliron_op( + name = "cute.scheduler_has_work", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], + attributes = (has_work_grid: CuteTileGridAttr) +)] +pub struct CuteSchedulerHasWorkOp; + +impl CuteSchedulerHasWorkOp { + /// Build a detached predicate comparing the linear cursor with the static output-grid size. + pub fn new(ctx: &mut Context, current: Value, grid: CuteTileGridAttr) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![bool_type(ctx)], + vec![current], + vec![], + 0, + ), + }; + operation.set_attr_has_work_grid(ctx, grid); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Current linear tile number carried by the loop. + #[must_use] + pub fn current(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// True while `current` is smaller than the static tile count. + #[must_use] + pub fn has_work(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Static logical output grid used by the bounds check. + #[must_use] + pub fn tile_grid(&self, ctx: &Context) -> Option { + self.get_attr_has_work_grid(ctx).map(|grid| *grid) + } +} + +impl Verify for CuteSchedulerHasWorkOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scheduler_has_work needs 1 operand and 1 result" + ); + } + if !is_u64(ctx, self.current(ctx)) || !is_bool(ctx, self.has_work(ctx)) { + return verify_err!( + op.loc(), + "cute.scheduler_has_work needs a u64 current value and an i1 result" + ); + } + let Some(grid) = self.tile_grid(ctx) else { + return verify_err!(op.loc(), "cute.scheduler_has_work must carry a tile grid"); + }; + grid.verify(ctx) + } +} + +/// Start this block's one-dimensional persistent schedule. +/// +/// Backend lowering reads `blockIdx.x` into `current` and `gridDim.x` into +/// `stride`. +/// Keeping those values as `u64` means an ordinary loop can carry them after +/// every semantic scheduler operation has been erased. +#[pliron_op( + name = "cute.scheduler_new_1d", + format, + interfaces = [NOpdsInterface<0>, NResultsInterface<2>], + attributes = (new_1d_grid: CuteTileGridAttr) +)] +pub struct CuteSchedulerNew1dOp; + +impl CuteSchedulerNew1dOp { + /// Build detached cursor/stride results from the physical one-dimensional launch grid. + pub fn new(ctx: &mut Context, grid: CuteTileGridAttr) -> Self { + let ty = u64_type(ctx); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ty, ty], + vec![], + vec![], + 0, + ), + }; + operation.set_attr_new_1d_grid(ctx, grid); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Linear tile number assigned to this physical block first. + #[must_use] + pub fn current(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Number of physical blocks between this block's logical tiles. + #[must_use] + pub fn stride(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Static logical output grid described by this schedule. + #[must_use] + pub fn tile_grid(&self, ctx: &Context) -> Option { + self.get_attr_new_1d_grid(ctx).map(|grid| *grid) + } +} + +impl Verify for CuteSchedulerNew1dOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 0 || op.get_num_results() != 2 { + return verify_err!( + op.loc(), + "cute.scheduler_new_1d needs 0 operands and 2 results" + ); + } + if !is_u64(ctx, self.current(ctx)) || !is_u64(ctx, self.stride(ctx)) { + return verify_err!( + op.loc(), + "cute.scheduler_new_1d current and stride must be unsigned 64-bit integers" + ); + } + let Some(grid) = self.tile_grid(ctx) else { + return verify_err!(op.loc(), "cute.scheduler_new_1d must carry a tile grid"); + }; + grid.verify(ctx) + } +} + +/// Expose the current linear number as a semantic work-tile handle. +/// +/// Call this only on the true result of [`CuteSchedulerHasWorkOp`]. The handle +/// carries no runtime bits; backend lowering resolves it back to `current`. +#[pliron_op( + name = "cute.scheduler_current", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], + attributes = (current_grid: CuteTileGridAttr) +)] +pub struct CuteSchedulerCurrentOp; + +impl CuteSchedulerCurrentOp { + /// Build a detached work-tile selection carrying the static grid and current linear cursor. + pub fn new(ctx: &mut Context, current: Value, grid: CuteTileGridAttr) -> Self { + let tile: TypeHandle = CuteWorkTileType::get(ctx, grid).into(); + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![tile], + vec![current], + vec![], + 0, + ), + }; + operation.set_attr_current_grid(ctx, grid); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Current linear tile number carried by the loop. + #[must_use] + pub fn current(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Semantic tile handle created on the `has_work` path. + #[must_use] + pub fn work_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Static logical output grid used by the bounds check and tile mapping. + #[must_use] + pub fn tile_grid(&self, ctx: &Context) -> Option { + self.get_attr_current_grid(ctx).map(|grid| *grid) + } +} + +impl Verify for CuteSchedulerCurrentOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scheduler_current needs 1 operand and 1 result" + ); + } + if !is_u64(ctx, self.current(ctx)) { + return verify_err!(op.loc(), "cute.scheduler_current input must be a u64 value"); + } + let Some(grid) = self.tile_grid(ctx) else { + return verify_err!(op.loc(), "cute.scheduler_current must carry a tile grid"); + }; + grid.verify(ctx)?; + let Some(tile) = work_tile_of(ctx, self.work_tile(ctx)) else { + return verify_err!( + op.loc(), + "cute.scheduler_current result must be a work tile" + ); + }; + if tile.grid != grid { + return verify_err!( + op.loc(), + "cute.scheduler_current work tile must use the operation's tile grid" + ); + } + tile.verify(ctx) + } +} + +/// Map one valid linear tile to M-fastest `(m, n, batch)` coordinates. +/// +/// Backend lowering uses the same two divide and multiply/subtract pairs as +/// the Rust scheduler. `linear` is returned unchanged. +#[pliron_op( + name = "cute.work_tile_coordinates", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<4>] +)] +pub struct CuteWorkTileCoordinatesOp; + +impl CuteWorkTileCoordinatesOp { + /// Build detached linear/M/N/batch results for one selected work tile. + pub fn new(ctx: &mut Context, tile: Value) -> Self { + let ty = u64_type(ctx); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![ty, ty, ty, ty], + vec![tile], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the selected compiler-only output tile consumed by coordinate extraction. + #[must_use] + pub fn work_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the linear logical tile-number result. + #[must_use] + pub fn linear(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Return the output M-tile coordinate result. + #[must_use] + pub fn m_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(1) + } + + /// Return the output N-tile coordinate result. + #[must_use] + pub fn n_tile(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(2) + } + + /// Return the batch-coordinate result. + #[must_use] + pub fn batch(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(3) + } +} + +impl Verify for CuteWorkTileCoordinatesOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 4 { + return verify_err!( + op.loc(), + "cute.work_tile_coordinates needs 1 operand and 4 results" + ); + } + let Some(tile) = work_tile_of(ctx, self.work_tile(ctx)) else { + return verify_err!( + op.loc(), + "cute.work_tile_coordinates input must be a work tile" + ); + }; + tile.verify(ctx)?; + if ![ + self.linear(ctx), + self.m_tile(ctx), + self.n_tile(ctx), + self.batch(ctx), + ] + .into_iter() + .all(|value| is_u64(ctx, value)) + { + return verify_err!( + op.loc(), + "cute.work_tile_coordinates results must be unsigned 64-bit integers" + ); + } + Ok(()) + } +} + +/// Move this block to its next logical tile. +/// +/// The result is `current.saturating_add(stride)`. Saturation prevents a +/// wrapped tile number from looking valid again. +#[pliron_op( + name = "cute.scheduler_advance", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>], + attributes = (advance_grid: CuteTileGridAttr) +)] +pub struct CuteSchedulerAdvanceOp; + +impl CuteSchedulerAdvanceOp { + /// Build a detached linear-cursor update by the physical-grid stride. + pub fn new(ctx: &mut Context, current: Value, stride: Value, grid: CuteTileGridAttr) -> Self { + let operation = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![u64_type(ctx)], + vec![current, stride], + vec![], + 0, + ), + }; + operation.set_attr_advance_grid(ctx, grid); + operation + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } + + /// Return the input linear tile cursor before scheduler advancement. + #[must_use] + pub fn current(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(0) + } + + /// Return the physical-grid stride added to the linear cursor. + #[must_use] + pub fn stride(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_operand(1) + } + + /// Return the advanced linear tile cursor. + #[must_use] + pub fn next(&self, ctx: &Context) -> Value { + self.get_operation().deref(ctx).get_result(0) + } + + /// Static logical output grid described by this schedule. + #[must_use] + pub fn tile_grid(&self, ctx: &Context) -> Option { + self.get_attr_advance_grid(ctx).map(|grid| *grid) + } +} + +impl Verify for CuteSchedulerAdvanceOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.scheduler_advance needs 2 operands and 1 result" + ); + } + if !is_u64(ctx, self.current(ctx)) + || !is_u64(ctx, self.stride(ctx)) + || !is_u64(ctx, self.next(ctx)) + { + return verify_err!( + op.loc(), + "cute.scheduler_advance current, stride, and next must be unsigned 64-bit integers" + ); + } + let Some(grid) = self.tile_grid(ctx) else { + return verify_err!(op.loc(), "cute.scheduler_advance must carry a tile grid"); + }; + grid.verify(ctx) + } +} + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/tests.rs new file mode 100644 index 0000000000..10f95fd51b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/scheduler/tests.rs @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use dialect_mir::ops::MirUndefOp; + +/// Return result zero from the fixture operation for connecting a semantic chain. +fn result(ctx: &Context, op: Ptr, index: usize) -> Value { + op.deref(ctx).get_result(index) +} + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Scheduler construction, work test, tile selection, coordinates, and advance remain explicit operations. +#[test] +fn persistent_schedule_stays_five_composable_steps() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let grid = CuteTileGridAttr::new(16, 16, 1); + + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid); + assert!(start.verify(&ctx).is_ok()); + let current = start.current(&ctx); + let stride = start.stride(&ctx); + + let has_work = CuteSchedulerHasWorkOp::new(&mut ctx, current, grid); + assert!(has_work.verify(&ctx).is_ok()); + + let selected = CuteSchedulerCurrentOp::new(&mut ctx, current, grid); + assert!(selected.verify(&ctx).is_ok()); + let tile = selected.work_tile(&ctx); + + let coordinates = CuteWorkTileCoordinatesOp::new(&mut ctx, tile); + assert!(coordinates.verify(&ctx).is_ok()); + + let advance = CuteSchedulerAdvanceOp::new(&mut ctx, current, stride, grid); + assert!(advance.verify(&ctx).is_ok()); + assert!(is_u64(&ctx, advance.next(&ctx))); +} + +/// Selected work-tile metadata must agree with the scheduler's static grid. +#[test] +fn current_rejects_a_work_tile_for_a_different_grid() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u64_ty = u64_type(&ctx); + let current = undef(&mut ctx, u64_ty); + let grid = CuteTileGridAttr::new(16, 16, 1); + let selected = CuteSchedulerCurrentOp::new(&mut ctx, current, grid); + selected.set_attr_current_grid(&ctx, CuteTileGridAttr::new(8, 32, 1)); + + assert!( + selected + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("operation's tile grid") + ); +} + +/// Coordinate extraction requires a semantic work-tile ghost, not an untyped scalar. +#[test] +fn coordinates_reject_an_ordinary_runtime_value() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u64_ty = u64_type(&ctx); + let scalar = undef(&mut ctx, u64_ty); + let coordinates = CuteWorkTileCoordinatesOp::new(&mut ctx, scalar); + + assert!( + coordinates + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("input must be a work tile") + ); +} + +/// Grid dimensions must be positive and their product must fit the supported integer size. +#[test] +fn grid_checks_zero_and_overflow() { + let ctx = Context::new(); + assert!(CuteTileGridAttr::new(3, 2, 1).verify(&ctx).is_ok()); + assert_eq!(CuteTileGridAttr::new(3, 2, 1).total_tiles(), Some(6)); + assert!(CuteTileGridAttr::new(0, 2, 1).verify(&ctx).is_err()); + assert!(CuteTileGridAttr::new(u64::MAX, 2, 1).verify(&ctx).is_err()); +} + +/// The coordinate ABI exposes linear, M, N, and batch results in that exact order. +#[test] +fn result_order_matches_the_visual_contract() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let grid = CuteTileGridAttr::new(3, 2, 1); + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid); + let current = start.current(&ctx); + let has_work = CuteSchedulerHasWorkOp::new(&mut ctx, current, grid); + let selected = CuteSchedulerCurrentOp::new(&mut ctx, current, grid); + let tile = selected.work_tile(&ctx); + let coordinates = CuteWorkTileCoordinatesOp::new(&mut ctx, tile); + + assert_eq!( + coordinates.linear(&ctx), + result(&ctx, coordinates.get_operation(), 0) + ); + assert_eq!( + coordinates.m_tile(&ctx), + result(&ctx, coordinates.get_operation(), 1) + ); + assert_eq!( + coordinates.n_tile(&ctx), + result(&ctx, coordinates.get_operation(), 2) + ); + assert_eq!( + coordinates.batch(&ctx), + result(&ctx, coordinates.get_operation(), 3) + ); + assert!(is_bool(&ctx, has_work.has_work(&ctx))); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/mod.rs new file mode 100644 index 0000000000..ba34d8147c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/mod.rs @@ -0,0 +1,75 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 tensor-memory, collector-aware tiled MMA, and cluster pipelines. +//! +//! These operations keep ordinary MIR scalar carriers and retain the static +//! layout and collective contracts as typed attributes. A paired MMA is one +//! semantic operation: its A collector is filled and last-used within each +//! K=16 step, so a backend cannot accidentally carry it across K iterations. +//! +//! Cluster TMA loads are collective across all 32 producer-warp lanes. Native +//! CuTe lowering owns lane election and warp synchronization; callers must +//! not place a load inside an elected-lane branch. TMA stores retain a single +//! issuer, and accumulator release fences all lanes before electing an arrival. +//! +//! All operations have zero SSA results. Their docs give the operand order for +//! `new`; effects update shared memory, tensor memory, or hardware barriers. +//! TMEM addresses are encoded `u32` addresses, not byte pointers. Ring stages +//! are zero-based indices, and phases are parity bits maintained by the caller. + +use dialect_mir::ops::MirConstantOp; +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::attribute::Attribute; +use pliron::builtin::types::{FP32Type, IntegerType}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::{Op, OpId}; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::Typed; +use pliron::value::Value; +use pliron::{verify_err, verify_err_noloc}; +use pliron_derive::{pliron_attr, pliron_op}; + +use crate::attributes::{CuteComposedLayoutAttr, CuteTmaStorePipelineAttr}; +use crate::layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}; + +/// Build a zero-offset FP16 row-major layout with the requested hardware swizzle. +fn row_major(rows: u32, columns: u32, swizzle_bits: u32) -> ComposedLayout { + let outer: Layout = format!("({rows},{columns}):({columns},1)") + .parse() + .expect("fixed rank-two layout"); + ComposedLayout::new( + Swizzle::new(swizzle_bits, 3, 3), + 0, + outer, + OffsetUnit::Elements, + ) + .expect("fixed SM100 FP16 swizzle") +} + +mod operations; +mod plans; +mod verify; + +pub use operations::{ + CuteSm100AccumulatorAcquireOp, CuteSm100AccumulatorCommitOp, CuteSm100AccumulatorInitOp, + CuteSm100AccumulatorReleaseOp, CuteSm100AccumulatorWaitOp, CuteSm100ClusterTmaLoadOp, + CuteSm100PipelineAcquireOp, CuteSm100PipelineExpectOp, CuteSm100PipelineInitOp, + CuteSm100PipelineReleaseOp, CuteSm100PipelineWaitOp, CuteSm100StoreAcquireOp, + CuteSm100StoreCommitOp, CuteSm100StoreTailOp, CuteSm100TiledMmaOp, CuteSm100TmaStoreOp, + CuteSm100TmemAllocOp, CuteSm100TmemDeallocOp, CuteSm100TmemEpilogueOp, register, semantic_ids, + verify_local, +}; +pub use plans::{ + CuteSm100ClusterCopyPlanAttr, CuteSm100CollectorAttr, CuteSm100EpiloguePlanAttr, + CuteSm100MmaPlanAttr, CuteSm100PipelinePlanAttr, CuteSm100TmemPlanAttr, +}; +pub use verify::verify_story; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/operations.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/operations.rs new file mode 100644 index 0000000000..4912d0cea1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/operations.rs @@ -0,0 +1,409 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 operation definitions, local contracts, and registration. + +use super::*; + +/// Runtime scalar or pointer contract required at one positional operand. +#[derive(Clone, Copy)] +enum Carrier { + /// Unsigned 32-bit index or encoded tensor-memory address. + U32, + /// One-bit predicate selecting whether MMA retains the previous accumulator. + Bool, + /// FP32 scalar storage or accumulation. + F32, + /// Writable shared pointer to a 32-bit allocation token. + SharedU32, + /// Writable shared pointer to a 64-bit hardware barrier. + SharedBarrier, + /// Writable shared FP16 tile pointer. + SharedF16, + /// Readable shared FP16 tile pointer. + ReadSharedF16, + /// Immutable tensor-map pointer in generic, global, or constant space. + Descriptor, +} +/// Check an operand's scalar width or pointer storage/address-space/mutability contract. +fn carrier_matches(ctx: &Context, value: Value, carrier: Carrier) -> bool { + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let integer = |width| { + ty.downcast_ref::() + .is_some_and(|i| i.width() == width && (width == 1 || i.is_unsigned())) + }; + match carrier { + Carrier::U32 => integer(32), + Carrier::Bool => integer(1), + Carrier::F32 => ty.downcast_ref::().is_some(), + _ => { + let Some(ptr) = ty.downcast_ref::() else { + return false; + }; + if matches!(carrier, Carrier::Descriptor) { + return !ptr.is_mutable + && matches!( + ptr.address_space, + address_space::GENERIC | address_space::GLOBAL | address_space::CONSTANT + ); + } + if !matches!( + ptr.address_space, + address_space::SHARED | address_space::GENERIC + ) { + return false; + } + if !matches!(carrier, Carrier::ReadSharedF16) && !ptr.is_mutable { + return false; + } + let pointee = ptr.pointee.deref(ctx); + match carrier { + Carrier::SharedU32 => pointee + .downcast_ref::() + .is_some_and(|i| i.width() == 32 && i.is_unsigned()), + Carrier::SharedBarrier => pointee + .downcast_ref::() + .is_some_and(|i| i.width() == 64 && i.is_unsigned()), + Carrier::SharedF16 | Carrier::ReadSharedF16 => { + pointee.downcast_ref::().is_some() + } + _ => false, + } + } + } +} + +/// Generate zero-result operations with a typed plan and local carrier verification. +macro_rules! semantic_op { + ($(#[$doc:meta])* $ty:ident, $name:literal, $key:ident, $plan:ty, [$($carrier:ident),*], $check:expr) => { + #[pliron_op(name = $name, format, attributes = ($key: $plan))] + pub struct $ty; + // Keep contracts on the impl and constructor: pliron's generated + // attribute-table links do not compose correctly with local struct docs. + $(#[$doc])* + impl $ty { + $(#[$doc])* + pub fn new(ctx: &mut Context, operands: Vec, plan: $plan) -> Self { + let op = Self { op: Operation::new(ctx, Self::get_concrete_op_info(), vec![], operands, vec![], 0) }; + op.get_operation().deref_mut(ctx).attributes.set(stringify!($key).try_into().expect("plan key"), plan); + op + } + /// View an existing operation without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { Self { op } } + /// Borrow the typed plan while retaining the operation's context borrow. + pub fn get_attr_plan<'a>(&self, ctx: &'a Context) -> Option> { + std::cell::Ref::filter_map(self.get_operation().deref(ctx), |op| + op.attributes.get::<$plan>(&stringify!($key).try_into().expect("plan key"))).ok() + } + /// Clone the typed plan, or return `None` when the attribute is missing. + pub fn plan(&self, ctx: &Context) -> Option<$plan> { self.get_attr_plan(ctx).map(|attr| (*attr).clone()) } + } + impl Verify for $ty { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + let expected: &[Carrier] = &[$(Carrier::$carrier),*]; + if op.get_num_operands() != expected.len() || op.get_num_results() != 0 { + return verify_err!(op.loc(), "{} needs {} operands and no results", $name, expected.len()); + } + let Some(plan) = self.plan(ctx) else { return verify_err!(op.loc(), "{} requires its typed plan", $name); }; + plan.verify(ctx)?; + if !($check)(&plan) { return verify_err!(op.loc(), "{} has the wrong pipeline role", $name); } + for (index, carrier) in expected.iter().enumerate() { + if !carrier_matches(ctx, op.get_operand(index), *carrier) { + return verify_err!(op.loc(), "{} operand {} has an invalid scalar/pointer carrier", $name, index); + } + } + Ok(()) + } + } + }; +} +semantic_op!( + /// Allocate TMEM columns and relinquish the allocation permit. + /// + /// Operands: `[token]`, a shared `u32` destination for the encoded address. + /// All lanes of the allocating warp in both CTAs participate; synchronize + /// readers before they load the token. + CuteSm100TmemAllocOp, + "cute.sm100_tmem_alloc", + sm100_tmem_alloc_plan, + CuteSm100TmemPlanAttr, + [SharedU32], + |_| true +); +semantic_op!( + /// Free a completed two-CTA TMEM allocation. + /// + /// Operands: `[tmem]`, the encoded base address. Both allocating warps + /// participate after all asynchronous users finish and the CTAs rendezvous. + CuteSm100TmemDeallocOp, + "cute.sm100_tmem_dealloc", + sm100_tmem_dealloc_plan, + CuteSm100TmemPlanAttr, + [U32], + |_| true +); +semantic_op!( + /// Issue the paired A×B0 and A×B1 products into TMEM. + /// + /// Operands: `[tmem, a, b0, b1, accumulate]`; A/B pointers address shared + /// FP16 tiles. One elected leader lane issues after input completion. + /// `accumulate` controls the first K step; subsequent steps accumulate. + /// All source tiles and destination TMEM stay live until MMA completion. + CuteSm100TiledMmaOp, + "cute.sm100_tiled_mma", + sm100_tiled_mma_plan, + CuteSm100MmaPlanAttr, + [U32, ReadSharedF16, ReadSharedF16, ReadSharedF16, Bool], + |_| true +); +semantic_op!( + /// Add a row bias to FP32 TMEM values and place FP16 results in shared memory. + /// + /// Operands: `[tmem_slice, destination, tid, bias]`. The encoded slice + /// selects the caller's 32 TMEM rows and 32 columns; `tid` is in `0..128`. + /// All lanes of four epilogue warps participate after accumulator wait. + /// Publish their shared writes before issuing a TMA store. + CuteSm100TmemEpilogueOp, + "cute.sm100_tmem_epilogue", + sm100_tmem_epilogue_plan, + CuteSm100EpiloguePlanAttr, + [U32, SharedF16, U32, F32], + |_| true +); +semantic_op!( + /// Start one CTA's contribution to a cluster input transfer. + /// + /// Operands: `[descriptor, destination, row, column, barrier, rank]`. + /// Coordinates count global FP16 elements; rank is the local CTA rank. + /// All producer-warp lanes participate. Each CTA's destination covers its + /// local tile; the leader's expected byte count covers both CTAs. + CuteSm100ClusterTmaLoadOp, + "cute.sm100_cluster_tma_load", + sm100_cluster_tma_load_plan, + CuteSm100ClusterCopyPlanAttr, + [Descriptor, SharedF16, U32, U32, SharedBarrier, U32], + |_| true +); +semantic_op!( + /// Initialize input barriers from one elected lane per CTA. + /// + /// Operands: `[full, empty]`, shared arrays of `plan.stages` barriers. + /// Publish initialization and synchronize the cluster before use. + CuteSm100PipelineInitOp, + "cute.sm100_pipeline_init", + sm100_pipeline_init_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + /// Wait until the producer may reuse an input stage. + /// + /// Operands: `[empty, stage, phase]`. The producer starts with phase one; + /// the stage must be inside the shared empty-barrier array. + CuteSm100PipelineAcquireOp, + "cute.sm100_pipeline_acquire", + sm100_pipeline_acquire_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + /// Set the full barrier's expected transaction bytes for one input stage. + /// + /// Operands: `[full, stage]`. One elected rank-zero lane calls after + /// acquire; the plan counts the copies from both CTAs. + CuteSm100PipelineExpectOp, + "cute.sm100_pipeline_expect", + sm100_pipeline_expect_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + /// Wait for an input stage and order subsequent tensor-core reads. + /// + /// Operands: `[full, stage, phase]`. Rank zero's MMA warp supplies its + /// current consumer stage and phase. + CuteSm100PipelineWaitOp, + "cute.sm100_pipeline_wait", + sm100_pipeline_wait_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + /// Release an input stage in both CTAs after pending MMA reads complete. + /// + /// Operands: `[empty, stage]`. One elected leader lane calls after all + /// MMAs reading that stage have been issued. + CuteSm100PipelineReleaseOp, + "cute.sm100_pipeline_release", + sm100_pipeline_release_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes > 0 +); +semantic_op!( + /// Initialize the single accumulator slot from one elected lane per CTA. + /// + /// Operands: `[empty, full]`, two distinct shared barriers. Publish their + /// initialization cluster-wide before either side uses the slot. + CuteSm100AccumulatorInitOp, + "cute.sm100_accumulator_init", + sm100_accumulator_init_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + /// Wait until epilogue warps have released the previous accumulator tile. + /// + /// Operands: `[empty, phase]`. The leader CTA's MMA warp acquires with + /// initial phase one. + CuteSm100AccumulatorAcquireOp, + "cute.sm100_accumulator_acquire", + sm100_accumulator_acquire_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + /// Publish the accumulator to both CTAs after pending MMA completes. + /// + /// Operands: `[full]`. One elected leader lane calls after the last K stage. + CuteSm100AccumulatorCommitOp, + "cute.sm100_accumulator_commit", + sm100_accumulator_commit_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + /// Wait for an accumulator tile and order subsequent TMEM loads. + /// + /// Operands: `[full, phase]`. All consumer-warp lanes supply their current + /// phase, initially zero. + CuteSm100AccumulatorWaitOp, + "cute.sm100_accumulator_wait", + sm100_accumulator_wait_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier, U32], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + /// Release one consumer warp's accumulator portion to the leader CTA. + /// + /// Operands: `[empty]`. All epilogue-warp lanes call after their TMEM reads; + /// lowering fences every lane and elects one arrival per warp. + CuteSm100AccumulatorReleaseOp, + "cute.sm100_accumulator_release", + sm100_accumulator_release_plan, + CuteSm100PipelinePlanAttr, + [SharedBarrier], + |p: &CuteSm100PipelinePlanAttr| p.transaction_bytes == 0 +); +semantic_op!( + /// Start an asynchronous shared-to-global FP16 output transfer. + /// + /// Operands: `[descriptor, source, row, column]`; coordinates count global + /// elements. One elected issuer calls after writers publish and synchronize. + /// Keep the source unchanged until its committed group finishes reading it. + CuteSm100TmaStoreOp, + "cute.sm100_tma_store", + sm100_tma_store_plan, + CuteSm100EpiloguePlanAttr, + [Descriptor, ReadSharedF16, U32, U32], + |_| true +); +semantic_op!( + /// Commit this issuer's pending TMA stores as one hardware group. + /// + /// No operands. The same issuer performs the matching acquire and tail. + CuteSm100StoreCommitOp, + "cute.sm100_store_commit", + sm100_store_commit_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); +semantic_op!( + /// Wait until at most `plan.stages - 1` store groups still read shared memory. + /// + /// No operands. The store issuer may then reuse the oldest buffer in its + /// ring; the newest group can still be pending. + CuteSm100StoreAcquireOp, + "cute.sm100_store_acquire", + sm100_store_acquire_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); +semantic_op!( + /// Wait until all committed stores stop reading shared memory. + /// + /// No operands. The store issuer calls before releasing its output buffers. + CuteSm100StoreTailOp, + "cute.sm100_store_tail", + sm100_store_tail_plan, + CuteTmaStorePipelineAttr, + [], + |_| true +); + +/// Apply a callback to the closed paired-SM100 operation set. +macro_rules! for_all_ops { + ($callback:ident) => { + $callback!( + CuteSm100TmemAllocOp, + CuteSm100TmemDeallocOp, + CuteSm100TiledMmaOp, + CuteSm100TmemEpilogueOp, + CuteSm100ClusterTmaLoadOp, + CuteSm100PipelineInitOp, + CuteSm100PipelineAcquireOp, + CuteSm100PipelineExpectOp, + CuteSm100PipelineWaitOp, + CuteSm100PipelineReleaseOp, + CuteSm100AccumulatorInitOp, + CuteSm100AccumulatorAcquireOp, + CuteSm100AccumulatorCommitOp, + CuteSm100AccumulatorWaitOp, + CuteSm100AccumulatorReleaseOp, + CuteSm100TmaStoreOp, + CuteSm100StoreCommitOp, + CuteSm100StoreAcquireOp, + CuteSm100StoreTailOp + ) + }; +} +/// All operations belonging to this semantic family. +pub fn semantic_ids() -> Vec { + /// Collect the registered operation IDs from the supplied type list. + macro_rules! ids { ($($ty:ty),*) => { vec![$(<$ty>::get_opid_static()),*] }; } + for_all_ops!(ids) +} +/// Register the SM100 operations and their typed plans. +pub fn register(ctx: &mut Context) { + /// Register every listed semantic operation in the supplied IR context. + macro_rules! register_ops { ($($ty:ty),*) => { $(<$ty>::register(ctx);)* }; } + for_all_ops!(register_ops); + CuteSm100CollectorAttr::register(ctx); + CuteSm100TmemPlanAttr::register(ctx); + CuteSm100MmaPlanAttr::register(ctx); + CuteSm100ClusterCopyPlanAttr::register(ctx); + CuteSm100EpiloguePlanAttr::register(ctx); + CuteSm100PipelinePlanAttr::register(ctx); +} +/// Locally verify a recognized operation; `None` means another family. +pub fn verify_local(ctx: &Context, op: Ptr) -> Option> { + let id = Operation::get_opid(op, ctx); + /// Select and run the local verifier for one recognized SM100 operation. + macro_rules! verify_ops { ($($ty:ty),*) => { $(if id == <$ty>::get_opid_static() { return Some(<$ty>::wrap(op).verify(ctx)); })* }; } + for_all_ops!(verify_ops); + None +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/plans.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/plans.rs new file mode 100644 index 0000000000..2676b9018b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/plans.rs @@ -0,0 +1,265 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Typed plans for SM100 tensor memory, MMA, TMA, and collective pipelines. + +use super::*; + +/// The first N partition fills A, and the second consumes its last use. +#[pliron_attr(name = "cute.sm100_collector", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum CuteSm100CollectorAttr { + /// Fill the A collector for the first product and consume its last use in the second. + AFillLastUse, +} + +/// Capacity of one two-CTA tensor-memory allocation, in columns per CTA. +/// +/// Each CTA has 128 rows. The supported column count is a power of two in +/// `32..=512`; allocation and deallocation must use the same plan. +#[pliron_attr( + name = "cute.sm100_tmem_plan", + format = "`<` $columns `,` $cta_group `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteSm100TmemPlanAttr { + /// Tensor-memory columns allocated per CTA. + pub columns: u32, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, +} +impl Verify for CuteSm100TmemPlanAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.cta_group != 2 + || !(32..=512).contains(&self.columns) + || !self.columns.is_power_of_two() + { + return verify_err_noloc!( + "SM100 tensor memory requires CTA group 2 and power-of-two columns in 32..=512" + ); + } + Ok(()) + } +} + +/// Two FP16 MMA products sharing A, accumulated into adjacent TMEM partitions. +/// +/// `m`, `n0`, `n1`, and `k` describe the complete two-CTA products. The shared +/// layouts describe each CTA's portion: A is `(m / 2, k)`, and each B is +/// `(n / 2, k)`, all K-major with a 128-byte swizzle. The current path fixes +/// `m = 256`, `k = 64`, and `tmem_columns = 512`; each N is a multiple of 16 +/// in `16..=256`. B1's output follows B0's `n0` columns. +#[pliron_attr( + name = "cute.sm100_mma_plan", + format = "`<` $m `,` $n0 `,` $n1 `,` $k `,` $tmem_columns `,` $cta_group `,` $collector `,` $a_layout `,` $b0_layout `,` $b1_layout `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100MmaPlanAttr { + /// Collective output-row extent M. + pub m: u32, + /// Output columns in the first product, starting at TMEM column zero. + pub n0: u32, + /// Output columns in the second product, immediately following the first. + pub n1: u32, + /// Reduction extent K, in logical elements. + pub k: u32, + /// TMEM capacity in columns available to both output partitions. + pub tmem_columns: u32, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, + /// Lifetime policy for the shared A collector across the two products. + pub collector: CuteSm100CollectorAttr, + /// Physical shared placement of each CTA's A tile. + pub a_layout: CuteComposedLayoutAttr, + /// Physical shared placement of each CTA's first B tile. + pub b0_layout: CuteComposedLayoutAttr, + /// Physical shared placement of each CTA's second B tile. + pub b1_layout: CuteComposedLayoutAttr, +} + +impl CuteSm100MmaPlanAttr { + /// Bytes copied by both CTAs for one complete A/B0/B1 stage. + pub fn transaction_bytes(&self) -> Option { + self.m + .checked_add(self.n0)? + .checked_add(self.n1)? + .checked_mul(self.k)? + .checked_mul(2) + } +} + +impl Verify for CuteSm100MmaPlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + CuteSm100TmemPlanAttr { + columns: self.tmem_columns, + cta_group: self.cta_group, + } + .verify(ctx)?; + if self.m != 256 + || self.k != 64 + || self.tmem_columns != 512 + || !(16..=256).contains(&self.n0) + || !(16..=256).contains(&self.n1) + || !self.n0.is_multiple_of(16) + || !self.n1.is_multiple_of(16) + || self.n0 + self.n1 > self.tmem_columns + { + return verify_err_noloc!( + "SM100 paired FP16 MMA requires M256/K64, two N partitions in 16..=256 divisible by 16, and 512 TMEM columns" + ); + } + for (name, layout, rows) in [ + ("A", &self.a_layout, 128), + ("B0", &self.b0_layout, self.n0 / 2), + ("B1", &self.b1_layout, self.n1 / 2), + ] { + layout.verify(ctx)?; + if layout.0 != row_major(rows, self.k, 3) { + return verify_err_noloc!( + "SM100 {name} MMA layout must be the matching K-major FP16 B128 shared tile" + ); + } + } + Ok(()) + } +} + +/// One CTA's FP16 input tile in a two-CTA cluster TMA copy. +/// +/// The layout is K-major with 64 columns and a 128-byte swizzle; its row count +/// is a multiple of eight in `8..=128`. Completion is routed to leader rank +/// zero. The caller's barrier expectation covers bytes copied by both CTAs. +#[pliron_attr( + name = "cute.sm100_cluster_copy_plan", + format = "`<` $layout `,` $cta_group `,` $leader_rank `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100ClusterCopyPlanAttr { + /// Static coordinate layout used to interpret the storage. + pub layout: CuteComposedLayoutAttr, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, + /// CTA rank receiving collective completion notifications. + pub leader_rank: u32, +} +impl Verify for CuteSm100ClusterCopyPlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.layout.verify(ctx)?; + let dimensions = self.layout.0.outer().shape.flatten(); + if self.cta_group != 2 + || self.leader_rank != 0 + || dimensions.len() != 2 + || dimensions[0] < 8 + || dimensions[0] > 128 + || dimensions[0] % 8 != 0 + || dimensions[1] != 64 + || self.layout.0 != row_major(dimensions[0] as u32, 64, 3) + { + return verify_err_noloc!( + "SM100 cluster TMA requires CTA group 2, leader 0, and a K64 FP16 B128 tile with 8..=128 rows divisible by 8" + ); + } + Ok(()) + } +} + +/// A 128-by-32 FP16 shared output tile, with a 64-byte swizzle. +/// +/// The epilogue fills this tile from FP32 TMEM, and the store operation sends +/// the same physical layout to global memory. Rows and columns count elements. +#[pliron_attr( + name = "cute.sm100_epilogue_plan", + format = "`<` $layout `,` $rows `,` $columns `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100EpiloguePlanAttr { + /// Static coordinate layout used to interpret the storage. + pub layout: CuteComposedLayoutAttr, + /// Row extent of the shared output tile, in elements. + pub rows: u32, + /// Allocated or copied column extent, in elements. + pub columns: u32, +} +impl Verify for CuteSm100EpiloguePlanAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.layout.verify(ctx)?; + if self.rows != 128 || self.columns != 32 || self.layout.0 != row_major(128, 32, 2) { + return verify_err_noloc!("SM100 epilogue requires the 128x32 FP16 B64 shared layout"); + } + Ok(()) + } +} + +/// Static barrier configuration for either input stages or one accumulator. +/// +/// Positive `transaction_bytes` selects the input pipeline: it counts all +/// A/B0/B1 bytes copied by both CTAs per stage, with one MMA completion arrival +/// releasing the stage. Zero selects the accumulator pipeline: one stage and +/// one release arrival per consumer warp across both CTAs. The current +/// epilogue has four consumer warps per CTA, hence eight arrivals. +/// Both forms use CTA group two, leader rank zero, and multicast mask `0b11`. +/// +/// ``` +/// use dialect_cute::ops::sm100::CuteSm100PipelinePlanAttr; +/// use pliron::{common_traits::Verify, context::Context}; +/// +/// let accumulator = CuteSm100PipelinePlanAttr { +/// stages: 1, +/// transaction_bytes: 0, +/// consumer_arrivals: 8, +/// cta_group: 2, +/// leader_rank: 0, +/// multicast_mask: 0b11, +/// }; +/// accumulator.verify(&Context::new()).unwrap(); +/// ``` +#[pliron_attr( + name = "cute.sm100_pipeline_plan", + format = "`<` $stages `,` $transaction_bytes `,` $consumer_arrivals `,` $cta_group `,` $leader_rank `,` $multicast_mask `>`" +)] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub struct CuteSm100PipelinePlanAttr { + /// Number of reusable slots in the barrier or output-buffer ring. + pub stages: u32, + /// Bytes expected by the full barrier for one complete input stage; zero selects accumulator synchronization. + pub transaction_bytes: u32, + /// Slot-release arrivals: one MMA completion for input, or one per epilogue warp for accumulators. + pub consumer_arrivals: u32, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, + /// CTA rank receiving collective completion notifications. + pub leader_rank: u32, + /// Bit mask of CTA ranks notified by a collective completion. + pub multicast_mask: u32, +} +impl Verify for CuteSm100PipelinePlanAttr { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.stages == 0 + || self.cta_group != 2 + || self.leader_rank != 0 + || self.multicast_mask != 3 + { + return verify_err_noloc!( + "SM100 pipeline requires positive stages, CTA group 2, leader 0, and multicast mask 3" + ); + } + if self.transaction_bytes == 0 { + if self.stages != 1 + || self.consumer_arrivals == 0 + || self.consumer_arrivals > 64 + || !self.consumer_arrivals.is_multiple_of(2) + { + return verify_err_noloc!( + "SM100 accumulator pipeline requires one stage and an even number of consumer arrivals in 2..=64" + ); + } + } else if self.consumer_arrivals != 1 || !self.transaction_bytes.is_multiple_of(16) { + return verify_err_noloc!( + "SM100 input pipeline requires one MMA completion arrival and a transaction size divisible by 16" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/tests.rs new file mode 100644 index 0000000000..519a3192a3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/tests.rs @@ -0,0 +1,395 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use dialect_mir::ops::{MirConstantOp, MirUndefOp}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::attributes::IntegerAttr; +use pliron::builtin::op_interfaces::{NOpdsInterface, NResultsInterface}; +use pliron::builtin::ops::ModuleOp; +use pliron::builtin::types::Signedness; +use pliron::linked_list::ContainsLinkedList; +use pliron::r#type::TypeHandle; +use pliron::utils::apint::APInt; +use std::num::NonZeroUsize; + +// Match importer-produced NVVM op IDs without a production dependency. +/// Test-only canonical two-CTA allocation recognized by the semantic lifecycle audit. +#[pliron_op(name = "nvvm.tcgen05_alloc_cg2", format, verifier = "succ", interfaces = [NOpdsInterface<2>, NResultsInterface<0>])] +struct TestAllocCg2Op; +/// Test-only canonical two-CTA deallocation recognized by the lifecycle audit. +#[pliron_op(name = "nvvm.tcgen05_dealloc_cg2", format, verifier = "succ", interfaces = [NOpdsInterface<2>, NResultsInterface<0>])] +struct TestDeallocCg2Op; +/// Test-only marker proving the two-CTA allocation permit was relinquished. +#[pliron_op(name = "nvvm.tcgen05_relinquish_alloc_permit_cg2", format, verifier = "succ", interfaces = [NOpdsInterface<0>, NResultsInterface<0>])] +struct TestRelinquishCg2Op; +/// Test-only one-CTA allocation used to reject the wrong collective mode. +#[pliron_op(name = "nvvm.tcgen05_alloc", format, verifier = "succ", interfaces = [NOpdsInterface<2>, NResultsInterface<0>])] +struct TestAllocCg1Op; +/// Test-only one-CTA deallocation used to reject a mismatched collective mode. +#[pliron_op(name = "nvvm.tcgen05_dealloc", format, verifier = "succ", interfaces = [NOpdsInterface<2>, NResultsInterface<0>])] +struct TestDeallocCg1Op; + +/// Append a direct unsigned 32-bit MIR literal for static-capacity verification tests. +fn literal_u32(ctx: &mut Context, block: Ptr, value: u32) -> Value { + let ty = IntegerType::get(ctx, 32, Signedness::Unsigned); + let op = Operation::new( + ctx, + MirConstantOp::get_concrete_op_info(), + vec![ty.into()], + vec![], + vec![], + 0, + ); + MirConstantOp::new(op).set_attr_value( + ctx, + IntegerAttr::new(ty, APInt::from_u32(value, NonZeroUsize::new(32).unwrap())), + ); + op.insert_at_back(block, ctx); + op.deref(ctx).get_result(0) +} + +/// Append a test-only canonical CTM/NVVM operation and include it in the audited operation list. +fn append_ctm_op( + ctx: &mut Context, + block: Ptr, + ops: &mut Vec>, + operands: Vec, +) { + let op = Operation::new(ctx, T::get_concrete_op_info(), vec![], operands, vec![], 0); + op.insert_at_back(block, ctx); + ops.push(op); +} + +/// Replace fixture allocation markers with configurable CTM lifecycle sites for compatibility and failure cases. +fn replace_cute_tmem_with_ctm( + ctx: &mut Context, + ops: &mut Vec>, + columns: u32, + group: u32, + dealloc: bool, + relinquish: bool, +) { + let block = ops[0].deref(ctx).get_parent_block().unwrap(); + let shared = ops[0].deref(ctx).get_operand(0); + let scalar = ops + .iter() + .find(|op| Operation::get_opid(**op, ctx) == CuteSm100TmemDeallocOp::get_opid_static()) + .unwrap() + .deref(ctx) + .get_operand(0); + let is_cute_tmem = |op| { + [ + CuteSm100TmemAllocOp::get_opid_static(), + CuteSm100TmemDeallocOp::get_opid_static(), + ] + .contains(&Operation::get_opid(op, ctx)) + }; + for op in ops.iter().copied().filter(|op| is_cute_tmem(*op)) { + op.unlink(ctx); + } + ops.retain(|op| !is_cute_tmem(*op)); + let count = literal_u32(ctx, block, columns); + if group == 2 { + append_ctm_op::(ctx, block, ops, vec![shared, count]); + if relinquish { + append_ctm_op::(ctx, block, ops, vec![]); + } + if dealloc { + append_ctm_op::(ctx, block, ops, vec![scalar, count]); + } + } else { + append_ctm_op::(ctx, block, ops, vec![shared, count]); + if dealloc { + append_ctm_op::(ctx, block, ops, vec![scalar, count]); + } + } +} + +/// Build the supported paired M256/N192+160/K64 plan with canonical shared layouts. +fn mma_plan() -> CuteSm100MmaPlanAttr { + CuteSm100MmaPlanAttr { + m: 256, + n0: 192, + n1: 160, + k: 64, + tmem_columns: 512, + cta_group: 2, + collector: CuteSm100CollectorAttr::AFillLastUse, + a_layout: CuteComposedLayoutAttr(row_major(128, 64, 3)), + b0_layout: CuteComposedLayoutAttr(row_major(96, 64, 3)), + b1_layout: CuteComposedLayoutAttr(row_major(80, 64, 3)), + } +} +/// Build an input or accumulator pipeline fixture from expected bytes and arrival count. +fn pipeline(bytes: u32, arrivals: u32) -> CuteSm100PipelinePlanAttr { + CuteSm100PipelinePlanAttr { + stages: if bytes == 0 { 1 } else { 5 }, + transaction_bytes: bytes, + consumer_arrivals: arrivals, + cta_group: 2, + leader_rank: 0, + multicast_mask: 3, + } +} +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, block: Ptr, ty: TypeHandle) -> Value { + let op = MirUndefOp::new(ctx, ty).get_operation(); + op.insert_at_back(block, ctx); + op.deref(ctx).get_result(0) +} +/// Construct one complete paired-SM100 lifecycle, returning the module and all semantic operation sites. +fn complete_story(ctx: &mut Context, bytes: u32) -> (Ptr, Vec>) { + dialect_mir::register(ctx); + crate::register(ctx); + let module = ModuleOp::new(ctx, "sm100_test".try_into().unwrap()).get_operation(); + let block = module + .deref(ctx) + .get_region(0) + .deref(ctx) + .iter(ctx) + .next() + .unwrap(); + let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let u8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let bool_ty: TypeHandle = IntegerType::get(ctx, 1, Signedness::Signless).into(); + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(ctx).into(); + let token_ty: TypeHandle = MirPtrType::get(ctx, u32_ty, true, address_space::SHARED).into(); + let barrier_ty: TypeHandle = MirPtrType::get(ctx, u64_ty, true, address_space::SHARED).into(); + let shared_ty: TypeHandle = MirPtrType::get(ctx, f16_ty, true, address_space::SHARED).into(); + let desc_ty: TypeHandle = MirPtrType::get_generic(ctx, u8_ty, false).into(); + let scalar = undef(ctx, block, u32_ty); + let condition = undef(ctx, block, bool_ty); + let bias = undef(ctx, block, f32_ty); + let token = undef(ctx, block, token_ty); + let full = undef(ctx, block, barrier_ty); + let empty = undef(ctx, block, barrier_ty); + let afull = undef(ctx, block, barrier_ty); + let aempty = undef(ctx, block, barrier_ty); + let shared = undef(ctx, block, shared_ty); + let desc = undef(ctx, block, desc_ty); + let mut ops = Vec::new(); + /// Append one typed semantic operation to the test block and collected operation list. + macro_rules! append { ($ty:ident,[$($arg:expr),*],$plan:expr) => {{ + let operation=$ty::new(ctx,vec![$($arg),*],$plan).get_operation(); + operation.insert_at_back(block,ctx);ops.push(operation); + }}; } + let tmem = CuteSm100TmemPlanAttr { + columns: 512, + cta_group: 2, + }; + let input = pipeline(bytes, 1); + let acc = pipeline(0, 8); + let epi = CuteSm100EpiloguePlanAttr { + layout: CuteComposedLayoutAttr(row_major(128, 32, 2)), + rows: 128, + columns: 32, + }; + append!(CuteSm100TmemAllocOp, [token], tmem); + append!(CuteSm100PipelineInitOp, [full, empty], input); + append!(CuteSm100AccumulatorInitOp, [aempty, afull], acc); + append!(CuteSm100PipelineAcquireOp, [empty, scalar, scalar], input); + append!(CuteSm100PipelineExpectOp, [full, scalar], input); + for rows in [128, 96, 80] { + append!( + CuteSm100ClusterTmaLoadOp, + [desc, shared, scalar, scalar, full, scalar], + CuteSm100ClusterCopyPlanAttr { + layout: CuteComposedLayoutAttr(row_major(rows, 64, 3)), + cta_group: 2, + leader_rank: 0 + } + ); + } + append!(CuteSm100AccumulatorAcquireOp, [aempty, scalar], acc); + append!(CuteSm100PipelineWaitOp, [full, scalar, scalar], input); + append!( + CuteSm100TiledMmaOp, + [scalar, shared, shared, shared, condition], + mma_plan() + ); + append!(CuteSm100PipelineReleaseOp, [empty, scalar], input); + append!(CuteSm100AccumulatorCommitOp, [afull], acc); + append!(CuteSm100AccumulatorWaitOp, [afull, scalar], acc); + append!( + CuteSm100TmemEpilogueOp, + [scalar, shared, scalar, bias], + epi.clone() + ); + append!(CuteSm100TmaStoreOp, [desc, shared, scalar, scalar], epi); + append!(CuteSm100StoreCommitOp, [], CuteTmaStorePipelineAttr::new(2)); + append!( + CuteSm100StoreAcquireOp, + [], + CuteTmaStorePipelineAttr::new(2) + ); + append!(CuteSm100AccumulatorReleaseOp, [aempty], acc); + append!(CuteSm100StoreTailOp, [], CuteTmaStorePipelineAttr::new(2)); + append!(CuteSm100TmemDeallocOp, [scalar], tmem); + (module, ops) +} + +/// The stage-byte promise includes A and both B partitions across both CTAs. +#[test] +fn paired_plan_derives_complete_two_cta_transaction_bytes() { + let ctx = Context::new(); + let plan = mma_plan(); + plan.verify(&ctx).unwrap(); + assert_eq!(plan.transaction_bytes(), Some(77_824)); + let mut bad = plan.clone(); + bad.b0_layout = plan.b1_layout.clone(); + assert!(bad.verify(&ctx).is_err()); + let mut bad = plan.clone(); + bad.n1 = 176; + assert!(bad.verify(&ctx).is_err()); + let mut bad = plan; + bad.cta_group = 1; + assert!(bad.verify(&ctx).is_err()); +} +/// Unsupported allocation capacities and collective routing configurations fail static verification. +#[test] +fn tmem_and_cluster_collective_contracts_fail_closed() { + let ctx = Context::new(); + assert!( + CuteSm100TmemPlanAttr { + columns: 48, + cta_group: 2 + } + .verify(&ctx) + .is_err() + ); + assert!( + CuteSm100TmemPlanAttr { + columns: 512, + cta_group: 1 + } + .verify(&ctx) + .is_err() + ); + let mut plan = pipeline(77_824, 1); + plan.multicast_mask = 1; + assert!(plan.verify(&ctx).is_err()); + let mut plan = pipeline(0, 8); + plan.stages = 2; + assert!(plan.verify(&ctx).is_err()); +} +/// The backend-neutral entry point accepts a complete matching paired-SM100 lifecycle. +#[test] +fn complete_sm100_lifecycle_is_verified_by_shared_entrypoint() { + let mut ctx = Context::new(); + let (module, _) = complete_story(&mut ctx, 77_824); + crate::verify::verify_cute_semantics(&ctx, module).unwrap(); +} +/// Canonical CTM two-CTA allocation/deallocation may supply storage for semantic CuTe MMA. +#[test] +fn accepts_ctm_group2_tmem_lifecycle_with_cute_mma() { + let mut ctx = Context::new(); + let (module, mut ops) = complete_story(&mut ctx, 77_824); + replace_cute_tmem_with_ctm(&mut ctx, &mut ops, 512, 2, true, true); + verify_story(&ctx, &ops).unwrap(); + crate::verify::verify_cute_semantics(&ctx, module).unwrap(); +} +/// A recognized allocation cannot omit the matching release of TMEM capacity. +#[test] +fn rejects_ctm_tmem_lifecycle_missing_deallocation() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + replace_cute_tmem_with_ctm(&mut ctx, &mut ops, 512, 2, false, true); + let error = verify_story(&ctx, &ops).unwrap_err(); + assert!(error.contains("dealloc"), "{error}"); +} +/// CTM allocation must relinquish its allocation permit before the lifecycle is accepted. +#[test] +fn rejects_ctm_tmem_lifecycle_missing_allocation_permit() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + replace_cute_tmem_with_ctm(&mut ctx, &mut ops, 512, 2, true, false); + let error = verify_story(&ctx, &ops).unwrap_err(); + assert!(error.contains("permit"), "{error}"); +} +/// One-CTA allocation/deallocation cannot satisfy a two-CTA MMA storage contract. +#[test] +fn rejects_ctm_tmem_lifecycle_with_group1_operations() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + replace_cute_tmem_with_ctm(&mut ctx, &mut ops, 512, 1, true, false); + let error = verify_story(&ctx, &ops).unwrap_err(); + assert!(error.contains("CTA group 2"), "{error}"); +} +/// A directly known CTM column count must match the paired MMA's TMEM capacity. +#[test] +fn rejects_ctm_tmem_lifecycle_with_known_wrong_capacity() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + replace_cute_tmem_with_ctm(&mut ctx, &mut ops, 256, 2, true, true); + let error = verify_story(&ctx, &ops).unwrap_err(); + assert!(error.contains("capacity"), "{error}"); +} +/// The input barrier must expect both CTAs' copies, not only the local tile bytes. +#[test] +fn rejects_only_one_ctas_expected_transaction_bytes() { + let mut ctx = Context::new(); + let (module, _) = complete_story(&mut ctx, 38_912); + let error = crate::verify::verify_cute_semantics(&ctx, module) + .unwrap_err() + .to_string(); + assert!( + error.contains("promises 38912 bytes") && error.contains("77824"), + "{error}" + ); +} +/// Every shared MMA operand placement must have a corresponding cluster input-copy layout. +#[test] +fn rejects_mma_input_layout_without_matching_cluster_copy() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + ops.retain(|op| Operation::get_opid(*op, &ctx) != CuteSm100ClusterTmaLoadOp::get_opid_static()); + assert!( + verify_story(&ctx, &ops) + .unwrap_err() + .contains("no matching cluster TMA copy") + ); +} +/// The epilogue must release its accumulator slot before the paired lifecycle is complete. +#[test] +fn rejects_missing_accumulator_release() { + let mut ctx = Context::new(); + let (_, mut ops) = complete_story(&mut ctx, 77_824); + ops.retain(|op| { + Operation::get_opid(*op, &ctx) != CuteSm100AccumulatorReleaseOp::get_opid_static() + }); + assert!( + verify_story(&ctx, &ops) + .unwrap_err() + .contains("sm100_accumulator_release") + ); +} +/// Malformed positional scalar types are rejected by the shared verifier before backend emission. +#[test] +fn rejects_bad_scalar_carriers_before_backend_lowering() { + let mut ctx = Context::new(); + let (_, ops) = complete_story(&mut ctx, 77_824); + let mma = ops + .iter() + .find(|op| Operation::get_opid(**op, &ctx) == CuteSm100TiledMmaOp::get_opid_static()) + .copied() + .unwrap(); + let scalar = mma.deref(&ctx).get_operand(0); + let malformed = CuteSm100TiledMmaOp::new(&mut ctx, vec![scalar; 5], mma_plan()); + assert!(malformed.verify(&ctx).is_err()); + let malformed = CuteSm100TmemAllocOp::new( + &mut ctx, + vec![], + CuteSm100TmemPlanAttr { + columns: 512, + cta_group: 2, + }, + ); + assert!(malformed.verify(&ctx).is_err()); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/verify.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/verify.rs new file mode 100644 index 0000000000..773bcb5791 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100/verify.rs @@ -0,0 +1,199 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Cross-operation layout, capacity, and collective lifecycle verification. + +use super::*; + +/// Find and clone the operation attribute of the requested static plan type. +fn static_plan(ctx: &Context, op: Ptr) -> Option { + op.deref(ctx) + .attributes + .0 + .values() + .find_map(|attr| attr.downcast_ref::().cloned()) +} + +// CTM lowers its typed group-2 TMEM guard into these canonical NVVM ops. +// Read the column count when MIR still carries it as a direct constant; the +// group-2 op name is the static group contract in all cases. +/// Read a direct literal column count from a canonical CTM/NVVM allocation operation. +fn nvvm_tmem_columns(ctx: &Context, op: Ptr) -> Option { + let value = op.deref(ctx).get_operand(1); + let constant = value.defining_op()?; + if Operation::get_opid(constant, ctx) != MirConstantOp::get_opid_static() { + return None; + } + let attribute = MirConstantOp::new(constant).get_attr_value(ctx)?; + (attribute.value().bw() <= 64).then(|| attribute.value().to_u64()) +} + +/// Check static contracts across the scalar-carrier SM100 graph in one function. +/// +/// This checks matching tile layouts, copy transaction sizes, TMEM capacity, +/// and complete lifecycle membership. Dynamic phase arithmetic, lane election, +/// and pointer aliasing remain the documented unsafe Rust preconditions. +pub fn verify_story(ctx: &Context, ops: &[Ptr]) -> Result<(), String> { + let id = |op| Operation::get_opid(op, ctx); + let find = |wanted| ops.iter().copied().find(|op| id(*op) == wanted); + let Some(mma_op) = find(CuteSm100TiledMmaOp::get_opid_static()) else { + return Ok(()); + }; + let mma = CuteSm100TiledMmaOp::wrap(mma_op) + .plan(ctx) + .ok_or("SM100 MMA plan missing")?; + let nvvm = |name: &str| { + ops.iter() + .copied() + .filter(|op| id(*op).to_string() == name) + .collect::>() + }; + let ctm_alloc = nvvm("nvvm.tcgen05_alloc_cg2"); + let ctm_dealloc = nvvm("nvvm.tcgen05_dealloc_cg2"); + if find(CuteSm100TmemAllocOp::get_opid_static()).is_none() && ctm_alloc.is_empty() { + if !nvvm("nvvm.tcgen05_alloc").is_empty() { + return Err("SM100 tiled MMA requires CTA group 2 TMEM allocation".into()); + } + return Err( + "SM100 tiled MMA is missing required lifecycle operation cute.sm100_tmem_alloc or nvvm.tcgen05_alloc_cg2" + .into(), + ); + } + if find(CuteSm100TmemDeallocOp::get_opid_static()).is_none() && ctm_dealloc.is_empty() { + if !nvvm("nvvm.tcgen05_dealloc").is_empty() { + return Err("SM100 tiled MMA requires CTA group 2 TMEM deallocation".into()); + } + return Err( + "SM100 tiled MMA is missing required lifecycle operation cute.sm100_tmem_dealloc or nvvm.tcgen05_dealloc_cg2" + .into(), + ); + } + if !ctm_alloc.is_empty() && nvvm("nvvm.tcgen05_relinquish_alloc_permit_cg2").is_empty() { + return Err( + "SM100 CTM group-2 TMEM allocation is missing its allocation permit relinquish".into(), + ); + } + for op in ctm_alloc.into_iter().chain(ctm_dealloc) { + if op.deref(ctx).get_num_operands() != 2 { + return Err("SM100 CTM TMEM lifecycle operation must have two operands".into()); + } + if nvvm_tmem_columns(ctx, op).is_some_and(|columns| columns != mma.tmem_columns as u64) { + return Err("SM100 allocation capacity disagrees with tiled MMA".into()); + } + } + for needed in [ + CuteSm100PipelineInitOp::get_opid_static(), + CuteSm100PipelineAcquireOp::get_opid_static(), + CuteSm100PipelineExpectOp::get_opid_static(), + CuteSm100PipelineWaitOp::get_opid_static(), + CuteSm100PipelineReleaseOp::get_opid_static(), + CuteSm100AccumulatorInitOp::get_opid_static(), + CuteSm100AccumulatorAcquireOp::get_opid_static(), + CuteSm100AccumulatorCommitOp::get_opid_static(), + CuteSm100AccumulatorWaitOp::get_opid_static(), + CuteSm100AccumulatorReleaseOp::get_opid_static(), + CuteSm100TmemEpilogueOp::get_opid_static(), + CuteSm100TmaStoreOp::get_opid_static(), + CuteSm100StoreCommitOp::get_opid_static(), + CuteSm100StoreAcquireOp::get_opid_static(), + CuteSm100StoreTailOp::get_opid_static(), + ] { + if find(needed.clone()).is_none() { + return Err(format!( + "SM100 tiled MMA is missing required lifecycle operation {needed}" + )); + } + } + let mut copied_layouts = Vec::new(); + let mut input_plan = None; + let mut accumulator_plan = None; + let mut store_plan = None; + let input_ids = [ + CuteSm100PipelineInitOp::get_opid_static(), + CuteSm100PipelineAcquireOp::get_opid_static(), + CuteSm100PipelineExpectOp::get_opid_static(), + CuteSm100PipelineWaitOp::get_opid_static(), + CuteSm100PipelineReleaseOp::get_opid_static(), + ]; + let accumulator_ids = [ + CuteSm100AccumulatorInitOp::get_opid_static(), + CuteSm100AccumulatorAcquireOp::get_opid_static(), + CuteSm100AccumulatorCommitOp::get_opid_static(), + CuteSm100AccumulatorWaitOp::get_opid_static(), + CuteSm100AccumulatorReleaseOp::get_opid_static(), + ]; + for &op in ops { + let opid = id(op); + if opid == CuteSm100TiledMmaOp::get_opid_static() + && CuteSm100TiledMmaOp::wrap(op).plan(ctx).as_ref() != Some(&mma) + { + return Err("SM100 tiled MMA plans in one function must agree".into()); + } + if opid == CuteSm100TmemAllocOp::get_opid_static() + || opid == CuteSm100TmemDeallocOp::get_opid_static() + { + let plan = + static_plan::(ctx, op).ok_or("SM100 TMEM plan missing")?; + if plan.columns != mma.tmem_columns || plan.cta_group != mma.cta_group { + return Err("SM100 allocation capacity disagrees with tiled MMA".into()); + } + } + if input_ids.contains(&opid) || accumulator_ids.contains(&opid) { + let plan = static_plan::(ctx, op) + .ok_or("SM100 pipeline plan missing")?; + let target = if input_ids.contains(&opid) { + &mut input_plan + } else { + &mut accumulator_plan + }; + if target.is_some_and(|previous| previous != plan) { + return Err("SM100 pipeline lifecycle plans disagree".into()); + } + *target = Some(plan); + if plan.transaction_bytes > 0 && Some(plan.transaction_bytes) != mma.transaction_bytes() + { + return Err(format!( + "SM100 pipeline promises {} bytes but the MMA operand tiles need {}", + plan.transaction_bytes, + mma.transaction_bytes().unwrap() + )); + } + if plan.transaction_bytes == 0 && plan.consumer_arrivals != 8 { + return Err( + "SM100 128x32 epilogue requires four consumer warps in each CTA (8 arrivals)" + .into(), + ); + } + } + if opid == CuteSm100ClusterTmaLoadOp::get_opid_static() { + copied_layouts.push( + CuteSm100ClusterTmaLoadOp::wrap(op) + .plan(ctx) + .ok_or("SM100 copy plan missing")? + .layout, + ); + } + if [ + CuteSm100StoreCommitOp::get_opid_static(), + CuteSm100StoreAcquireOp::get_opid_static(), + CuteSm100StoreTailOp::get_opid_static(), + ] + .contains(&opid) + { + let plan = static_plan::(ctx, op) + .ok_or("SM100 store plan missing")?; + if store_plan.is_some_and(|previous| previous != plan) { + return Err("SM100 store lifecycle stage counts disagree".into()); + } + store_plan = Some(plan); + } + } + for required in [&mma.a_layout, &mma.b0_layout, &mma.b1_layout] { + if !copied_layouts.contains(required) { + return Err("SM100 MMA operand layout has no matching cluster TMA copy".into()); + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100_attention.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100_attention.rs new file mode 100644 index 0000000000..1a412c8435 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/sm100_attention.rs @@ -0,0 +1,405 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent SM100 MMA and rank-three tiled-copy contracts. +//! +//! Each operation owns its layout/shape contract. QK and PV may have different +//! shapes and layouts within one kernel. Pointer lifetimes, barrier phases and +//! asynchronous alias handoffs are unsafe source obligations, not implied by +//! the order of operations in a single thread's IR. +//! +//! MMA shared operands carry allocation bases plus trailing u32 byte offsets +//! (A then B for SS, B for TS). Offsets select complete tiles and must preserve +//! the layout's full swizzle period. Their validity within the live allocation +//! is an unsafe source obligation, like shared-pointer validity; the typed API +//! constructs these offsets only from bounded, complete-tile array indices. +//! +//! Every operation has zero SSA results; the operand lists below describe +//! `new`'s positional values. TMEM operands are encoded `u32` addresses. +//! TMA's three coordinates address the descriptor's physical axes: within-span +//! column, row, and column-span index. For a 64-element span, physical +//! `(row, column) = (7, 128)` becomes `(c0, c1, c2) = (0, 7, 2)`. +use crate::{ + attributes::CuteComposedLayoutAttr, + layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}, +}; +use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; +use pliron::{ + attribute::Attribute, + builtin::types::IntegerType, + common_traits::Verify, + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + result::Error, + r#type::Typed, + value::Value, + verify_err, verify_err_noloc, +}; +use pliron_derive::{pliron_attr, pliron_op}; + +/// Canonical tiled FP16 layout, with all offsets expressed in elements. +/// +/// `mn = false` makes K/columns contiguous; `true` makes M-or-N/rows contiguous. +/// The contiguous extent is 32 or a positive multiple of 64, and the other +/// extent is a positive multiple of eight. A span of 32 elements uses a +/// 64-byte swizzle; larger extents use tiled spans of 64 and a 128-byte swizzle. +/// +/// ``` +/// use dialect_cute::ops::sm100_attention::{CuteSm100MmaAttr, shared_layout}; +/// use pliron::{common_traits::Verify, context::Context}; +/// +/// // A collective M128/N128/K64 product: each CTA stores 64 rows of A and B. +/// let layout = shared_layout(64, 64, false).unwrap(); +/// let plan = CuteSm100MmaAttr { +/// m: 128, n: 128, k: 64, cta_group: 2, +/// a_mn: false, b_mn: false, +/// a_layout: layout.clone(), b_layout: layout, +/// }; +/// plan.verify(&Context::new()).unwrap(); +/// ``` +pub fn shared_layout(rows: u32, cols: u32, mn: bool) -> Result { + let (contiguous, outer) = if mn { (rows, cols) } else { (cols, rows) }; + if contiguous < 32 + || (contiguous != 32 && !contiguous.is_multiple_of(64)) + || outer == 0 + || !outer.is_multiple_of(8) + { + return Err( + "SM100 shared layout needs contiguous 32 or a multiple of 64 and outer multiple of 8" + .into(), + ); + } + let span = contiguous.min(64); + let (shape, stride) = if mn { + ( + format!("(({span},{}),(8,{}))", rows / span, cols / 8), + format!("((1,{}),({span},{}))", span * cols, span * 8), + ) + } else { + ( + format!("((8,{}),({span},{}))", rows / 8, cols / span), + format!("(({span},{}),(1,{}))", span * 8, span * rows), + ) + }; + let outer: Layout = format!("{shape}:{stride}") + .parse() + .map_err(|e| format!("{e}"))?; + ComposedLayout::new( + Swizzle::new(if span == 32 { 2 } else { 3 }, 3, 3), + 0, + outer, + OffsetUnit::Elements, + ) + .map(CuteComposedLayoutAttr) + .map_err(|e| e.to_string()) +} + +/// One independent FP16 MMA product with FP32 accumulation. +/// +/// Dimensions describe the whole CTA group: `M = 128`, N is a multiple of 16 +/// in `16..=256`, and K is a positive multiple of 16. Each CTA's shared A/B +/// layouts have `M / cta_group` and `N / cta_group` rows, respectively. +/// `a_mn`/`b_mn` select contiguous rows instead of contiguous K; both layouts +/// must match [`shared_layout`]. TMEM-A operations additionally require group +/// one and K-major A; their Rust importer accepts K=32 or a multiple of 64. +#[pliron_attr( + name = "cute.sm100_mma", + format = "`<` $m `,` $n `,` $k `,` $cta_group `,` $a_mn `,` $b_mn `,` $a_layout `,` $b_layout `>`" +)] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100MmaAttr { + /// Collective output-row extent M. + pub m: u32, + /// Collective output-column extent N. + pub n: u32, + /// Reduction extent K, in logical elements. + pub k: u32, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, + /// Whether A's M dimension, rather than K, is physically contiguous. + pub a_mn: bool, + /// Whether B's N dimension, rather than K, is physically contiguous. + pub b_mn: bool, + /// Physical shared placement of each CTA's A tile. + pub a_layout: CuteComposedLayoutAttr, + /// Physical shared placement of each CTA's B tile. + pub b_layout: CuteComposedLayoutAttr, +} +impl Verify for CuteSm100MmaAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + if !matches!(self.cta_group, 1 | 2) + || self.m != 128 + || !(16..=256).contains(&self.n) + || !self.n.is_multiple_of(16) + || self.k == 0 + || !self.k.is_multiple_of(16) + { + return verify_err_noloc!( + "SM100 independent FP16 MMA requires M128, N16..256 by16, K by16, CTA group1 or2" + ); + } + for (layout, rows, mn) in [ + (&self.a_layout, self.m / self.cta_group, self.a_mn), + (&self.b_layout, self.n / self.cta_group, self.b_mn), + ] { + layout.verify(ctx)?; + let expected = + shared_layout(rows, self.k, mn).map_err(|e| pliron::verify_error_noloc!("{e}"))?; + if *layout != expected { + return verify_err_noloc!( + "SM100 MMA physical operand layout disagrees with shape or major order" + ); + } + } + Ok(()) + } +} +/// One FP16 TMA tile's physical shared placement and participating CTA group. +/// +/// `layout` is one CTA's tile, in element offsets, as produced by +/// [`shared_layout`]. Group one or two is supported; the descriptor's three +/// axes encode the tile's contiguous span, rows, and remaining columns. +#[pliron_attr(name = "cute.sm100_tma", format = "`<` $layout `,` $cta_group `>`")] +#[derive(PartialEq, Eq, Clone, Debug, Hash)] +pub struct CuteSm100TmaAttr { + /// Static coordinate layout used to interpret the storage. + pub layout: CuteComposedLayoutAttr, + /// Number of CTAs participating in the collective contract. + pub cta_group: u32, +} +impl Verify for CuteSm100TmaAttr { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.layout.verify(ctx)?; + if !matches!(self.cta_group, 1 | 2) { + return verify_err_noloc!("SM100 TMA requires CTA group1 or2"); + } + let sw = self.layout.0.inner(); + if !matches!(sw.num_bits, 2 | 3) + || sw.num_base != 3 + || sw.num_shift != 3 + || self.layout.0.offset() != 0 + { + return verify_err_noloc!("SM100 TMA requires zero-offset tiled B64/B128 FP16 layout"); + } + let modes = self.layout.0.outer().modes(); + if modes.len() != 2 { + return verify_err_noloc!("SM100 TMA layout must have two logical modes"); + } + let extent = |index: usize| { + modes[index] + .checked_size() + .and_then(|n| u32::try_from(n).ok()) + }; + let (Some(rows), Some(cols)) = (extent(0), extent(1)) else { + return verify_err_noloc!("SM100 TMA layout extents must fit u32"); + }; + if ![false, true] + .into_iter() + .any(|mn| shared_layout(rows, cols, mn).is_ok_and(|canonical| canonical == self.layout)) + { + return verify_err_noloc!( + "SM100 TMA layout must match a canonical FP16 tiled shared layout" + ); + } + Ok(()) + } +} +/// Scalar and address-space requirements for one independent attention operand. +#[derive(Clone, Copy)] +enum Carrier { + /// Unsigned 32-bit index or encoded tensor-memory address. + U32, + /// Signed 32-bit physical tensor-map coordinate. + I32, + /// One-bit predicate selecting whether MMA retains the previous accumulator. + Bool, + /// FP16 pointer in CTA shared memory, with either read-only or writable access. + Shared, + /// Writable FP16 pointer in CTA shared memory. + SharedWrite, + /// Immutable tensor-map pointer in generic, global, or constant space. + Descriptor, + /// Writable CTA-shared pointer to a 64-bit completion barrier. + Barrier, +} +/// Check scalar width or shared/descriptor pointer requirements for an operand. +fn carrier(ctx: &Context, value: Value, want: Carrier) -> bool { + let ty = value.get_type(ctx).deref(ctx); + match want { + Carrier::U32 | Carrier::I32 | Carrier::Bool => ty + .downcast_ref::() + .is_some_and(|i| i.width() == if matches!(want, Carrier::Bool) { 1 } else { 32 }), + _ => { + let Some(p) = ty.downcast_ref::() else { + return false; + }; + if matches!(want, Carrier::Descriptor) { + return !p.is_mutable + && matches!( + p.address_space, + address_space::GENERIC | address_space::GLOBAL | address_space::CONSTANT + ); + } + if p.address_space != address_space::SHARED { + return false; + } + if matches!(want, Carrier::Barrier) { + return p.is_mutable + && p.pointee + .deref(ctx) + .downcast_ref::() + .is_some_and(|i| i.width() == 64); + } + (!matches!(want, Carrier::SharedWrite) || p.is_mutable) + && p.pointee.deref(ctx).downcast_ref::().is_some() + } + } +} +/// Generate independent zero-result MMA/TMA operations with typed plans and carrier checks. +macro_rules! operation{ + ($(#[$doc:meta])* $ty:ident,$name:literal,$key:ident,$attr:ty,[$($arg:ident),*])=>{ + #[pliron_op(name=$name,format,attributes=($key:$attr))]pub struct $ty; + // Keep contracts on the impl and constructor: pliron's generated + // attribute-table links do not compose correctly with local struct docs. + $(#[$doc])* + impl $ty{ + $(#[$doc])* + pub fn new(ctx:&mut Context,operands:Vec,plan:$attr)->Self{let op=Self{op:Operation::new(ctx,Self::get_concrete_op_info(),vec![],operands,vec![],0)};op.get_operation().deref_mut(ctx).attributes.set(stringify!($key).try_into().unwrap(),plan);op} + /// View an existing operation without allocating or verifying it. + pub fn wrap(op:Ptr)->Self{Self{op}} + /// Clone the typed MMA/TMA plan, or return `None` if the attribute is missing. + pub fn plan(&self,ctx:&Context)->Option<$attr>{self.get_operation().deref(ctx).attributes.get::<$attr>(&stringify!($key).try_into().unwrap()).cloned()} + } + impl Verify for $ty{ + fn verify(&self,ctx:&Context)->Result<(),Error>{ + let op=self.get_operation().deref(ctx);let expected=&[$(Carrier::$arg),*]; + if op.get_num_results()!=0||op.get_num_operands()!=expected.len(){return verify_err!(op.loc(),"{} has invalid operand/result count",$name);} + let Some(plan)=self.plan(ctx)else{return verify_err!(op.loc(),"{} lost its typed plan",$name)};plan.verify(ctx)?; + for (i,c) in expected.iter().enumerate(){if !carrier(ctx,op.get_operand(i),*c){return verify_err!(op.loc(),"{} operand {} violates its carrier/address-space contract",$name,i);}} + Ok(()) + } + } + } +} +operation!( + /// Issue shared A × shared B into FP32 TMEM from one elected leader lane. + /// + /// Operands: `[destination, a_base, b_base, accumulate, a_bytes, b_bytes]`. + /// Byte offsets select complete shared tiles and preserve the full swizzle + /// period. `accumulate` controls the first K step; later steps accumulate. + /// Issue after input completion and keep all allocations live until MMA completes. + CuteSm100MmaSsOp, + "cute.sm100_mma_ss", + sm100_attention_mma_ss_plan, + CuteSm100MmaAttr, + [U32, Shared, Shared, Bool, U32, U32] +); +operation!( + /// Issue TMEM A × shared B into distinct FP32 TMEM from one elected lane. + /// + /// Operands: `[destination, a_tmem, b_base, accumulate, b_bytes]`. + /// A contains packed FP16 values; B's byte offset selects a complete tile. + /// A and destination must not overlap. Inputs must be visible to MMA, and + /// all allocations must remain live until completion. + CuteSm100MmaTsOp, + "cute.sm100_mma_ts", + sm100_attention_mma_ts_plan, + CuteSm100MmaAttr, + [U32, U32, Shared, Bool, U32] +); +operation!( + /// Start a rank-three global-to-shared transfer from all producer-warp lanes. + /// + /// Operands: `[descriptor, destination, c0, c1, c2, barrier, rank]`. + /// Coordinates are signed element indices on the descriptor's physical + /// axes. Rank is the local CTA rank; group-two completion routes to rank + /// zero, whose barrier expectation covers the whole group's copies. + CuteSm100TmaLoadOp, + "cute.sm100_attention_tma_load", + sm100_attention_tma_load_plan, + CuteSm100TmaAttr, + [Descriptor, SharedWrite, I32, I32, I32, Barrier, U32] +); +operation!( + /// Start a rank-three shared-to-global transfer from one elected issuer. + /// + /// Operands: `[descriptor, source, c0, c1, c2]`, using signed element + /// coordinates on the descriptor's physical axes. Shared writes must be + /// visible to the async proxy; keep the source live and unchanged until + /// its committed store group finishes reading it. + CuteSm100TmaStoreOp, + "cute.sm100_attention_tma_store", + sm100_attention_tma_store_plan, + CuteSm100TmaAttr, + [Descriptor, Shared, I32, I32, I32] +); +/// Apply a callback to the closed independent attention operation set. +macro_rules! each_op { + ($m:ident) => { + $m!( + CuteSm100MmaSsOp, + CuteSm100MmaTsOp, + CuteSm100TmaLoadOp, + CuteSm100TmaStoreOp + ) + }; +} +/// Register independent attention operations and their typed MMA/TMA attributes. +pub fn register(ctx: &mut Context) { + /// Register each listed attention operation in the supplied IR context. + macro_rules! reg{($($t:ty),*)=>{$(<$t>::register(ctx);)*}} + each_op!(reg); + CuteSm100MmaAttr::register(ctx); + CuteSm100TmaAttr::register(ctx); +} +/// Return the operation IDs included in the independent attention family. +pub fn semantic_ids() -> Vec { + /// Collect the registered operation IDs from the supplied type list. + macro_rules! ids{($($t:ty),*)=>{vec![$(<$t>::get_opid_static()),*]}} + each_op!(ids) +} +/// Run this family's local verifier, returning `None` for an operation from another family. +pub fn verify_local(ctx: &Context, op: Ptr) -> Option> { + let id = Operation::get_opid(op, ctx); + /// Dispatch a family-local verifier when the operation ID matches. + macro_rules! check{($($t:ty),*)=>{$(if id==<$t>::get_opid_static(){return Some(<$t>::wrap(op).verify(ctx));})*}} + each_op!(check); + None +} +/// Enforce one hardware CTA-group mode while permitting distinct per-resource +/// shapes within one function. Unlike the paired GEMM family, QK does not +/// constrain PV's layout. Group consistency across helper calls and allocation +/// lifetimes remain unsafe caller obligations. +pub fn verify_story(ctx: &Context, ops: &[Ptr]) -> Result<(), String> { + let mut group = None; + for &op in ops { + let id = Operation::get_opid(op, ctx); + let plan = if id == CuteSm100MmaSsOp::get_opid_static() { + CuteSm100MmaSsOp::wrap(op).plan(ctx) + } else if id == CuteSm100MmaTsOp::get_opid_static() { + CuteSm100MmaTsOp::wrap(op).plan(ctx) + } else { + None + }; + if let Some(p) = plan { + if group.is_some_and(|g| g != p.cta_group) { + return Err("SM100 MMA operations in one kernel must agree on CTA group".into()); + } + group = Some(p.cta_group); + if id == CuteSm100MmaTsOp::get_opid_static() { + if p.cta_group != 1 || p.a_mn { + return Err( + "SM100 TMEM-A attention supports group1 K-major probabilities".into(), + ); + } + let o = op.deref(ctx); + if o.get_operand(0) == o.get_operand(1) { + return Err("SM100 TMEM-A source cannot be its destination".into()); + } + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/mod.rs new file mode 100644 index 0000000000..ce6769101b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/mod.rs @@ -0,0 +1,797 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! High-level operations for one-dimensional tensor views. +//! +//! These operations preserve the story written by the kernel: +//! +//! ```text +//! pointer + length +//! │ +//! ▼ +//! tensor_make → tensor_zipped_divide → tensor_slice +//! │ +//! ┌───────────────┼───────────────┐ +//! ▼ ▼ ▼ +//! is_full/base load/store tail store +//! ``` +//! +//! They do not prescribe CUDA instructions or MLIR spelling. A backend first +//! chooses how to lower the preserved view. + +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::builtin::{ + op_interfaces::{NOpdsInterface, NResultsInterface, OneOpdInterface, OneResultInterface}, + types::{IntegerType, Signedness}, +}; +use pliron::common_traits::Verify; +use pliron::context::{Context, Ptr}; +use pliron::op::Op; +use pliron::operation::Operation; +use pliron::result::Error; +use pliron::r#type::{TypeHandle, Typed}; +use pliron::value::Value; +use pliron::verify_err; +use pliron_derive::pliron_op; + +use crate::attributes::{ + CuteAlignmentAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorLayoutAttr, +}; +use crate::types::CuteTensorViewType; + +/// Read tensor-view metadata from an SSA operand's type. +fn view_of(ctx: &Context, value: Value) -> Option { + let ty = value.get_type(ctx); + ty.deref(ctx).downcast_ref::().cloned() +} + +/// Preserve tensor storage facts while replacing its static selection layout. +fn transformed_view_type( + ctx: &mut Context, + view: Value, + layout: CuteTensorLayoutAttr, +) -> TypeHandle { + let source = + view_of(ctx, view).expect("cute tensor-view builder needs a CuteTensorViewType input"); + source.with_layout(ctx, layout).into() +} + +/// Check that the SSA value has an unsigned 64-bit integer type. +fn is_u64(ctx: &Context, value: Value) -> bool { + let ty = value.get_type(ctx); + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) +} + +/// Check for a one-bit integer predicate type. +fn is_i1(ctx: &Context, ty: TypeHandle) -> bool { + ty.deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 1) +} + +/// Compare tensor storage and access facts independently of the selection layout. +fn same_view_facts(left: &CuteTensorViewType, right: &CuteTensorViewType) -> bool { + left.logical == right.logical + && left.storage == right.storage + && left.format == right.format + && left.role == right.role + && left.space == right.space + && left.access == right.access + && left.alignment == right.alignment +} + +/// Check the register carrier's pointer pointee and required write capability. +fn carrier_pointer(ctx: &Context, value: Value, pointee: TypeHandle, needs_write: bool) -> bool { + let ty = value.get_type(ctx); + let ty_ref = ty.deref(ctx); + let Some(pointer) = ty_ref.downcast_ref::() else { + return false; + }; + pointer.pointee == pointee + && (!needs_write || pointer.is_mutable) + && [address_space::GENERIC, address_space::LOCAL].contains(&pointer.address_space) +} + +/// Require a selected tile whose byte width and alignment support one vector transfer. +fn valid_full_tile_transfer( + view: &CuteTensorViewType, + ctx: &Context, + alignment_bytes: u64, +) -> bool { + let Some(tile_size) = view.selected_tile_size() else { + return false; + }; + let Some(storage_bytes) = view.storage_bytes(ctx) else { + return false; + }; + let Some(tile_bytes) = tile_size.checked_mul(storage_bytes) else { + return false; + }; + matches!(tile_bytes, 4 | 8 | 16) + && alignment_bytes.is_power_of_two() + && alignment_bytes >= tile_bytes +} + +/// Start a contiguous tensor view from a data pointer and element count. +/// +/// ```text +/// %data, %len ──► cute.tensor_make ──► contiguous tensor view +/// ``` +/// +/// The result is a ghost value: this operation does not load or allocate. +#[pliron_op( + name = "cute.tensor_make", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteTensorMakeOp; + +impl CuteTensorMakeOp { + /// Build a global-memory tensor view. `logical` and `storage` are kept as + /// separate arguments even though v0 requires them to match. + #[allow(clippy::too_many_arguments)] + pub fn new( + ctx: &mut Context, + data: Value, + len: Value, + logical: TypeHandle, + storage: TypeHandle, + access: CuteTensorAccessAttr, + alignment_bytes: u64, + ) -> Self { + let view: TypeHandle = CuteTensorViewType::get( + ctx, + logical, + storage, + CuteTensorAddressSpaceAttr::Gmem, + access, + alignment_bytes, + CuteTensorLayoutAttr::Contiguous1D, + ) + .into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![view], + vec![data, len], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorMakeOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tensor_make needs 2 operands and 1 result"); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(view) = result_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tensor_make result must be a tensor view"); + }; + if view.layout != CuteTensorLayoutAttr::Contiguous1D { + return verify_err!(op.loc(), "cute.tensor_make result must use Contiguous1D"); + } + if view.space != CuteTensorAddressSpaceAttr::Gmem { + return verify_err!(op.loc(), "cute.tensor_make result must be global memory"); + } + view.verify(ctx)?; + + let data_ty = op.get_operand(0).get_type(ctx); + let data_ty_ref = data_ty.deref(ctx); + let Some(data) = data_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tensor_make data must be a MIR pointer"); + }; + if data.pointee != view.storage { + return verify_err!( + op.loc(), + "cute.tensor_make data pointee must match tensor storage" + ); + } + if ![address_space::GENERIC, address_space::GLOBAL].contains(&data.address_space) { + return verify_err!( + op.loc(), + "cute.tensor_make data must point at global or generic memory" + ); + } + if view.access == CuteTensorAccessAttr::ReadWrite && !data.is_mutable { + return verify_err!( + op.loc(), + "a writable cute.tensor_make view needs a mutable data pointer" + ); + } + if !is_u64(ctx, op.get_operand(1)) { + return verify_err!( + op.loc(), + "cute.tensor_make length must be an unsigned 64-bit integer" + ); + } + Ok(()) + } +} + +/// Split one contiguous row into fixed-width tiles. +/// +/// ```text +/// [0 1 2 3 4 5 6 7] ── zipped_divide<4> ──► [0 1 2 3] [4 5 6 7] +/// ``` +/// +/// Only the layout changes. Element, storage, space, access, and alignment +/// must stay exactly the same. +#[pliron_op( + name = "cute.tensor_zipped_divide", + format, + interfaces = [ + NOpdsInterface<1>, + NResultsInterface<1>, + OneOpdInterface, + OneResultInterface + ], + attributes = (zipped_layout: CuteTensorLayoutAttr) +)] +pub struct CuteTensorZippedDivideOp; + +impl CuteTensorZippedDivideOp { + /// Build a detached subdivision into fixed-width element tiles without accessing memory. + pub fn new(ctx: &mut Context, tensor: Value, tile_size: u64) -> Self { + let layout = CuteTensorLayoutAttr::Zipped1D(tile_size); + let result_ty = transformed_view_type(ctx, tensor, layout); + let divide = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![tensor], + vec![], + 0, + ), + }; + divide.set_attr_zipped_layout(ctx, layout); + divide + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorZippedDivideOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide needs 1 operand and 1 result" + ); + } + let Some(input) = view_of(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide input must be a tensor view" + ); + }; + if input.layout != CuteTensorLayoutAttr::Contiguous1D { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide input must use Contiguous1D" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide result must be a tensor view" + ); + }; + let Some(layout) = self.get_attr_zipped_layout(ctx).map(|layout| *layout) else { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide must carry its zipped layout" + ); + }; + if !matches!(layout, CuteTensorLayoutAttr::Zipped1D(size) if size > 0) { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide layout must be Zipped1D with a positive tile size" + ); + } + if result.layout != layout { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide result layout must match its layout attribute" + ); + } + if !same_view_facts(&input, result) { + return verify_err!( + op.loc(), + "cute.tensor_zipped_divide may change only the tensor layout" + ); + } + result.verify(ctx) + } +} + +/// Select tile number `tile_index` from a zipped view. +/// +/// ```text +/// zipped tiles: [tile 0] [tile 1] [tile 2] +/// ^ slice(..., 1) +/// ``` +#[pliron_op( + name = "cute.tensor_slice", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<1>, OneResultInterface] +)] +pub struct CuteTensorSliceOp; + +impl CuteTensorSliceOp { + /// Build a detached selection of one tile, retaining the original tensor storage facts. + pub fn new(ctx: &mut Context, tensor: Value, tile_index: Value) -> Self { + let input = view_of(ctx, tensor) + .expect("cute.tensor_slice builder needs a CuteTensorViewType input"); + let tile_size = match input.layout { + CuteTensorLayoutAttr::Zipped1D(size) => size, + _ => panic!("cute.tensor_slice builder needs a Zipped1D input"), + }; + let result_ty = input + .with_layout(ctx, CuteTensorLayoutAttr::Tile1D(tile_size)) + .into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![tensor, tile_index], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorSliceOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tensor_slice needs 2 operands and 1 result"); + } + let Some(input) = view_of(ctx, op.get_operand(0)) else { + return verify_err!(op.loc(), "cute.tensor_slice input must be a tensor view"); + }; + let CuteTensorLayoutAttr::Zipped1D(tile_size) = input.layout else { + return verify_err!(op.loc(), "cute.tensor_slice input must use Zipped1D"); + }; + if tile_size == 0 { + return verify_err!(op.loc(), "cute.tensor_slice tile size must be positive"); + } + if !is_u64(ctx, op.get_operand(1)) { + return verify_err!( + op.loc(), + "cute.tensor_slice tile index must be an unsigned 64-bit integer" + ); + } + let result_ty = op.get_result(0).get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(result) = result_ty_ref.downcast_ref::() else { + return verify_err!(op.loc(), "cute.tensor_slice result must be a tensor view"); + }; + if result.layout != CuteTensorLayoutAttr::Tile1D(tile_size) { + return verify_err!( + op.loc(), + "cute.tensor_slice result tile width must match its zipped input" + ); + } + if !same_view_facts(&input, result) { + return verify_err!( + op.loc(), + "cute.tensor_slice may change only the tensor layout" + ); + } + result.verify(ctx) + } +} + +/// Ask whether every position in a selected tile is in bounds. +#[pliron_op( + name = "cute.tensor_is_full", + format, + interfaces = [ + NOpdsInterface<1>, + NResultsInterface<1>, + OneOpdInterface, + OneResultInterface + ] +)] +pub struct CuteTensorIsFullOp; + +impl CuteTensorIsFullOp { + /// Build a detached predicate checking whether the selected tile fits the runtime extent. + pub fn new(ctx: &mut Context, tensor: Value) -> Self { + let result_ty: TypeHandle = IntegerType::get(ctx, 1, Signedness::Signless).into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![tensor], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorIsFullOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tensor_is_full needs 1 operand and 1 result"); + } + let Some(view) = view_of(ctx, op.get_operand(0)) else { + return verify_err!(op.loc(), "cute.tensor_is_full input must be a tensor view"); + }; + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(size) if size > 0) { + return verify_err!( + op.loc(), + "cute.tensor_is_full input must be a selected Tile1D" + ); + } + if !is_i1(ctx, op.get_result(0).get_type(ctx)) { + return verify_err!(op.loc(), "cute.tensor_is_full result must be i1"); + } + Ok(()) + } +} + +/// Return the selected tile's first absolute element index. +#[pliron_op( + name = "cute.tensor_base", + format, + interfaces = [ + NOpdsInterface<1>, + NResultsInterface<1>, + OneOpdInterface, + OneResultInterface + ] +)] +pub struct CuteTensorBaseOp; + +impl CuteTensorBaseOp { + /// Build a detached query for the selected tile's absolute element offset. + pub fn new(ctx: &mut Context, tensor: Value) -> Self { + let result_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty], + vec![tensor], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorBaseOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 1 || op.get_num_results() != 1 { + return verify_err!(op.loc(), "cute.tensor_base needs 1 operand and 1 result"); + } + let Some(view) = view_of(ctx, op.get_operand(0)) else { + return verify_err!(op.loc(), "cute.tensor_base input must be a tensor view"); + }; + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(size) if size > 0) { + return verify_err!(op.loc(), "cute.tensor_base input must be a selected Tile1D"); + } + if !is_u64(ctx, op.get_result(0)) { + return verify_err!( + op.loc(), + "cute.tensor_base result must be an unsigned 64-bit integer" + ); + } + Ok(()) + } +} + +/// Load one full tile into a thread-local register carrier. +/// +/// The operation keeps the tile as an operand. A backend can lower it to an +/// aligned native vector copy or a tensor-to-register mapping. +/// +/// The verifier checks the static tile width and the alignment promise. At +/// runtime, the caller must also prove both facts below before executing it: +/// +/// ```text +/// tensor_is_full(tile) == true +/// tile address % promised alignment == 0 +/// ``` +#[pliron_op( + name = "cute.tensor_load_into", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<0>], + attributes = (load_alignment_bytes: CuteAlignmentAttr) +)] +pub struct CuteTensorLoadIntoOp; + +impl CuteTensorLoadIntoOp { + /// Build a detached full-tile transfer into registers with a selected-address alignment promise. + pub fn new( + ctx: &mut Context, + tensor: Value, + destination: Value, + assumed_alignment_bytes: u64, + ) -> Self { + let load = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![tensor, destination], + vec![], + 0, + ), + }; + load.set_attr_load_alignment_bytes(ctx, CuteAlignmentAttr(assumed_alignment_bytes)); + load + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorLoadIntoOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tensor_load_into needs 2 operands and 0 results" + ); + } + let Some(view) = view_of(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.tensor_load_into source must be a tensor view" + ); + }; + if view.access != CuteTensorAccessAttr::ReadOnly { + return verify_err!(op.loc(), "cute.tensor_load_into source must be read-only"); + } + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(size) if size > 0) { + return verify_err!( + op.loc(), + "cute.tensor_load_into source must be a selected Tile1D" + ); + } + if !carrier_pointer(ctx, op.get_operand(1), view.storage, true) { + return verify_err!( + op.loc(), + "cute.tensor_load_into destination must be a mutable local/generic pointer to storage" + ); + } + let Some(alignment) = self + .get_attr_load_alignment_bytes(ctx) + .map(|attribute| attribute.0) + else { + return verify_err!( + op.loc(), + "cute.tensor_load_into must carry an alignment promise" + ); + }; + if !valid_full_tile_transfer(&view, ctx, alignment) { + return verify_err!( + op.loc(), + "cute.tensor_load_into needs a supported full-tile width and sufficient power-of-two alignment" + ); + } + Ok(()) + } +} + +/// Store one thread-local register carrier into a full writable tile. +/// +/// The verifier checks the static tile width, writable access, and alignment +/// promise. At runtime the selected tile must be full, its address must meet +/// that promise, and no conflicting memory access may race this store. +#[pliron_op( + name = "cute.tensor_store_from", + format, + interfaces = [NOpdsInterface<2>, NResultsInterface<0>], + attributes = (store_alignment_bytes: CuteAlignmentAttr) +)] +pub struct CuteTensorStoreFromOp; + +impl CuteTensorStoreFromOp { + /// Operand order follows the data flow: register carrier, then tensor. + pub fn new( + ctx: &mut Context, + source: Value, + tensor: Value, + assumed_alignment_bytes: u64, + ) -> Self { + let store = Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![source, tensor], + vec![], + 0, + ), + }; + store.set_attr_store_alignment_bytes(ctx, CuteAlignmentAttr(assumed_alignment_bytes)); + store + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorStoreFromOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 2 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tensor_store_from needs 2 operands and 0 results" + ); + } + let Some(view) = view_of(ctx, op.get_operand(1)) else { + return verify_err!( + op.loc(), + "cute.tensor_store_from destination must be a tensor view" + ); + }; + if view.access != CuteTensorAccessAttr::ReadWrite { + return verify_err!( + op.loc(), + "cute.tensor_store_from destination must be writable" + ); + } + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(size) if size > 0) { + return verify_err!( + op.loc(), + "cute.tensor_store_from destination must be a selected Tile1D" + ); + } + if !carrier_pointer(ctx, op.get_operand(0), view.storage, false) { + return verify_err!( + op.loc(), + "cute.tensor_store_from source must be a local/generic pointer to storage" + ); + } + let Some(alignment) = self + .get_attr_store_alignment_bytes(ctx) + .map(|attribute| attribute.0) + else { + return verify_err!( + op.loc(), + "cute.tensor_store_from must carry an alignment promise" + ); + }; + if !valid_full_tile_transfer(&view, ctx, alignment) { + return verify_err!( + op.loc(), + "cute.tensor_store_from needs a supported full-tile width and sufficient power-of-two alignment" + ); + } + Ok(()) + } +} + +/// Store one scalar at an absolute index inside a selected writable tile. +/// +/// This is the short-tail path. Keeping it explicit lets both backends see +/// that the scalar store still belongs to the same selected tensor tile. +/// The absolute index must be inside both the selected tile and the tensor's +/// valid element range. No conflicting memory access may race this store. +#[pliron_op( + name = "cute.tensor_store_element_abs", + format, + interfaces = [NOpdsInterface<3>, NResultsInterface<0>] +)] +pub struct CuteTensorStoreElementAbsOp; + +impl CuteTensorStoreElementAbsOp { + /// Build a detached scalar-tail store retaining its selected tile and absolute element index. + pub fn new(ctx: &mut Context, tensor: Value, absolute_index: Value, value: Value) -> Self { + Self { + op: Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![], + vec![tensor, absolute_index, value], + vec![], + 0, + ), + } + } + + /// View an existing operation through this wrapper without allocating or verifying it. + pub fn wrap(op: Ptr) -> Self { + Self { op } + } +} + +impl Verify for CuteTensorStoreElementAbsOp { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); + if op.get_num_operands() != 3 || op.get_num_results() != 0 { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs needs 3 operands and 0 results" + ); + } + let Some(view) = view_of(ctx, op.get_operand(0)) else { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs destination must be a tensor view" + ); + }; + if view.access != CuteTensorAccessAttr::ReadWrite { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs destination must be writable" + ); + } + if !matches!(view.layout, CuteTensorLayoutAttr::Tile1D(size) if size > 0) { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs destination must be a selected Tile1D" + ); + } + if !is_u64(ctx, op.get_operand(1)) { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs index must be an unsigned 64-bit integer" + ); + } + if op.get_operand(2).get_type(ctx) != view.logical { + return verify_err!( + op.loc(), + "cute.tensor_store_element_abs value must match the tensor logical type" + ); + } + Ok(()) + } +} + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/tests.rs new file mode 100644 index 0000000000..a44f3de91f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tensor/tests.rs @@ -0,0 +1,187 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression fixtures for this operation family. + +use super::*; +use dialect_mir::ops::MirUndefOp; +use pliron::builtin::types::FP32Type; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Create a context and typed global/register pointers for the elementwise tensor-flow tests. +fn setup() -> (Context, TypeHandle, Value, Value, Value, Value) { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let global_ro_ty: TypeHandle = + MirPtrType::get(&mut ctx, f32_ty, false, address_space::GLOBAL).into(); + let global_rw_ty: TypeHandle = + MirPtrType::get(&mut ctx, f32_ty, true, address_space::GLOBAL).into(); + let local_rw_ty: TypeHandle = + MirPtrType::get(&mut ctx, f32_ty, true, address_space::LOCAL).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let global_ro = undef(&mut ctx, global_ro_ty); + let global_rw = undef(&mut ctx, global_rw_ty); + let local = undef(&mut ctx, local_rw_ty); + let u64_value = undef(&mut ctx, u64_ty); + (ctx, f32_ty, global_ro, global_rw, local, u64_value) +} + +/// The supported make/divide/select/load/store/tail operations compose while preserving tensor metadata. +#[test] +fn all_eight_ops_verify_as_one_visible_tensor_flow() { + let (mut ctx, f32_ty, global_ro, global_rw, local, index) = setup(); + + let read = CuteTensorMakeOp::new( + &mut ctx, + global_ro, + index, + f32_ty, + f32_ty, + CuteTensorAccessAttr::ReadOnly, + 4, + ); + assert!(read.verify(&ctx).is_ok()); + let read_contiguous = read.get_operation().deref(&ctx).get_result(0); + let read_zipped = CuteTensorZippedDivideOp::new(&mut ctx, read_contiguous, 4); + assert!(read_zipped.verify(&ctx).is_ok()); + let read_zipped_value = read_zipped.get_operation().deref(&ctx).get_result(0); + let read_tile = CuteTensorSliceOp::new(&mut ctx, read_zipped_value, index); + assert!(read_tile.verify(&ctx).is_ok()); + let read_tile_value = read_tile.get_operation().deref(&ctx).get_result(0); + + assert!( + CuteTensorIsFullOp::new(&mut ctx, read_tile_value) + .verify(&ctx) + .is_ok() + ); + assert!( + CuteTensorBaseOp::new(&mut ctx, read_tile_value) + .verify(&ctx) + .is_ok() + ); + assert!( + CuteTensorLoadIntoOp::new(&mut ctx, read_tile_value, local, 16) + .verify(&ctx) + .is_ok() + ); + + let write = CuteTensorMakeOp::new( + &mut ctx, + global_rw, + index, + f32_ty, + f32_ty, + CuteTensorAccessAttr::ReadWrite, + 4, + ); + let write_contiguous = write.get_operation().deref(&ctx).get_result(0); + let write_zipped = CuteTensorZippedDivideOp::new(&mut ctx, write_contiguous, 4); + let write_zipped_value = write_zipped.get_operation().deref(&ctx).get_result(0); + let write_tile = CuteTensorSliceOp::new(&mut ctx, write_zipped_value, index); + let write_tile_value = write_tile.get_operation().deref(&ctx).get_result(0); + + assert!( + CuteTensorStoreFromOp::new(&mut ctx, local, write_tile_value, 16) + .verify(&ctx) + .is_ok() + ); + let scalar = undef(&mut ctx, f32_ty); + assert!( + CuteTensorStoreElementAbsOp::new(&mut ctx, write_tile_value, index, scalar) + .verify(&ctx) + .is_ok() + ); +} + +/// A transfer cannot strengthen view access or accept an invalid alignment promise. +#[test] +fn verifiers_reject_access_and_alignment_drift() { + let (mut ctx, f32_ty, global_ro, _global_rw, local, index) = setup(); + let read = CuteTensorMakeOp::new( + &mut ctx, + global_ro, + index, + f32_ty, + f32_ty, + CuteTensorAccessAttr::ReadOnly, + 4, + ); + let contiguous = read.get_operation().deref(&ctx).get_result(0); + let zipped = CuteTensorZippedDivideOp::new(&mut ctx, contiguous, 4); + let zipped_value = zipped.get_operation().deref(&ctx).get_result(0); + let tile = CuteTensorSliceOp::new(&mut ctx, zipped_value, index); + let tile_value = tile.get_operation().deref(&ctx).get_result(0); + + let weak_alignment = CuteTensorLoadIntoOp::new(&mut ctx, tile_value, local, 8); + assert!( + weak_alignment + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("sufficient power-of-two alignment") + ); + + let wrong_access = CuteTensorStoreFromOp::new(&mut ctx, local, tile_value, 16); + assert!( + wrong_access + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("must be writable") + ); +} + +/// Dividing a view may change its selection layout but must preserve its storage facts. +#[test] +fn zipped_divide_rejects_result_fact_drift() { + let (mut ctx, f32_ty, global_ro, _global_rw, _local, index) = setup(); + let read = CuteTensorMakeOp::new( + &mut ctx, + global_ro, + index, + f32_ty, + f32_ty, + CuteTensorAccessAttr::ReadOnly, + 4, + ); + let input = read.get_operation().deref(&ctx).get_result(0); + let wrong_result: TypeHandle = CuteTensorViewType::get( + &mut ctx, + f32_ty, + f32_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadWrite, + 4, + CuteTensorLayoutAttr::Zipped1D(4), + ) + .into(); + let raw = Operation::new( + &mut ctx, + CuteTensorZippedDivideOp::get_concrete_op_info(), + vec![wrong_result], + vec![input], + vec![], + 0, + ); + let divide = CuteTensorZippedDivideOp::wrap(raw); + divide.set_attr_zipped_layout(&ctx, CuteTensorLayoutAttr::Zipped1D(4)); + + assert!( + divide + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("may change only") + ); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tests.rs new file mode 100644 index 0000000000..341e9ea0e7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/ops/tests.rs @@ -0,0 +1,789 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe operation, type, and semantic verification contracts. + +use cute_layout::{ + ComposedLayout, IntTuple, Layout, MAX_STATIC_VALIDATION_BYTES, OffsetUnit, Swizzle, +}; +use dialect_mir::{ + ops::MirUndefOp, + types::{MirArrayType, MirFP16Type, MirPtrType, address_space}, +}; +use pliron::{ + builtin::types::{FP32Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + op::Op, + operation::Operation, + r#type::TypeHandle, + value::Value, +}; + +use super::*; +use crate::attributes::{CuteDivisibilityAttr, CuteMatrixRoleAttr}; + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, ty: TypeHandle) -> Value { + MirUndefOp::new(ctx, ty) + .get_operation() + .deref(ctx) + .get_result(0) +} + +/// Build a matrix-load fixture with a selected operand role and otherwise valid shared carriers. +fn matrix_load( + ctx: &mut Context, + role: CuteMatrixRoleAttr, + source_space: u32, + malformed: &str, +) -> CuteLdmatrixOp { + let elem: TypeHandle = if malformed == "element" { + FP32Type::get(ctx).into() + } else { + MirFP16Type::get(ctx).into() + }; + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let source: TypeHandle = MirPtrType::get(ctx, elem, false, source_space).into(); + let words = if role == CuteMatrixRoleAttr::B { 2 } else { 4 }; + let registers: TypeHandle = MirArrayType::get(ctx, u32_ty, words).into(); + let slot: TypeHandle = MirPtrType::get_generic(ctx, registers, malformed != "slot").into(); + let operands = [ + undef(ctx, source), + undef(ctx, if malformed == "row" { u32_ty } else { u64_ty }), + undef( + ctx, + if malformed == "column" { + u32_ty + } else { + u64_ty + }, + ), + undef(ctx, if malformed == "lane" { u64_ty } else { u32_ty }), + undef(ctx, slot), + ]; + let layout = ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + "(64,64):(64,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(); + CuteLdmatrixOp::new(ctx, operands, role, layout, elem) +} + +/// All supported A/B load roles accept their matching shared FP16 and register-carrier ABI. +#[test] +fn matrix_load_roles_accept_the_shared_f16_fragment_abi() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + for role in [ + CuteMatrixRoleAttr::A, + CuteMatrixRoleAttr::B, + CuteMatrixRoleAttr::BPair, + CuteMatrixRoleAttr::BPairKMajor, + ] { + for space in [address_space::SHARED, address_space::GENERIC] { + matrix_load(&mut ctx, role, space, "").verify(&ctx).unwrap(); + } + } +} + +/// The non-transposed K-major B-pair variant still checks pointer roles and scalar carrier widths. +#[test] +fn k_major_b_pair_keeps_matrix_load_operand_checks() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + for (malformed, message) in [ + ("element", "supports f16 elements only"), + ("row", "warp tile row must be an unsigned 64-bit integer"), + ( + "column", + "warp tile column must be an unsigned 64-bit integer", + ), + ("lane", "lane id must be an unsigned 32-bit integer"), + ("slot", "fragment slot must be mutable"), + ] { + let load = matrix_load( + &mut ctx, + CuteMatrixRoleAttr::BPairKMajor, + address_space::SHARED, + malformed, + ); + let error = load.verify(&ctx).unwrap_err().to_string(); + assert!(error.contains(message), "{malformed}: {error}"); + } + let load = matrix_load( + &mut ctx, + CuteMatrixRoleAttr::BPairKMajor, + address_space::GLOBAL, + "", + ); + let error = load.verify(&ctx).unwrap_err().to_string(); + assert!( + error.contains("source must point at shared/generic f16"), + "{error}" + ); +} + +/// Return mutually compatible thread, value, tile, and shared layouts for one cooperative-copy partition. +fn valid_layouts() -> (Layout, Layout, Layout, ComposedLayout) { + // Six rows, four columns. Each row/thread owns one 16-byte f32 atom. + // + // thread 2 -> (2,0) (2,1) (2,2) (2,3) + let thread: Layout = "(6,1):(1,0)".parse().unwrap(); + let value: Layout = "(1,4):(0,1)".parse().unwrap(); + let tile: Layout = "(6,4):(4,1)".parse().unwrap(); + // Match the decoder canary: an eight-element prefix plus a composed + // swizzle. The tile is small enough that the source field is zero, + // but the complete composition and byte recast are still exercised. + let smem = + ComposedLayout::new(Swizzle::new(3, 4, 3), 8, tile.clone(), OffsetUnit::Elements).unwrap(); + (thread, value, tile, smem) +} + +/// Build a cooperative-copy fixture with controllable pointers, atom width, and static layouts. +#[allow(clippy::too_many_arguments)] +fn cooperative_copy( + ctx: &mut Context, + src_mutable: bool, + src_address_space: u32, + dst_mutable: bool, + dst_address_space: u32, + atom_bytes: u32, + thread_layout: Layout, + value_layout: Layout, + tile_layout: Layout, + smem_layout: ComposedLayout, +) -> CuteCopyG2SOp { + cooperative_copy_with_integer_abi( + ctx, + src_mutable, + src_address_space, + dst_mutable, + dst_address_space, + 64, + Signedness::Unsigned, + 32, + Signedness::Unsigned, + atom_bytes, + thread_layout, + value_layout, + tile_layout, + smem_layout, + ) +} + +/// Build the copy fixture with independently selected index/lane integer types for negative ABI cases. +#[allow(clippy::too_many_arguments)] +fn cooperative_copy_with_integer_abi( + ctx: &mut Context, + src_mutable: bool, + src_address_space: u32, + dst_mutable: bool, + dst_address_space: u32, + usize_width: u32, + usize_signedness: Signedness, + tidx_width: u32, + tidx_signedness: Signedness, + atom_bytes: u32, + thread_layout: Layout, + value_layout: Layout, + tile_layout: Layout, + smem_layout: ComposedLayout, +) -> CuteCopyG2SOp { + let elem: TypeHandle = FP32Type::get(ctx).into(); + let src_ty: TypeHandle = MirPtrType::get(ctx, elem, src_mutable, src_address_space).into(); + let dst_ty: TypeHandle = MirPtrType::get(ctx, elem, dst_mutable, dst_address_space).into(); + let usize_ty: TypeHandle = IntegerType::get(ctx, usize_width, usize_signedness).into(); + let tidx_ty: TypeHandle = IntegerType::get(ctx, tidx_width, tidx_signedness).into(); + let operands = [ + undef(ctx, src_ty), + undef(ctx, usize_ty), + undef(ctx, usize_ty), + undef(ctx, usize_ty), + undef(ctx, usize_ty), + undef(ctx, usize_ty), + undef(ctx, dst_ty), + undef(ctx, usize_ty), + undef(ctx, tidx_ty), + ]; + CuteCopyG2SOp::new( + ctx, + operands, + atom_bytes, + thread_layout, + value_layout, + tile_layout, + smem_layout, + 4, + elem, + ) +} + +/// Capture a failed local verification diagnostic for a malformed fixture operation. +fn verifier_error(copy: &CuteCopyG2SOp, ctx: &Context, expected: &str) { + let error = copy.verify(ctx).unwrap_err().to_string(); + assert!( + error.contains(expected), + "expected `{expected}` in `{error}`" + ); +} + +/// Build a raw TMA load fixture with the supplied destination and completion barrier. +fn tma_copy(ctx: &mut Context, dst_address_space: u32) -> CuteTmaLoad2dOp { + let elem: TypeHandle = FP32Type::get(ctx).into(); + let u8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let destination: TypeHandle = MirPtrType::get(ctx, elem, true, dst_address_space).into(); + let barrier: TypeHandle = MirPtrType::get(ctx, u64_ty, true, address_space::SHARED).into(); + let tensor_map: TypeHandle = MirPtrType::get_generic(ctx, u8_ty, false).into(); + let operands = [ + undef(ctx, destination), + undef(ctx, barrier), + undef(ctx, tensor_map), + undef(ctx, u64_ty), + undef(ctx, u64_ty), + ]; + let smem = ComposedLayout::from_layout("(4,4):(4,1)".parse().unwrap(), OffsetUnit::Elements); + CuteTmaLoad2dOp::new(ctx, operands, smem, elem) +} + +/// Build a raw TMA store fixture with the supplied shared source. +fn tma_store(ctx: &mut Context, src_address_space: u32) -> CuteTmaStore2dOp { + let elem: TypeHandle = FP32Type::get(ctx).into(); + let u8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let source: TypeHandle = MirPtrType::get(ctx, elem, true, src_address_space).into(); + let tensor_map: TypeHandle = MirPtrType::get_generic(ctx, u8_ty, false).into(); + let operands = [ + undef(ctx, source), + undef(ctx, tensor_map), + undef(ctx, u64_ty), + undef(ctx, u64_ty), + ]; + let smem = ComposedLayout::from_layout("(4,4):(4,1)".parse().unwrap(), OffsetUnit::Elements); + CuteTmaStore2dOp::new(ctx, operands, smem, elem) +} + +/// Malformed operand/result counts report errors instead of indexing missing values and panicking. +#[test] +fn custom_verifiers_check_cardinality_before_indexing() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let malformed = |ctx: &mut Context, info| Operation::new(ctx, info, vec![], vec![], vec![], 0); + + let g2s = CuteCopyG2SOp::wrap(malformed(&mut ctx, CuteCopyG2SOp::get_concrete_op_info())); + verifier_error(&g2s, &ctx, "exactly 9 operands and 0 results"); + + let copy = CuteCopyOp::wrap(malformed(&mut ctx, CuteCopyOp::get_concrete_op_info())); + let error = copy.verify(&ctx).unwrap_err().to_string(); + assert!( + error.contains("exactly 2 operands and 0 results"), + "{error}" + ); + + let assume = + CuteAssumeDivOp::wrap(malformed(&mut ctx, CuteAssumeDivOp::get_concrete_op_info())); + let error = assume.verify(&ctx).unwrap_err().to_string(); + assert!(error.contains("exactly 1 operand and 1 result"), "{error}"); + + let tma_load = + CuteTmaLoad2dOp::wrap(malformed(&mut ctx, CuteTmaLoad2dOp::get_concrete_op_info())); + let error = tma_load.verify(&ctx).unwrap_err().to_string(); + assert!( + error.contains("exactly 5 operands and 0 results"), + "{error}" + ); + + let tma_store = CuteTmaStore2dOp::wrap(malformed( + &mut ctx, + CuteTmaStore2dOp::get_concrete_op_info(), + )); + let error = tma_store.verify(&ctx).unwrap_err().to_string(); + assert!( + error.contains("exactly 4 operands and 0 results"), + "{error}" + ); +} + +/// A complete, non-overlapping thread/value partition satisfies the cooperative-copy contract. +#[test] +fn cooperative_abi_accepts_one_exact_well_formed_partition() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let (thread_layout, value_layout, tile_layout, smem_layout) = valid_layouts(); + + let copy = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + thread_layout, + value_layout, + tile_layout, + smem_layout, + ); + copy.verify(&ctx).unwrap(); +} + +/// Generic pointers cannot stand in for a proven CTA-local shared TMA destination. +#[test] +fn tma_abi_requires_an_explicit_cta_shared_destination() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + tma_copy(&mut ctx, address_space::SHARED) + .verify(&ctx) + .unwrap(); + let generic = tma_copy(&mut ctx, address_space::GENERIC); + let error = generic.verify(&ctx).unwrap_err().to_string(); + assert!(error.contains("mutable CTA-shared elem memory"), "{error}"); +} + +/// TMA stores require a proven CTA-local shared source rather than an ambiguous generic pointer. +#[test] +fn tma_store_abi_requires_an_explicit_cta_shared_source() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + tma_store(&mut ctx, address_space::SHARED) + .verify(&ctx) + .unwrap(); + let generic = tma_store(&mut ctx, address_space::GENERIC); + let error = generic.verify(&ctx).unwrap_err().to_string(); + assert!(error.contains("CTA-shared elem memory"), "{error}"); +} + +/// Row-pitch divisibility must preserve alignment for every vector copy atom. +#[test] +fn cooperative_abi_rejects_a_pitch_promise_weaker_than_the_atom() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let (thread, value, tile, smem) = valid_layouts(); + let copy = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + thread, + value, + tile, + smem, + ); + + // f32 x a two-element pitch promise protects only 8 bytes: + // + // row start ── 8-byte guarantee ──X── 16-byte cp.async + copy.set_attr_leading_dim_divisor(&ctx, CuteDivisibilityAttr(2)); + verifier_error(©, &ctx, "promise is weaker than its copy atom"); +} + +/// A zero-width copy atom fails validation before any alignment division occurs. +#[test] +fn cooperative_abi_rejects_zero_atom_before_alignment_arithmetic() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let (thread, value, tile, smem) = valid_layouts(); + let copy = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 0, + thread, + value, + tile, + smem, + ); + verifier_error(©, &ctx, "atom must be 4, 8, or 16 bytes"); +} + +/// Global/shared pointer roles and destination writability are verified independently of layout shape. +#[test] +fn cooperative_abi_rejects_pointer_role_and_mutability_errors() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let (thr, val, tile, smem) = valid_layouts(); + let shared_source = cooperative_copy( + &mut ctx, + false, + address_space::SHARED, + true, + address_space::SHARED, + 16, + thr, + val, + tile, + smem, + ); + verifier_error( + &shared_source, + &ctx, + "global source has the wrong address space", + ); + + let (thr, val, tile, smem) = valid_layouts(); + let global_destination = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::GLOBAL, + 16, + thr, + val, + tile, + smem, + ); + verifier_error( + &global_destination, + &ctx, + "shared destination must be mutable shared/generic memory", + ); + + let (thr, val, tile, smem) = valid_layouts(); + let immutable_destination = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + false, + address_space::SHARED, + 16, + thr, + val, + tile, + smem, + ); + verifier_error( + &immutable_destination, + &ctx, + "shared destination must be mutable shared/generic memory", + ); +} + +/// Lane and index carriers must use the exact unsigned widths expected by lowering. +#[test] +fn cooperative_abi_rejects_integer_width_or_signedness_drift() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let make = |ctx: &mut Context, usize_width, usize_signedness, tidx_width, tidx_signedness| { + let (thread, value, tile, smem) = valid_layouts(); + cooperative_copy_with_integer_abi( + ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + usize_width, + usize_signedness, + tidx_width, + tidx_signedness, + 16, + thread, + value, + tile, + smem, + ) + }; + + let narrow_usize = make(&mut ctx, 32, Signedness::Unsigned, 32, Signedness::Unsigned); + verifier_error( + &narrow_usize, + &ctx, + "row count must be an unsigned 64-bit integer", + ); + + let signed_usize = make(&mut ctx, 64, Signedness::Signed, 32, Signedness::Unsigned); + verifier_error( + &signed_usize, + &ctx, + "row count must be an unsigned 64-bit integer", + ); + + let signed_tidx = make(&mut ctx, 64, Signedness::Unsigned, 32, Signedness::Signed); + verifier_error( + &signed_tidx, + &ctx, + "thread index must be an unsigned 32-bit integer", + ); +} + +/// Thread/value coverage cannot contain overlap, gaps, or an incompatible atom grouping. +#[test] +fn cooperative_abi_rejects_bad_tv_and_atom_partitions() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let sparse_threads = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(2,2):(0,1)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(4,4):(4,1)".parse().unwrap(), + ComposedLayout::from_layout("(4,4):(4,1)".parse().unwrap(), OffsetUnit::Elements), + ); + verifier_error(&sparse_threads, &ctx, "invalid thread layout"); + + let partial_atom = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(6,1):(1,0)".parse().unwrap(), + "(1,2):(0,1)".parse().unwrap(), + "(6,2):(2,1)".parse().unwrap(), + ComposedLayout::from_layout("(6,2):(2,1)".parse().unwrap(), OffsetUnit::Elements), + ); + verifier_error(&partial_atom, &ctx, "invalid copy atom"); + + let incomplete_tile = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(6,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(6,3):(3,1)".parse().unwrap(), + ComposedLayout::from_layout("(6,3):(3,1)".parse().unwrap(), OffsetUnit::Elements), + ); + verifier_error(&incomplete_tile, &ctx, "assignments"); + + let aliased_tile = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(6,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(6,4):(0,1)".parse().unwrap(), + ComposedLayout::from_layout("(6,4):(4,1)".parse().unwrap(), OffsetUnit::Elements), + ); + verifier_error(&aliased_tile, &ctx, "invalid tile layout"); + + let (thread, value, _, _) = valid_layouts(); + let flat_tile = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + thread, + value, + Layout::contiguous(24), + ComposedLayout::from_layout(Layout::contiguous(24), OffsetUnit::Elements), + ); + verifier_error(&flat_tile, &ctx, "exactly two row/column tile modes"); +} + +/// Each instruction must address one contiguous aligned atom after shared swizzling. +#[test] +fn cooperative_abi_rejects_scattered_or_swizzle_split_atoms() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let scattered_source = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(6,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(6,4):(1,6)".parse().unwrap(), + ComposedLayout::from_layout("(6,4):(4,1)".parse().unwrap(), OffsetUnit::Elements), + ); + verifier_error( + &scattered_source, + &ctx, + "tile layout is not canonical row-major", + ); + + let split_destination = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(6,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(6,4):(4,1)".parse().unwrap(), + ComposedLayout::new( + Swizzle::new(1, 3, 3), + 64, + "(6,4):(16,4)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ); + verifier_error( + &split_destination, + &ctx, + "shared-memory copy atom is not physically contiguous", + ); + + let aliased_destination = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(2,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(2,4):(4,1)".parse().unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + 0, + "(2,4):(0,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ); + verifier_error( + &aliased_destination, + &ctx, + "invalid physical shared-memory map", + ); + + let negative_destination = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(2,1):(1,0)".parse().unwrap(), + "(1,4):(0,1)".parse().unwrap(), + "(2,4):(4,1)".parse().unwrap(), + ComposedLayout::new( + Swizzle::IDENTITY, + -4, + "(2,4):(4,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ); + verifier_error(&negative_destination, &ctx, "negative smem offset"); +} + +/// A vector atom cannot straddle logical rows even when its flattened offsets look contiguous. +#[test] +fn source_atoms_cannot_cross_a_logical_tile_row() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + // The four f32 values are 16 contiguous bytes in a dense tile, but + // they span two logical matrix rows: + // + // row 0: [ A B ] + // row 1: [ C D ] <- one atom must not bridge this boundary + let tile: Layout = "(2,2):(2,1)".parse().unwrap(); + let crossing = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + "(1,1):(1,1)".parse().unwrap(), + "(2,2):(2,1)".parse().unwrap(), + tile.clone(), + ComposedLayout::from_layout(tile, OffsetUnit::Elements), + ); + verifier_error(&crossing, &ctx, "global copy atom crosses a tile row"); +} + +/// Excessive enumeration sizes and overflowing address calculations fail conservatively. +#[test] +fn cooperative_abi_bounds_exhaustive_work_and_checked_offsets() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + + let elements = MAX_STATIC_VALIDATION_BYTES / 4 + 1; + let shape = IntTuple::Tuple(vec![IntTuple::Leaf(elements), IntTuple::Leaf(1)]); + let thread = Layout::new( + shape.clone(), + IntTuple::Tuple(vec![IntTuple::Leaf(1), IntTuple::Leaf(0)]), + ); + let value: Layout = "(1,1):(0,0)".parse().unwrap(); + let tile = Layout::new( + shape, + IntTuple::Tuple(vec![IntTuple::Leaf(1), IntTuple::Leaf(0)]), + ); + let oversized = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 4, + thread, + value, + tile.clone(), + ComposedLayout::from_layout(tile, OffsetUnit::Elements), + ); + verifier_error(&oversized, &ctx, "static byte-map limit"); + + let (thread, value, tile, _) = valid_layouts(); + let overflowing_smem = cooperative_copy( + &mut ctx, + false, + address_space::GLOBAL, + true, + address_space::SHARED, + 16, + thread, + value, + tile, + ComposedLayout::new( + Swizzle::IDENTITY, + i64::MAX - 15, + "(6,4):(16,4)".parse().unwrap(), + OffsetUnit::Bytes, + ) + .unwrap(), + ); + verifier_error(&overflowing_smem, &ctx, "composed offset"); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/epilogue.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/epilogue.rs new file mode 100644 index 0000000000..6233e245ea --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/epilogue.rs @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared output placement and asynchronous store synchronization contracts. + +use dialect_mir::types::MirFP16Type; +use pliron::{ + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypeHandle, TypedHandle}, + verify_err, +}; + +use crate::attributes::CuteEpiloguePlanAttr; + +/// Static meaning of the shared result allocation used by one tiled MMA. +/// +/// This type is carried only by operation attributes. Runtime values remain +/// the same ordinary pointer, warp/lane scalars, and FP32 accumulator used by +/// the Rust kernel: +/// +/// ```text +/// AS3 f16 pointer + 64xf32 accumulator +/// │ TypeAttr> +/// ▼ +/// two 128x64 shared halves +/// ``` +#[pliron_type(name = "cute.epilogue_tile", format = "`<` $storage `,` $plan `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteEpilogueTileType { + /// MIR element type physically stored in memory. + pub storage: TypeHandle, + /// Static layout and instruction contract used by this fixture or type. + pub plan: CuteEpiloguePlanAttr, +} + +impl CuteEpilogueTileType { + /// Attach the current epilogue placement to its physical storage type. + pub fn get( + ctx: &mut Context, + storage: TypeHandle, + plan: CuteEpiloguePlanAttr, + ) -> TypedHandle { + Type::instantiate(Self { storage, plan }, ctx) + } + + /// Return the element capacity of one output half, or `None` on overflow. + #[must_use] + pub fn half_elements(&self) -> Option { + self.plan.half_elements() + } + + /// Return the element capacity of the complete output tile, or `None` on overflow. + #[must_use] + pub fn full_elements(&self) -> Option { + self.plan.full_elements() + } + + /// Return one output half's byte capacity, or `None` for an invalid size. + #[must_use] + pub fn half_bytes(&self) -> Option { + self.plan.half_bytes(2) + } + + /// Return the complete output tile's byte capacity, or `None` for an invalid size. + #[must_use] + pub fn full_bytes(&self) -> Option { + self.plan.full_bytes(2) + } +} + +impl Verify for CuteEpilogueTileType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + if self + .storage + .deref(ctx) + .downcast_ref::() + .is_none() + { + return verify_err!( + Location::Unknown, + "cute.epilogue_tile v0 storage must be f16" + ); + } + self.plan.verify(ctx)?; + if self.half_elements() != Some(8192) + || self.full_elements() != Some(16384) + || self.half_bytes() != Some(16 * 1024) + || self.full_bytes() != Some(32 * 1024) + { + return verify_err!( + Location::Unknown, + "cute.epilogue_tile v0 must be two 16 KiB halves in one 32 KiB allocation" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/gemv.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/gemv.rs new file mode 100644 index 0000000000..e7adbf7355 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/gemv.rs @@ -0,0 +1,205 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Packed-value and scale-view contracts for the block-scaled GEMV flow. + +use pliron::{ + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypeHandle, TypedHandle}, + verify_err, +}; + +use super::tensor::CuteTensorViewType; +use crate::attributes::{ + CuteScaledLayoutAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, +}; + +/// A packed value tensor bound to its scale tensor. +/// +/// This remains a view. It owns no storage and loads nothing: +/// +/// ```text +/// E2M1 tensor ──┐ +/// ├── !cute.scaled_view<..., Full> +/// UE8M0 tensor ─┘ +/// ``` +#[pliron_type( + name = "cute.scaled_view", + format = "`<` $values `,` $scales `,` $role `,` $layout `>`" +)] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteScaledViewType { + /// Type of the packed value tensor bound to the scale tensor. + pub values: TypeHandle, + /// Type of the scale tensor paired with the packed values. + pub scales: TypeHandle, + /// Logical M/N coordinate role assigned to the tensor. + pub role: CuteTensorRoleAttr, + /// Static coordinate layout used to interpret the storage. + pub layout: CuteScaledLayoutAttr, +} + +impl CuteScaledViewType { + /// Bind value and scale tensor types at one visible selection stage. + pub fn get( + ctx: &mut Context, + values: TypeHandle, + scales: TypeHandle, + role: CuteTensorRoleAttr, + layout: CuteScaledLayoutAttr, + ) -> TypedHandle { + Type::instantiate( + Self { + values, + scales, + role, + layout, + }, + ctx, + ) + } + + /// Copy the bound tensors and role, changing only the selection stage. + #[must_use] + pub fn with_layout( + &self, + ctx: &mut Context, + layout: CuteScaledLayoutAttr, + ) -> TypedHandle { + Self::get(ctx, self.values, self.scales, self.role, layout) + } + + /// Return how many logical K values share one scale. + #[must_use] + pub fn values_per_scale(&self, ctx: &Context) -> Option { + let scales = self.scales.deref(ctx); + let view = scales.downcast_ref::()?; + view.layout.values_per_scale() + } + + /// Return the selected logical K width, if this is a K tile. + #[must_use] + pub const fn k_width(&self) -> Option { + match self.layout { + CuteScaledLayoutAttr::KTile(width) => Some(width), + _ => None, + } + } +} + +impl Verify for CuteScaledViewType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let values_ty = self.values.deref(ctx); + let Some(values) = values_ty.downcast_ref::() else { + return verify_err!( + Location::Unknown, + "cute.scaled_view values must be a tensor view" + ); + }; + let scales_ty = self.scales.deref(ctx); + let Some(scales) = scales_ty.downcast_ref::() else { + return verify_err!( + Location::Unknown, + "cute.scaled_view scales must be a tensor view" + ); + }; + values.verify(ctx)?; + scales.verify(ctx)?; + + if values.format != CuteTensorFormatAttr::E2M1 + || values.layout != CuteTensorLayoutAttr::KMajor + { + return verify_err!( + Location::Unknown, + "cute.scaled_view values must be an E2M1 KMajor tensor" + ); + } + if !matches!( + scales.format, + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3 + ) || scales.layout != CuteTensorLayoutAttr::BlockScaleKMajor(16) + { + return verify_err!( + Location::Unknown, + "cute.scaled_view scales must be a UE8M0 or E4M3 BlockScaleKMajor<16> tensor" + ); + } + if self.role == CuteTensorRoleAttr::Generic + || values.role != self.role + || scales.role != self.role + { + return verify_err!( + Location::Unknown, + "cute.scaled_view values, scales, and result must share one Mkl or Nkl role" + ); + } + if values.logical != scales.logical { + return verify_err!( + Location::Unknown, + "cute.scaled_view values and scales must decode to the same logical type" + ); + } + self.layout.verify(ctx)?; + if let CuteScaledLayoutAttr::KTile(width) = self.layout + && width % 16 != 0 + { + return verify_err!( + Location::Unknown, + "cute.scaled_view K tile width must contain whole 16-value scale groups" + ); + } + Ok(()) + } +} + +/// One loaded block-scaled tile held in registers. +/// +/// The source type keeps the packed format, scale format, role, and K width +/// available without copying those facts into a second list. +#[pliron_type(name = "cute.fragment", format = "`<` $source `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteFragmentType { + /// Type of the selected block-scaled view from which the fragment was loaded. + pub source: TypeHandle, +} + +impl CuteFragmentType { + /// Describe the register fragment loaded from one scaled K tile. + pub fn get(ctx: &mut Context, source: TypeHandle) -> TypedHandle { + Type::instantiate(Self { source }, ctx) + } + + /// Read the source view carrying this fragment's semantic facts. + #[must_use] + pub fn source_view(&self, ctx: &Context) -> Option { + self.source + .deref(ctx) + .downcast_ref::() + .cloned() + } +} + +impl Verify for CuteFragmentType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let Some(source) = self.source_view(ctx) else { + return verify_err!( + Location::Unknown, + "cute.fragment source must be a scaled view" + ); + }; + source.verify(ctx)?; + if !matches!(source.layout, CuteScaledLayoutAttr::KTile(width) if width > 0) { + return verify_err!( + Location::Unknown, + "cute.fragment source must be a selected K tile" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mma.rs new file mode 100644 index 0000000000..104ca1a7d6 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mma.rs @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared tensor placement, register carriers, and tiled matrix-instruction contracts. + +use dialect_mir::types::{MirArrayType, MirStructType, MirTupleType}; +use pliron::{ + builtin::types::{FP32Type, IntegerType}, + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypeHandle, TypedHandle}, + verify_err, +}; + +use super::tensor::CuteTensorViewType; +use crate::attributes::{ + CuteComposedLayoutAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, + CuteTensorLayoutAttr, +}; + +/// Physical scalar registers inside one ordinary MIR aggregate. +/// +/// MMA operations keep runtime values in ordinary Rust arrays and structs. +/// This count lets the verifier check their register requirements. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MmaCarrierShape { + /// Count of packed 32-bit registers in the ordinary carrier. + pub u32_registers: u64, + /// Count of FP32 registers in the ordinary carrier. + pub f32_registers: u64, +} + +impl MmaCarrierShape { + /// Describe a carrier containing only the given number of packed 32-bit registers. + #[must_use] + pub const fn u32(registers: u64) -> Self { + Self { + u32_registers: registers, + f32_registers: 0, + } + } + + /// Describe a carrier containing only the given number of FP32 registers. + #[must_use] + pub const fn f32(registers: u64) -> Self { + Self { + u32_registers: 0, + f32_registers: registers, + } + } + + /// Bytes carried by all counted registers. + #[must_use] + pub const fn bytes(self) -> Option { + let Some(registers) = self.u32_registers.checked_add(self.f32_registers) else { + return None; + }; + registers.checked_mul(4) + } + + /// Combine register counts for aggregate fields, returning `None` on overflow. + fn checked_add(self, other: Self) -> Option { + Some(Self { + u32_registers: self.u32_registers.checked_add(other.u32_registers)?, + f32_registers: self.f32_registers.checked_add(other.f32_registers)?, + }) + } + + /// Scale register counts for an array of carriers, returning `None` on overflow. + fn checked_repeat(self, count: u64) -> Option { + Some(Self { + u32_registers: self.u32_registers.checked_mul(count)?, + f32_registers: self.f32_registers.checked_mul(count)?, + }) + } +} + +/// Count the u32 and f32 registers in a Rust-shaped MIR carrier. +/// +/// Arrays, tuples, and structs are walked recursively. Empty marker structs +/// such as `PhantomData` contribute nothing. Any other scalar or +/// aggregate kind returns `None`. +#[must_use] +pub fn mma_carrier_shape(ctx: &Context, ty: TypeHandle) -> Option { + let ty = ty.deref(ctx); + if let Some(integer) = ty.downcast_ref::() { + return (integer.width() == 32 && integer.is_unsigned()).then_some(MmaCarrierShape::u32(1)); + } + if ty.downcast_ref::().is_some() { + return Some(MmaCarrierShape::f32(1)); + } + if let Some(array) = ty.downcast_ref::() { + return mma_carrier_shape(ctx, array.element_ty)?.checked_repeat(array.size); + } + let children = if let Some(tuple) = ty.downcast_ref::() { + Some(tuple.types.as_slice()) + } else { + ty.downcast_ref::() + .map(|structure| structure.field_types.as_slice()) + }?; + children + .iter() + .try_fold(MmaCarrierShape::default(), |sum, child| { + sum.checked_add(mma_carrier_shape(ctx, *child)?) + }) +} + +/// Static meaning attached to a shared-memory pointer/capacity pair. +/// +/// This type lives in a `TypeAttr` on `cute.smem_tensor_overlay`; it is not an +/// SSA result and therefore never has to cross a loop edge: +/// +/// ```text +/// ordinary AS3 pointer + ordinary u64 capacity +/// │ {view = !cute.smem_tensor<...>} +/// ▼ +/// ordinary AS3 pointer + ordinary u64 capacity +/// ``` +/// +/// `placement` explains where logical rows and K positions live inside one +/// shared stage. The wrapped tensor explains whether those carriers are E2M1 +/// values or UE8M0 scales and whether their rows follow M or N. +#[pliron_type(name = "cute.smem_tensor", format = "`<` $tensor `,` $placement `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteSmemTensorType { + /// Wrapped tensor-view type containing storage, format, role, and access facts. + pub tensor: TypeHandle, + /// Mapping from logical coordinates to physical shared storage offsets. + pub placement: CuteComposedLayoutAttr, +} + +impl CuteSmemTensorType { + /// Attach one shared-memory placement to packed tensor facts. + pub fn get( + ctx: &mut Context, + tensor: TypeHandle, + placement: crate::layout::ComposedLayout, + ) -> TypedHandle { + Type::instantiate( + Self { + tensor, + placement: CuteComposedLayoutAttr(placement), + }, + ctx, + ) + } + + /// Return the packed tensor facts wrapped by this placement. + #[must_use] + pub fn tensor_view(&self, ctx: &Context) -> Option { + self.tensor + .deref(ctx) + .downcast_ref::() + .cloned() + } + + /// Number of physical carrier elements in one shared stage. + #[must_use] + pub fn storage_elements(&self) -> Option { + u64::try_from(self.placement.0.outer().checked_size()?).ok() + } + + /// Physical byte size of one shared stage. + #[must_use] + pub fn storage_bytes(&self, ctx: &Context) -> Option { + self.storage_elements()? + .checked_mul(self.tensor_view(ctx)?.storage_bytes(ctx)?) + } +} + +impl Verify for CuteSmemTensorType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let Some(tensor) = self.tensor_view(ctx) else { + return verify_err!( + Location::Unknown, + "cute.smem_tensor must wrap a cute.tensor_view" + ); + }; + tensor.verify(ctx)?; + if tensor.space != CuteTensorAddressSpaceAttr::Smem + || tensor.access != CuteTensorAccessAttr::ReadOnly + { + return verify_err!( + Location::Unknown, + "cute.smem_tensor v0 must describe a read-only shared tensor" + ); + } + if !matches!( + (tensor.format, tensor.layout), + (CuteTensorFormatAttr::E2M1, CuteTensorLayoutAttr::KMajor) + | ( + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3, + CuteTensorLayoutAttr::BlockScaleKMajor(_) + ) + ) { + return verify_err!( + Location::Unknown, + "cute.smem_tensor needs packed E2M1 values or UE8M0/E4M3 block scales" + ); + } + if self.placement.0.unit() != crate::layout::OffsetUnit::Elements + || self.placement.0.offset() != 0 + || self.storage_elements().is_none() + { + return verify_err!( + Location::Unknown, + "cute.smem_tensor placement must be a positive zero-based element layout" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mod.rs new file mode 100644 index 0000000000..e7feda4915 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/mod.rs @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor-view types used to check and lower CuTe operations. +//! +//! A tensor view is a compiler description, not a runtime struct: +//! +//! ```text +//! pointer + length ── cute.tensor_make ──► !cute.tensor_view<...> +//! │ +//! └─ gone before LLVM lowering +//! ``` +//! +//! A selected backend consumes the description while lowering to native +//! operations or MLIR. A view has no runtime storage and must be removed +//! before the backend produces a kernel signature. +//! +//! After MIR preparation, the verifier follows each view through its uses: +//! +//! ```text +//! make -> zipped_divide -> slice -> load/store +//! ``` +//! +//! At this stage, a view cannot be stored inside a pointer or aggregate, +//! passed between blocks, or appear in a function signature. Ordinary scalar +//! values may remain in compiler-owned local storage whose address does not +//! escape, or pass between blocks. The verifier follows those values too. + +/// Shared output placement and asynchronous store synchronization contracts. +mod epilogue; +/// Packed-value and scale-view contracts for the block-scaled GEMV flow. +mod gemv; +/// Shared tensor placement, register carriers, and tiled matrix-instruction contracts. +mod mma; +/// Barrier-ring configuration and producer/consumer cursor contracts. +mod pipeline; +/// Persistent output-tile grids and their ordinary cursor provenance. +mod scheduler; +/// Tensor storage, access, format, and selection metadata. +mod tensor; +/// Descriptor-backed raw transport views and their direct copy contracts. +mod tma; + +pub use epilogue::CuteEpilogueTileType; +pub use gemv::{CuteFragmentType, CuteScaledViewType}; +pub use mma::{CuteSmemTensorType, MmaCarrierShape, mma_carrier_shape}; +pub use pipeline::CuteTmaLoadPipelineType; +pub use scheduler::CuteWorkTileType; +pub use tensor::CuteTensorViewType; +pub use tma::CuteTmaViewType; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/pipeline.rs new file mode 100644 index 0000000000..f013857dce --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/pipeline.rs @@ -0,0 +1,105 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Barrier-ring configuration and producer/consumer cursor contracts. + +use pliron::{ + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypedHandle}, + verify_err, +}; + +/// The shared barrier ring used by TMA loads. +/// +/// This is a compiler-only handle. The real storage is still the shared +/// pointer consumed by `cute.tma_load_pipeline_make`: +/// +/// ```text +/// shared barrier base ── make<3, 8 warps, 17408 bytes> +/// │ +/// ▼ +/// !cute.tma_load_pipeline<3, 8, 17408> +/// ``` +/// +/// The handle must disappear before ABI or LLVM lowering. Slot and phase are +/// deliberately *not* part of this type; ordinary `u32` values carry them +/// through the loop. +#[pliron_type( + name = "cute.tma_load_pipeline", + format = "`<` $stages `,` $consumer_warps `,` $transaction_bytes `>`" +)] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteTmaLoadPipelineType { + /// Number of reusable slots in the barrier or output-buffer ring. + pub stages: u64, + /// Consumer warps that release each full stage after reading it. + pub consumer_warps: u32, + /// Total bytes expected by each full barrier for one input stage. + pub transaction_bytes: u32, +} + +impl CuteTmaLoadPipelineType { + /// Create a compiler-only TMA load-pipeline handle. + pub fn get( + ctx: &mut Context, + stages: u64, + consumer_warps: u32, + transaction_bytes: u32, + ) -> TypedHandle { + Type::instantiate( + Self { + stages, + consumer_warps, + transaction_bytes, + }, + ctx, + ) + } + + /// Return bytes occupied by `[full; stages]` and `[empty; stages]`. + /// + /// Each hardware barrier occupies eight bytes. + #[must_use] + pub fn storage_bytes(&self) -> Option { + self.stages + .checked_mul(2) + .and_then(|count| count.checked_mul(8)) + } +} + +impl Verify for CuteTmaLoadPipelineType { + fn verify(&self, _ctx: &Context) -> Result<(), Error> { + if self.stages == 0 || self.stages > u64::from(u32::MAX) { + return verify_err!( + Location::Unknown, + "cute.tma_load_pipeline stages must be between 1 and {}", + u32::MAX + ); + } + if !(1..=32).contains(&self.consumer_warps) { + return verify_err!( + Location::Unknown, + "cute.tma_load_pipeline consumer warps must be between 1 and 32" + ); + } + if self.transaction_bytes == 0 { + return verify_err!( + Location::Unknown, + "cute.tma_load_pipeline transaction bytes must be greater than zero" + ); + } + if self.storage_bytes().is_none() { + return verify_err!( + Location::Unknown, + "cute.tma_load_pipeline barrier storage size overflowed u64" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/scheduler.rs new file mode 100644 index 0000000000..a72cdac820 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/scheduler.rs @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent output-tile grids and their ordinary cursor provenance. + +use pliron::{ + common_traits::Verify, + context::Context, + derive::pliron_type, + result::Error, + r#type::{Type, TypedHandle}, +}; + +use crate::attributes::CuteTileGridAttr; + +/// One logical output tile selected by a static scheduler. +/// +/// The type remembers the static tile grid. Its defining scheduler operation +/// carries the current linear tile number; no `WorkTileInfo` struct reaches +/// the runtime ABI. +/// +/// ```text +/// current linear number + grid<16,16,1> +/// │ +/// ▼ +/// !cute.work_tile> +/// │ work_tile_coordinates +/// ▼ +/// linear, m, n, batch +/// ``` +#[pliron_type(name = "cute.work_tile", format = "`<` $grid `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteWorkTileType { + /// Static M/N/batch dimensions of the output tile grid. + pub grid: CuteTileGridAttr, +} + +impl CuteWorkTileType { + /// Create the short-lived semantic handle for one scheduler iteration. + pub fn get(ctx: &mut Context, grid: CuteTileGridAttr) -> TypedHandle { + Type::instantiate(Self { grid }, ctx) + } +} + +impl Verify for CuteWorkTileType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + self.grid.verify(ctx) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tensor.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tensor.rs new file mode 100644 index 0000000000..7835e2ab5b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tensor.rs @@ -0,0 +1,356 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor storage, access, format, and selection metadata. + +use dialect_mir::types::{MirFP16Type, MirStructType}; +use pliron::{ + builtin::types::{FP32Type, IntegerType}, + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypeHandle, TypedHandle}, + verify_err, +}; + +use crate::attributes::{ + CuteAlignmentAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, CuteTensorFormatAttr, + CuteTensorLayoutAttr, CuteTensorRoleAttr, +}; + +/// A typed view of tensor storage. +/// +/// Read the fields left to right: +/// +/// ```text +/// logical storage format role space access alignment layout +/// f32 u8 E2M1 Mkl Gmem RO 1 KMajor +/// │ │ │ │ │ │ │ │ +/// decoded carrier bits coords where allowed base order +/// ``` +/// +/// Plain elementwise views have matching logical and storage types. Packed +/// views decode `u8` storage into logical `f32` values and say whether their +/// rows follow the M or N coordinate. +/// +/// Alignment here describes the storage pointer known at view construction. +/// A later unsafe transfer may promise a stronger selected-address alignment +/// on its own operation. +/// +/// This is a ghost type. It records meaning for compiler passes; it is never a +/// runtime descriptor and cannot be part of a kernel ABI. +#[pliron_type( + name = "cute.tensor_view", + format = "`<` $logical `,` $storage `,` $format `,` $role `,` $space `,` $access `,` $alignment `,` $layout `>`" +)] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteTensorViewType { + /// Decoded element type visible to tensor arithmetic. + pub logical: TypeHandle, + /// MIR element type physically stored in memory. + pub storage: TypeHandle, + /// Logical interpretation of the bits in each physical storage element. + pub format: CuteTensorFormatAttr, + /// Logical M/N coordinate role assigned to the tensor. + pub role: CuteTensorRoleAttr, + /// Address space in which the tensor storage resides. + pub space: CuteTensorAddressSpaceAttr, + /// Read-only or writable capability of the tensor view. + pub access: CuteTensorAccessAttr, + /// Promised storage-base alignment, in bytes. + pub alignment: CuteAlignmentAttr, + /// Static coordinate layout used to interpret the storage. + pub layout: CuteTensorLayoutAttr, +} + +impl CuteTensorViewType { + /// Create a ghost tensor-view type. + #[allow(clippy::too_many_arguments)] + pub fn get( + ctx: &mut Context, + logical: TypeHandle, + storage: TypeHandle, + space: CuteTensorAddressSpaceAttr, + access: CuteTensorAccessAttr, + alignment_bytes: u64, + layout: CuteTensorLayoutAttr, + ) -> TypedHandle { + Self::get_with_facts( + ctx, + logical, + storage, + space, + access, + alignment_bytes, + CuteTensorFormatAttr::Plain, + CuteTensorRoleAttr::Generic, + layout, + ) + } + + /// Create a tensor view with packed-format and matrix-role facts. + #[allow(clippy::too_many_arguments)] + pub fn get_with_facts( + ctx: &mut Context, + logical: TypeHandle, + storage: TypeHandle, + space: CuteTensorAddressSpaceAttr, + access: CuteTensorAccessAttr, + alignment_bytes: u64, + format: CuteTensorFormatAttr, + role: CuteTensorRoleAttr, + layout: CuteTensorLayoutAttr, + ) -> TypedHandle { + Type::instantiate( + Self { + logical, + storage, + format, + role, + space, + access, + alignment: CuteAlignmentAttr(alignment_bytes), + layout, + }, + ctx, + ) + } + + /// Copy every tensor fact except its layout. + #[must_use] + pub fn with_layout( + &self, + ctx: &mut Context, + layout: CuteTensorLayoutAttr, + ) -> TypedHandle { + Self::get_with_facts( + ctx, + self.logical, + self.storage, + self.space, + self.access, + self.alignment.0, + self.format, + self.role, + layout, + ) + } + + /// Scalar representation used to transport this element's unchanged bits. + /// + /// Logical and storage types retain transparent numeric wrappers. Only + /// copying looks through a rustc-proven transparent scalar with one field, + /// zero offset, and exactly the scalar's size and alignment. Other structs + /// cannot acquire a scalar copy representation merely from their size. + #[must_use] + pub fn transport_scalar(&self, ctx: &Context) -> Option { + scalar_transport_layout(ctx, self.storage, 0).map(|(scalar, _)| scalar) + } + + /// Size of one supported storage element in bytes. + #[must_use] + pub fn storage_bytes(&self, ctx: &Context) -> Option { + scalar_transport_layout(ctx, self.storage, 0).map(|(_, bytes)| bytes) + } + + /// Return the selected tile width, if this is a tile view. + #[must_use] + pub const fn selected_tile_size(&self) -> Option { + match self.layout { + CuteTensorLayoutAttr::Tile1D(size) => Some(size), + _ => None, + } + } +} + +/// Resolve scalar bits only through wrappers with a fully checked identical layout. +fn scalar_transport_layout( + ctx: &Context, + ty: TypeHandle, + depth: usize, +) -> Option<(TypeHandle, u64)> { + if depth > 16 { + return None; + } + let element = ty.deref(ctx); + if element.downcast_ref::().is_some() { + Some((ty, 2)) + } else if element.downcast_ref::().is_some() { + Some((ty, 4)) + } else if let Some(integer) = element.downcast_ref::() { + let width = u64::from(integer.width()); + (width > 0 && width % 8 == 0).then_some((ty, width / 8)) + } else if let Some(wrapper) = element.downcast_ref::() { + if !wrapper.is_transparent_scalar() + || wrapper.field_count() != 1 + || wrapper.field_names.len() != 1 + || wrapper.field_offsets.as_slice() != [0] + || wrapper.memory_order().as_slice() != [0] + { + return None; + } + let (scalar, bytes) = scalar_transport_layout(ctx, wrapper.field_types[0], depth + 1)?; + (wrapper.total_size == bytes && wrapper.abi_align == bytes).then_some((scalar, bytes)) + } else { + None + } +} + +impl Verify for CuteTensorViewType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let Some(storage_bytes) = self.storage_bytes(ctx) else { + return verify_err!( + Location::Unknown, + "cute.tensor_view storage must use a whole number of bytes" + ); + }; + let alignment = self.alignment.0; + if alignment == 0 || !alignment.is_power_of_two() { + return verify_err!( + Location::Unknown, + "cute.tensor_view alignment must be a positive power of two" + ); + } + if alignment < storage_bytes { + return verify_err!( + Location::Unknown, + "cute.tensor_view alignment cannot be smaller than one storage element" + ); + } + if self.layout.tile_size().is_some_and(|size| size == 0) { + return verify_err!( + Location::Unknown, + "cute.tensor_view tile size must be greater than zero" + ); + } + + let logical_is_f32 = self.logical.deref(ctx).downcast_ref::().is_some(); + let storage_is_u8 = self + .storage + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 8 && integer.is_unsigned()); + let storage_is_u32 = self + .storage + .deref(ctx) + .downcast_ref::() + .is_some_and(|integer| integer.width() == 32 && integer.is_unsigned()); + let storage_is_f16 = self + .storage + .deref(ctx) + .downcast_ref::() + .is_some(); + + match self.format { + CuteTensorFormatAttr::Plain => { + if self.logical != self.storage { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view needs matching logical and storage types" + ); + } + // Keep the wrapper in the typed IR; inspect only its validated + // physical scalar when deciding which bit copies are supported. + let transport = self + .transport_scalar(ctx) + .expect("storage width checked above"); + let storage_ty = transport.deref(ctx); + let plain_integer_is_supported = storage_ty + .downcast_ref::() + .is_some_and(|integer| { + integer.is_unsigned() && matches!(integer.width(), 8 | 16) + }); + let plain_is_supported = storage_ty.downcast_ref::().is_some() + || storage_ty.downcast_ref::().is_some() + || plain_integer_is_supported; + if !plain_is_supported { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view supports u8, u16, f16, and f32 carriers or proven transparent scalar wrappers only" + ); + } + if self.role != CuteTensorRoleAttr::Generic { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view must use the Generic role" + ); + } + if !matches!( + self.layout, + CuteTensorLayoutAttr::Contiguous1D + | CuteTensorLayoutAttr::Zipped1D(_) + | CuteTensorLayoutAttr::Tile1D(_) + | CuteTensorLayoutAttr::Tma2D + ) { + return verify_err!( + Location::Unknown, + "plain cute.tensor_view needs an elementwise or TMA transport layout" + ); + } + } + CuteTensorFormatAttr::E2M1 => { + let storage_matches_space = match self.space { + CuteTensorAddressSpaceAttr::Gmem => storage_is_u8, + CuteTensorAddressSpaceAttr::Smem => storage_is_f16, + }; + if !logical_is_f32 || !storage_matches_space { + return verify_err!( + Location::Unknown, + "E2M1 cute.tensor_view needs logical f32 values in u8 global or f16 shared storage" + ); + } + if self.access != CuteTensorAccessAttr::ReadOnly { + return verify_err!( + Location::Unknown, + "E2M1 cute.tensor_view must be read-only" + ); + } + if !matches!(self.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) { + return verify_err!( + Location::Unknown, + "E2M1 cute.tensor_view needs an Mkl or Nkl role" + ); + } + if self.layout != CuteTensorLayoutAttr::KMajor { + return verify_err!(Location::Unknown, "E2M1 cute.tensor_view must use KMajor"); + } + } + CuteTensorFormatAttr::UE8M0 | CuteTensorFormatAttr::E4M3 => { + let storage_matches_space = match self.space { + CuteTensorAddressSpaceAttr::Gmem => storage_is_u8, + CuteTensorAddressSpaceAttr::Smem => storage_is_u32, + }; + if !logical_is_f32 || !storage_matches_space { + return verify_err!( + Location::Unknown, + "UE8M0/E4M3 cute.tensor_view needs logical f32 scales in u8 global or u32 shared storage" + ); + } + if self.access != CuteTensorAccessAttr::ReadOnly { + return verify_err!( + Location::Unknown, + "UE8M0/E4M3 cute.tensor_view must be read-only" + ); + } + if !matches!(self.role, CuteTensorRoleAttr::Mkl | CuteTensorRoleAttr::Nkl) { + return verify_err!( + Location::Unknown, + "UE8M0/E4M3 cute.tensor_view needs an Mkl or Nkl role" + ); + } + if !matches!(self.layout, CuteTensorLayoutAttr::BlockScaleKMajor(size) if size > 0) + { + return verify_err!( + Location::Unknown, + "UE8M0/E4M3 cute.tensor_view must use BlockScaleKMajor" + ); + } + } + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tests.rs new file mode 100644 index 0000000000..b9ee86d054 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tests.rs @@ -0,0 +1,355 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe operation, type, and semantic verification contracts. + +use dialect_mir::types::{MirArrayType, MirFP16Type, MirStructType, MirTupleType, StructAbiKind}; +use pliron::{ + builtin::types::{FP32Type, IntegerType, Signedness}, + common_traits::Verify, + context::Context, + r#type::TypeHandle, +}; + +use super::*; +use crate::{ + attributes::{ + CuteScaledLayoutAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, + }, + layout::{ComposedLayout, Layout, OffsetUnit, Swizzle}, +}; + +/// The tensor type retains element, storage, access, alignment, and selection information together. +#[test] +fn tensor_view_records_the_whole_elementwise_view() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let view = CuteTensorViewType::get( + &mut ctx, + f32_ty, + f32_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Tile1D(4), + ); + + assert_eq!(view.deref(&ctx).selected_tile_size(), Some(4)); + assert!(view.deref(&ctx).verify(&ctx).is_ok()); +} + +/// Unsupported formats, storage types, or selection sizes fail type verification before lowering. +#[test] +fn tensor_view_rejects_invalid_v0_facts() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(&ctx).into(); + + let mismatched = CuteTensorViewType::get( + &mut ctx, + f32_ty, + f16_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Contiguous1D, + ); + assert!( + mismatched + .deref(&ctx) + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("matching logical and storage") + ); + + let zero_tile = CuteTensorViewType::get( + &mut ctx, + f32_ty, + f32_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Tile1D(0), + ); + assert!( + zero_tile + .deref(&ctx) + .verify(&ctx) + .unwrap_err() + .to_string() + .contains("tile size") + ); +} + +/// Transparent scalar wrappers preserve their logical identity while copying scalar bits. +#[test] +fn tensor_view_accepts_transparent_scalar_storage() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u16_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let mut element = u16_ty; + // Nested transparent wrappers are equally valid; no defining type name + // should be required to copy their unchanged physical representation. + for name in ["BFloat16", "TypedBFloat16"] { + element = MirStructType::get_with_full_layout_and_abi( + &mut ctx, + name.into(), + vec!["0".into()], + vec![element], + vec![0], + vec![0], + 2, + 2, + StructAbiKind::TransparentScalar, + ) + .into(); + let view = CuteTensorViewType::get( + &mut ctx, + element, + element, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadWrite, + 2, + CuteTensorLayoutAttr::Tile1D(8), + ); + assert_eq!(view.deref(&ctx).storage_bytes(&ctx), Some(2)); + assert!(view.deref(&ctx).verify(&ctx).is_ok()); + assert_eq!(view.deref(&ctx).logical, element); + assert_eq!(view.deref(&ctx).storage, element); + } +} + +/// A transparent ABI flag alone cannot authorize a scalar copy with a different layout. +#[test] +fn tensor_view_rejects_unproven_scalar_wrappers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u16_ty: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + for (name, count, offsets, size, alignment, abi) in [ + ("ordinary", 1, vec![0], 2, 2, StructAbiKind::Aggregate), + ( + "padding", + 1, + vec![0], + 4, + 2, + StructAbiKind::TransparentScalar, + ), + ("offset", 1, vec![1], 2, 2, StructAbiKind::TransparentScalar), + ( + "alignment", + 1, + vec![0], + 2, + 4, + StructAbiKind::TransparentScalar, + ), + ( + "extra_field", + 2, + vec![0, 0], + 2, + 2, + StructAbiKind::TransparentScalar, + ), + ( + "unknown_offset", + 1, + vec![], + 2, + 2, + StructAbiKind::TransparentScalar, + ), + ( + "unknown_size", + 1, + vec![0], + 0, + 2, + StructAbiKind::TransparentScalar, + ), + ( + "unknown_alignment", + 1, + vec![0], + 2, + 0, + StructAbiKind::TransparentScalar, + ), + ] { + let element: TypeHandle = MirStructType::get_with_full_layout_and_abi( + &mut ctx, + name.into(), + (0..count).map(|i| i.to_string()).collect(), + vec![u16_ty; count], + (0..count).collect(), + offsets, + size, + alignment, + abi, + ) + .into(); + let view = CuteTensorViewType::get( + &mut ctx, + element, + element, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Contiguous1D, + ); + assert!(view.deref(&ctx).verify(&ctx).is_err(), "accepted {name}"); + } +} + +/// Packed value/scale formats and coordinate roles survive row selection and fragment loading. +#[test] +fn packed_views_keep_format_role_scale_and_fragment_facts() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + + let values: TypeHandle = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + u8_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 16, + CuteTensorFormatAttr::E2M1, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::KMajor, + ) + .into(); + let scales: TypeHandle = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + u8_ty, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorFormatAttr::UE8M0, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::BlockScaleKMajor(16), + ) + .into(); + let tile: TypeHandle = CuteScaledViewType::get( + &mut ctx, + values, + scales, + CuteTensorRoleAttr::Mkl, + CuteScaledLayoutAttr::KTile(64), + ) + .into(); + let fragment = CuteFragmentType::get(&mut ctx, tile); + + assert!(values.deref(&ctx).verify(&ctx).is_ok()); + assert!(scales.deref(&ctx).verify(&ctx).is_ok()); + assert!(tile.deref(&ctx).verify(&ctx).is_ok()); + assert!(fragment.deref(&ctx).verify(&ctx).is_ok()); + let source = fragment.deref(&ctx).source_view(&ctx).unwrap(); + assert_eq!(source.role, CuteTensorRoleAttr::Mkl); + assert_eq!(source.values_per_scale(&ctx), Some(16)); + assert_eq!(source.k_width(), Some(64)); +} + +/// Shared packed data retains its physical f16/u32 carrier types rather than pretending to store decoded FP32. +#[test] +fn shared_packed_views_keep_their_real_transport_carriers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let f16_ty: TypeHandle = MirFP16Type::get(&ctx).into(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + + let values: TypeHandle = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + f16_ty, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + 16, + CuteTensorFormatAttr::E2M1, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::KMajor, + ) + .into(); + let scales: TypeHandle = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + u32_ty, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorFormatAttr::UE8M0, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::BlockScaleKMajor(32), + ) + .into(); + assert!(values.deref(&ctx).verify(&ctx).is_ok()); + assert!(scales.deref(&ctx).verify(&ctx).is_ok()); + + let data_inner: Layout = "(128,32):(32,1)".parse().unwrap(); + let data_placement = + ComposedLayout::new(Swizzle::new(2, 3, 3), 0, data_inner, OffsetUnit::Elements).unwrap(); + let shared = CuteSmemTensorType::get(&mut ctx, values, data_placement); + assert!(shared.deref(&ctx).verify(&ctx).is_ok()); + assert_eq!(shared.deref(&ctx).storage_elements(), Some(4096)); + assert_eq!(shared.deref(&ctx).storage_bytes(&ctx), Some(8192)); + + let wrong_shared_carrier = CuteTensorViewType::get_with_facts( + &mut ctx, + f32_ty, + u8_ty, + CuteTensorAddressSpaceAttr::Smem, + CuteTensorAccessAttr::ReadOnly, + 1, + CuteTensorFormatAttr::E2M1, + CuteTensorRoleAttr::Mkl, + CuteTensorLayoutAttr::KMajor, + ); + assert!(wrong_shared_carrier.deref(&ctx).verify(&ctx).is_err()); +} + +/// Register counting traverses nested arrays/structs while zero-sized marker fields contribute no registers. +#[test] +fn mma_carrier_shape_walks_nested_rust_aggregates_and_ignores_markers() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + crate::register(&mut ctx); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let f32_ty: TypeHandle = FP32Type::get(&ctx).into(); + let words: TypeHandle = MirArrayType::get(&mut ctx, u32_ty, 4).into(); + let marker: TypeHandle = + MirStructType::get(&mut ctx, "PhantomData".into(), vec![], vec![]).into(); + let fragment: TypeHandle = MirTupleType::get(&mut ctx, vec![words, marker]).into(); + let two_fragments: TypeHandle = MirArrayType::get(&mut ctx, fragment, 2).into(); + assert_eq!( + mma_carrier_shape(&ctx, two_fragments), + Some(MmaCarrierShape::u32(8)) + ); + + let mixed: TypeHandle = MirTupleType::get(&mut ctx, vec![two_fragments, f32_ty]).into(); + assert_eq!( + mma_carrier_shape(&ctx, mixed), + Some(MmaCarrierShape { + u32_registers: 8, + f32_registers: 1, + }) + ); + assert_eq!(mma_carrier_shape(&ctx, mixed).unwrap().bytes(), Some(36)); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tma.rs new file mode 100644 index 0000000000..0bb6d6f1ad --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/types/tma.rs @@ -0,0 +1,140 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Descriptor-backed raw transport views and their direct copy contracts. + +use pliron::{ + common_traits::Verify, + context::Context, + derive::pliron_type, + location::Location, + result::Error, + r#type::{Type, TypeHandle, TypedHandle}, + verify_err, +}; + +use super::tensor::CuteTensorViewType; +use crate::attributes::{ + CuteComposedLayoutAttr, CuteTensorFormatAttr, CuteTensorLayoutAttr, CuteTensorRoleAttr, +}; + +/// A raw-carrier tensor view used by one two-dimensional TMA transfer. +/// +/// TMA only moves storage boxes. It does not decode FP4 values, apply block +/// scales, or decide whether bytes belong to A or B: +/// +/// ```text +/// !cute.tensor_view +/// │ +/// └── exact shared placement +/// !cute.tma_view<..., layout> +/// ``` +/// +/// The same type wraps the descriptor-backed global source and the +/// pointer-backed shared destination. The inner tensor facts distinguish the +/// address space and access. `smem_layout` records the tile shape, row-major +/// stride, and optional hardware swizzle shared by both sides. +#[pliron_type(name = "cute.tma_view", format = "`<` $tensor `,` $smem_layout `>`")] +#[derive(Hash, PartialEq, Eq, Debug, Clone)] +pub struct CuteTmaViewType { + /// Wrapped tensor-view type containing storage, format, role, and access facts. + pub tensor: TypeHandle, + /// Exact shared tile shape, strides, offset unit, and swizzle. + pub smem_layout: CuteComposedLayoutAttr, +} + +impl CuteTmaViewType { + /// Attach one exact TMA tile placement to a plain carrier tensor. + pub fn get( + ctx: &mut Context, + tensor: TypeHandle, + smem_layout: crate::layout::ComposedLayout, + ) -> TypedHandle { + Type::instantiate( + Self { + tensor, + smem_layout: CuteComposedLayoutAttr(smem_layout), + }, + ctx, + ) + } + + /// Return the plain tensor facts wrapped by this TMA view. + #[must_use] + pub fn tensor_view(&self, ctx: &Context) -> Option { + self.tensor + .deref(ctx) + .downcast_ref::() + .cloned() + } + + /// Return the physical carrier type moved by TMA. + #[must_use] + pub fn element(&self, ctx: &Context) -> Option { + self.tensor_view(ctx).map(|view| view.storage) + } + + /// Return the number of carrier elements in one tile. + #[must_use] + pub fn tile_elements(&self) -> Option { + u64::try_from(self.smem_layout.0.outer().checked_size()?).ok() + } + + /// Return the physical byte count moved by one complete tile. + #[must_use] + pub fn tile_bytes(&self, ctx: &Context) -> Option { + let elements = self.tile_elements()?; + let bytes = self.tensor_view(ctx)?.storage_bytes(ctx)?; + elements.checked_mul(bytes) + } +} + +impl Verify for CuteTmaViewType { + fn verify(&self, ctx: &Context) -> Result<(), Error> { + let Some(tensor) = self.tensor_view(ctx) else { + return verify_err!( + Location::Unknown, + "cute.tma_view must wrap a cute.tensor_view" + ); + }; + tensor.verify(ctx)?; + if tensor.format != CuteTensorFormatAttr::Plain + || tensor.role != CuteTensorRoleAttr::Generic + || tensor.layout != CuteTensorLayoutAttr::Tma2D + { + return verify_err!( + Location::Unknown, + "cute.tma_view must move a Plain Generic Tma2D carrier tensor" + ); + } + if self.smem_layout.0.unit() != crate::layout::OffsetUnit::Elements { + return verify_err!( + Location::Unknown, + "cute.tma_view shared layout offsets must count carrier elements" + ); + } + let Some(element_bytes) = tensor.storage_bytes(ctx) else { + return verify_err!( + Location::Unknown, + "cute.tma_view carrier must have a known byte width" + ); + }; + let Ok(element_bytes) = i64::try_from(element_bytes) else { + return verify_err!( + Location::Unknown, + "cute.tma_view carrier byte width does not fit i64" + ); + }; + if let Err(error) = + crate::layout::validate_tma_encodable(&self.smem_layout.0, element_bytes) + { + return verify_err!( + Location::Unknown, + "cute.tma_view layout is not encodable by TMA: {error}" + ); + } + Ok(()) + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs new file mode 100644 index 0000000000..6a4da1be59 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs @@ -0,0 +1,286 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Function scopes, reachability, dominance, and source-order protocol checks. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::ops::{MirGotoOp, function::MirFuncOp}; +use pliron::{ + basic_block::BasicBlock, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, +}; + +use super::{Result, invalid}; + +/// The semantic verifier treats every MIR function as an independent story +/// namespace. Tests and tools may also place operations directly below a +/// module, in which case the verifier root is their namespace. +fn semantic_scope( + ctx: &Context, + root: Ptr, + operation: Ptr, +) -> Ptr { + let mut current = operation; + while let Some(parent) = current.deref(ctx).get_parent_op(ctx) { + if Operation::get_opid(parent, ctx) == MirFuncOp::get_opid_static() { + return parent; + } + if parent == root { + return root; + } + current = parent; + } + root +} + +/// Partition operations by their enclosing MIR function, keeping standalone IR in the root scope. +pub(super) fn operations_by_scope( + ctx: &Context, + root: Ptr, + all_ops: &[Ptr], +) -> Vec>> { + let mut scopes: HashMap, Vec>> = HashMap::new(); + for operation in all_ops { + let scope = semantic_scope(ctx, root, *operation); + scopes.entry(scope).or_default().push(*operation); + } + scopes.into_values().collect() +} + +/// Test whether any successor path in the same region reaches the destination block. +pub(super) fn block_reaches( + ctx: &Context, + start: Ptr, + destination: Ptr, +) -> bool { + let Some(region) = start.deref(ctx).get_parent_region() else { + return false; + }; + if destination.deref(ctx).get_parent_region() != Some(region) { + return false; + } + let mut worklist = vec![start]; + let mut visited = HashSet::new(); + while let Some(block) = worklist.pop() { + if !visited.insert(block) { + continue; + } + if block == destination { + return true; + } + let Some(terminator) = block.deref(ctx).get_terminator(ctx) else { + continue; + }; + worklist.extend(terminator.deref(ctx).successors()); + } + false +} + +/// Locate an operation's parent block and source-order index, or return `None` if detached. +fn operation_position( + ctx: &Context, + operation: Ptr, +) -> Option<(Ptr, usize)> { + let block = operation.deref(ctx).get_parent_block()?; + let index = block + .deref(ctx) + .iter(ctx) + .position(|candidate| candidate == operation)?; + Some((block, index)) +} + +/// Existential CFG reachability with source order inside one block. This is +/// used only to connect role-specific semantic milestones; each role's local +/// sequence is checked separately by `require_linear_protocol_order` where the +/// importer promises a straight-line call chain. +pub(super) fn operation_reaches( + ctx: &Context, + source: Ptr, + destination: Ptr, +) -> bool { + let Some((source_block, source_index)) = operation_position(ctx, source) else { + return false; + }; + let Some((destination_block, destination_index)) = operation_position(ctx, destination) else { + return false; + }; + if source_block == destination_block { + return source_index < destination_index; + } + let Some(terminator) = source_block.deref(ctx).get_terminator(ctx) else { + return false; + }; + terminator + .deref(ctx) + .successors() + .any(|successor| block_reaches(ctx, successor, destination_block)) +} + +/// Require source order in one block or dominance from the region entry across blocks. +pub(super) fn operation_dominates( + ctx: &Context, + source: Ptr, + destination: Ptr, +) -> bool { + let Some((source_block, source_index)) = operation_position(ctx, source) else { + return false; + }; + let Some((destination_block, destination_index)) = operation_position(ctx, destination) else { + return false; + }; + if source_block == destination_block { + return source_index < destination_index; + } + let Some(region) = source_block.deref(ctx).get_parent_region() else { + return false; + }; + if destination_block.deref(ctx).get_parent_region() != Some(region) { + return false; + } + let blocks: Vec<_> = region.deref(ctx).iter(ctx).collect(); + let Some(entry) = blocks.first().copied() else { + return false; + }; + let all: HashSet<_> = blocks.iter().copied().collect(); + let mut dominators: HashMap, HashSet>> = blocks + .iter() + .copied() + .map(|block| { + if block == entry { + (block, HashSet::from([entry])) + } else { + (block, all.clone()) + } + }) + .collect(); + + loop { + let mut changed = false; + for block in blocks.iter().copied().filter(|block| *block != entry) { + let predecessors: Vec<_> = block + .uses(ctx) + .into_iter() + .filter_map(|edge| edge.user_op().deref(ctx).get_parent_block()) + .collect(); + let mut next = if let Some(first) = predecessors.first() { + dominators.get(first).cloned().unwrap_or_default() + } else { + HashSet::new() + }; + for predecessor in predecessors.iter().skip(1) { + let predecessor_dominators = + dominators.get(predecessor).cloned().unwrap_or_default(); + next.retain(|candidate| predecessor_dominators.contains(candidate)); + } + next.insert(block); + if dominators.get(&block) != Some(&next) { + dominators.insert(block, next); + changed = true; + } + } + if !changed { + break; + } + } + dominators + .get(&destination_block) + .is_some_and(|set| set.contains(&source_block)) +} + +/// Follow one straight-line source sequence across the block splits emitted +/// for Rust call terminators. Every crossed edge must be an unconditional +/// goto into a block with no side entry. +pub(super) fn require_linear_protocol_order( + ctx: &Context, + label: &str, + ordered: &[Ptr], +) -> Result>> { + let Some(first) = ordered.first() else { + return Err(invalid(format!("{label} has no operations to order"))); + }; + let first_block = first + .deref(ctx) + .get_parent_block() + .ok_or_else(|| invalid(format!("{label} contains an operation outside a block")))?; + let region = first_block + .deref(ctx) + .get_parent_region() + .ok_or_else(|| invalid(format!("{label} starts in a block outside a region")))?; + if ordered.iter().any(|operation| { + operation + .deref(ctx) + .get_parent_block() + .and_then(|block| block.deref(ctx).get_parent_region()) + != Some(region) + }) { + return Err(invalid(format!( + "{label} operations must stay in one region" + ))); + } + + let mut current = first_block; + let mut visited = HashSet::new(); + let mut next_expected = 0usize; + let mut range = Vec::new(); + loop { + if !visited.insert(current) { + return Err(invalid(format!( + "{label} loops before reaching its final operation" + ))); + } + let operations: Vec<_> = current.deref(ctx).iter(ctx).collect(); + let start = if current == first_block { + operations + .iter() + .position(|operation| operation == first) + .ok_or_else(|| { + invalid(format!( + "{label} contains an operation missing from its parent block" + )) + })? + } else { + 0 + }; + for operation in operations.into_iter().skip(start) { + range.push(operation); + if ordered.get(next_expected) == Some(&operation) { + next_expected += 1; + if next_expected == ordered.len() { + return Ok(range); + } + } else if ordered.contains(&operation) { + return Err(invalid(format!( + "{label} is not in the required source order" + ))); + } + } + + let terminator = current + .deref(ctx) + .get_terminator(ctx) + .ok_or_else(|| invalid(format!("{label} ends a block before its final operation")))?; + if Operation::get_opid(terminator, ctx) != MirGotoOp::get_opid_static() + || terminator.deref(ctx).get_num_successors() != 1 + { + return Err(invalid(format!( + "{label} crosses an extra exit instead of one unconditional successor" + ))); + } + let successor = terminator.deref(ctx).get_successor(0); + if successor.deref(ctx).get_parent_region() != Some(region) { + return Err(invalid(format!("{label} leaves its source region"))); + } + if successor.preds(ctx).len() != 1 { + return Err(invalid(format!( + "{label} crosses a block with a side entry" + ))); + } + current = successor; + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs new file mode 100644 index 0000000000..66aae6d026 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs @@ -0,0 +1,691 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared output placement and asynchronous store synchronization contracts. + +use dialect_mir::{ + ops::{ + MirAddOp, MirCallOp, MirCastOp, MirConstantOp, MirConstructStructOp, MirConstructTupleOp, + MirExtractFieldOp, MirGotoOp, MirInsertFieldOp, MirLoadOp, MirRefOp, MirStorageDeadOp, + MirStorageLiveOp, MirUndefOp, + }, + types::{MirPtrType, MirStructType, address_space}, +}; +use pliron::{ + builtin::types::IntegerType, + common_traits::Verify, + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{operation_reaches, require_linear_protocol_order}, + ghosts::epilogue_semantic_ids, + invalid, + mma::{MmaFact, build_mma_facts, one_mma_fact}, + provenance::{is_semantic_pointer_carrier_cast, local_cell, local_cell_is_closed}, + tma::direct_tma_view_producer, +}; +use crate::{ + attributes::CuteEpilogueSyncPhaseAttr, + ops::epilogue::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + ops::gemm_tma::{CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + ops::mma::CuteFragmentFillOp, +}; + +/// Accept only pointer casts that preserve shared carrier provenance. +fn is_epilogue_pointer_identity_cast(ctx: &Context, operation: Ptr) -> bool { + is_semantic_pointer_carrier_cast(ctx, operation) +} + +/// Recognize carrier-only operations allowed between writer synchronization anchors. +fn is_epilogue_writer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + opid == MirGotoOp::get_opid_static() + || opid == MirStorageLiveOp::get_opid_static() + || opid == MirStorageDeadOp::get_opid_static() + || opid == MirExtractFieldOp::get_opid_static() + || is_epilogue_pointer_identity_cast(ctx, operation) + || is_closed_local_cell_load(ctx, operation) +} + +/// A non-volatile `mir.load` from a closed compiler-created local cell only +/// re-materializes carrier state (e.g. the epilogue handle struct spilled to +/// an alloca slot); it cannot touch tile memory, so it is admissible plumbing +/// between the writer-protocol anchor operations. +fn is_closed_local_cell_load(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirLoadOp::get_opid_static() { + return false; + } + let load = MirLoadOp::new(operation); + { + let operation_ref = operation.deref(ctx); + if operation_ref.get_num_operands() != 1 || operation_ref.get_num_results() != 1 { + return false; + } + } + if load.is_volatile(ctx) { + return false; + } + let Some((root, _path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + local_cell_is_closed(ctx, root) +} + +/// Accept only the dead immutable reference Rust leaves around a no-inline +/// call on the fieldless, zero-sized TMA store-pipeline token. +fn is_dead_epilogue_store_pipeline_ref(ctx: &Context, operation: Ptr) -> bool { + if Operation::get_opid(operation, ctx) != MirRefOp::get_opid_static() { + return false; + } + let reference = MirRefOp::new(operation); + if reference.verify(ctx).is_err() + || reference + .get_attr_mutable(ctx) + .is_none_or(|mutable| mutable.0) + { + return false; + } + let operation_ref = operation.deref(ctx); + let referent = operation_ref.get_operand(0); + let result = operation_ref.get_result(0); + if result.is_used(ctx) { + return false; + } + + let referent_ty = referent.get_type(ctx); + let referent_ty_ref = referent_ty.deref(ctx); + let Some(store_pipeline) = referent_ty_ref.downcast_ref::() else { + return false; + }; + if store_pipeline.name != "PipelineTmaStore" + || !store_pipeline.field_names.is_empty() + || !store_pipeline.field_types.is_empty() + || !store_pipeline.mem_to_decl.is_empty() + || !store_pipeline.field_offsets.is_empty() + || store_pipeline.total_size != 0 + || store_pipeline.abi_align != 1 + { + return false; + } + let result_ty = result.get_type(ctx); + let result_ty_ref = result_ty.deref(ctx); + let Some(pointer) = result_ty_ref.downcast_ref::() else { + return false; + }; + if pointer.pointee != referent_ty + || pointer.is_mutable + || pointer.address_space != address_space::GENERIC + { + return false; + } + + let Some(construct) = referent.defining_op() else { + return false; + }; + if Operation::get_opid(construct, ctx) != MirConstructStructOp::get_opid_static() { + return false; + } + let construct_ref = construct.deref(ctx); + if construct_ref.get_num_operands() != 0 + || construct_ref.get_num_results() != 1 + || construct_ref.get_result(0) != referent + { + return false; + } + let uses = referent.uses(ctx); + uses.len() == 1 && uses[0].user_op() == operation && uses[0].find_index(ctx) == 0 +} + +/// Recognize harmless carrier assembly and indexing between TMA issuer anchors. +fn is_epilogue_issuer_plumbing(ctx: &Context, operation: Ptr) -> bool { + let opid = Operation::get_opid(operation, ctx); + is_epilogue_writer_plumbing(ctx, operation) + || opid == MirUndefOp::get_opid_static() + || opid == MirInsertFieldOp::get_opid_static() + || opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static() + || opid == MirConstantOp::get_opid_static() + || opid == MirAddOp::get_opid_static() + || is_dead_epilogue_store_pipeline_ref(ctx, operation) +} + +/// Follow constant field insertions to the requested carrier, rejecting overly deep chains. +fn resolve_inserted_epilogue_field( + ctx: &Context, + mut aggregate: Value, + wanted_index: u32, +) -> Result> { + for _ in 0..16 { + let Some(definer) = aggregate.defining_op() else { + return Ok(None); + }; + if Operation::get_opid(definer, ctx) != MirInsertFieldOp::get_opid_static() { + return Ok(None); + } + let insert = MirInsertFieldOp::new(definer); + let Some(index) = insert.get_attr_insert_index(ctx) else { + return Err(invalid( + "epilogue carrier has an insert_field without a constant index", + )); + }; + let operation = definer.deref(ctx); + if index.0 == wanted_index { + return Ok(Some(operation.get_operand(1))); + } + aggregate = operation.get_operand(0); + } + Err(invalid( + "epilogue carrier insert/extract chain is unexpectedly deep", + )) +} + +/// Peel only lossless SSA carrier plumbing. Memory, arithmetic, control flow, +/// and arbitrary aliases deliberately stop provenance resolution. +fn resolve_epilogue_ssa_carrier(ctx: &Context, mut value: Value) -> Result { + for _ in 0..16 { + let Some(definer) = value.defining_op() else { + return Ok(value); + }; + let opid = Operation::get_opid(definer, ctx); + if opid == MirCastOp::get_opid_static() { + if is_semantic_pointer_carrier_cast(ctx, definer) { + value = definer.deref(ctx).get_operand(0); + continue; + } + return Ok(value); + } + if opid == MirExtractFieldOp::get_opid_static() { + let extract = MirExtractFieldOp::new(definer); + let Some(index) = extract.get_attr_index(ctx) else { + return Err(invalid( + "epilogue carrier has an extract_field without a constant index", + )); + }; + let aggregate = definer.deref(ctx).get_operand(0); + if let Some(aggregate_definer) = aggregate.defining_op() { + let aggregate_id = Operation::get_opid(aggregate_definer, ctx); + if aggregate_id == MirConstructStructOp::get_opid_static() + || aggregate_id == MirConstructTupleOp::get_opid_static() + { + let aggregate_operation = aggregate_definer.deref(ctx); + if usize::try_from(index.0) + .ok() + .is_some_and(|index| index < aggregate_operation.get_num_operands()) + { + value = aggregate_operation.get_operand(index.0 as usize); + continue; + } + } + } + let Some(inserted) = resolve_inserted_epilogue_field(ctx, aggregate, index.0)? else { + return Ok(value); + }; + value = inserted; + continue; + } + return Ok(value); + } + Err(invalid( + "epilogue carrier cast/insert/extract chain is unexpectedly deep", + )) +} + +/// Reject carriers that differ after peeling lossless SSA forwarding. +fn require_same_epilogue_carrier( + ctx: &Context, + actual: Value, + expected: Value, + message: &str, +) -> Result { + if resolve_epilogue_ssa_carrier(ctx, actual)? != resolve_epilogue_ssa_carrier(ctx, expected)? { + return Err(invalid(message)); + } + Ok(()) +} + +/// Read a direct unsigned 64-bit MIR literal used for epilogue coordinates. +fn epilogue_u64_constant(ctx: &Context, value: Value) -> Option { + let integer = value.get_type(ctx); + let integer = integer.deref(ctx); + if !integer + .downcast_ref::() + .is_some_and(|integer| integer.width() == 64 && integer.is_unsigned()) + { + return None; + } + let constant = value.defining_op()?; + if Operation::get_opid(constant, ctx) != MirConstantOp::get_opid_static() { + return None; + } + let attribute = MirConstantOp::new(constant).get_attr_value(ctx)?; + (attribute.value().bw() <= 64).then(|| attribute.value().to_u64()) +} + +/// Require the right output half's column to be the left column plus one. +fn require_next_epilogue_column(ctx: &Context, left: Value, right: Value) -> Result { + let left = resolve_epilogue_ssa_carrier(ctx, left)?; + let right = resolve_epilogue_ssa_carrier(ctx, right)?; + let Some(add) = right.defining_op() else { + return Err(invalid( + "epilogue half 1 must target the tile column immediately after half 0", + )); + }; + if Operation::get_opid(add, ctx) != MirAddOp::get_opid_static() { + return Err(invalid("epilogue half 1 must target `half_0_column + 1`")); + } + let add = add.deref(ctx); + let lhs = resolve_epilogue_ssa_carrier(ctx, add.get_operand(0))?; + let rhs = resolve_epilogue_ssa_carrier(ctx, add.get_operand(1))?; + if !((lhs == left && epilogue_u64_constant(ctx, rhs) == Some(1)) + || (rhs == left && epilogue_u64_constant(ctx, lhs) == Some(1))) + { + return Err(invalid("epilogue half 1 must target `half_0_column + 1`")); + } + Ok(()) +} + +/// Require reachable hand-offs from acquire through writers, issuer, commit, and tail. +pub(super) fn require_epilogue_cfg_links( + ctx: &Context, + acquire: Ptr, + reusable_sync: Ptr, + ready_sync: Ptr, + half_zero: Ptr, + commit: Ptr, + tail: Ptr, +) -> Result { + if operation_reaches(ctx, acquire, reusable_sync) + && operation_reaches(ctx, ready_sync, half_zero) + && operation_reaches(ctx, commit, tail) + { + Ok(()) + } else { + Err(invalid( + "epilogue CFG must connect acquire -> writer hand-off -> TMA issuer -> commit -> tail inside one role-anchored story", + )) + } +} + +/// A semantic output copy and the direct views supplying its storage. +#[derive(Clone, Copy)] +struct EpilogueStoreFacts { + /// The semantic shared-to-global copy being checked. + operation: Ptr, + /// The direct shared view producer carrying the output half. + source_view: Ptr, + /// The direct descriptor-backed global view producer. + destination_view: Ptr, +} + +/// Resolve both copy operands to direct TMA views inside the current function scope. +fn epilogue_store_facts( + ctx: &Context, + operation: Ptr, + scope: &[Ptr], +) -> Result { + let store = CuteTmaStore2dSemanticOp::wrap(operation); + let source_view = direct_tma_view_producer( + ctx, + store.source(ctx), + &CuteTmaSmemViewOp::get_opid_static(), + "epilogue store source", + scope, + )?; + let destination_view = direct_tma_view_producer( + ctx, + store.destination(ctx), + &CuteTmaGmemViewOp::get_opid_static(), + "epilogue store destination", + scope, + )?; + Ok(EpilogueStoreFacts { + operation, + source_view, + destination_view, + }) +} + +/// Validate the single two-half output protocol, carrier origins, and role-local ordering. +pub(super) fn verify_epilogue_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = epilogue_semantic_ids(); + let mut overlays = Vec::new(); + let mut slices = Vec::new(); + let mut fragments = Vec::new(); + let mut syncs = Vec::new(); + let mut halves = Vec::new(); + let mut acquires = Vec::new(); + let mut commits = Vec::new(); + let mut tails = Vec::new(); + let mut stores = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if opid == ids[0] { + overlays.push(*operation); + } else if opid == ids[1] { + slices.push(*operation); + } else if opid == ids[2] { + fragments.push(*operation); + } else if opid == ids[3] { + syncs.push(*operation); + } else if opid == ids[4] { + halves.push(*operation); + } else if opid == ids[5] { + acquires.push(*operation); + } else if opid == ids[6] { + commits.push(*operation); + } else if opid == ids[7] { + tails.push(*operation); + } else if opid == ids[8] { + stores.push(*operation); + } + } + + let found = overlays.len() + + slices.len() + + fragments.len() + + syncs.len() + + halves.len() + + acquires.len() + + commits.len() + + tails.len() + + stores.len(); + if found == 0 { + return Ok(()); + } + if overlays.len() != 1 + || slices.len() != 1 + || fragments.len() != 1 + || syncs.len() != 2 + || halves.len() != 2 + || acquires.len() != 1 + || commits.len() != 1 + || tails.len() != 1 + || stores.len() != 2 + { + return Err(invalid(format!( + "epilogue story needs exactly 1 overlay, 1 warp slice, 1 fragment store, 2 syncs, 2 halves, 1 acquire, 2 TMA stores, 1 commit, and 1 tail; found {}, {}, {}, {}, {}, {}, {}, {}, {}", + overlays.len(), + slices.len(), + fragments.len(), + syncs.len(), + halves.len(), + acquires.len(), + stores.len(), + commits.len(), + tails.len() + ))); + } + + let overlay = overlays[0]; + let warp_slice = slices[0]; + let store_fragment = fragments[0]; + let acquire = acquires[0]; + let commit = commits[0]; + let tail = tails[0]; + let overlay_view = CuteEpilogueSmemOverlayOp::wrap(overlay); + let warp_view = CuteEpilogueWarpSliceOp::wrap(warp_slice); + let fragment = CuteEpilogueStoreFragmentOp::wrap(store_fragment); + let overlay_tile = overlay_view + .tile(ctx) + .expect("locally verified epilogue overlay has a tile"); + if warp_view.tile(ctx).as_ref() != Some(&overlay_tile) + || fragment.tile(ctx).as_ref() != Some(&overlay_tile) + { + return Err(invalid( + "epilogue overlay, warp slice, and fragment store must carry the same tile plan", + )); + } + let mma_facts = build_mma_facts(ctx, all_ops); + let accumulator_seed = match one_mma_fact( + &mma_facts, + fragment.accumulator(ctx), + "epilogue fragment accumulator", + )? { + MmaFact::Accumulator(seed) => seed, + _ => { + return Err(invalid( + "epilogue fragment accumulator must be the closed tiled-GEMM recurrence", + )); + } + }; + if CuteFragmentFillOp::wrap(accumulator_seed).plan(ctx) + != Some(overlay_tile.plan.tiled_mma.clone()) + { + return Err(invalid( + "epilogue accumulator plan differs from its shared-MMA story", + )); + } + require_same_epilogue_carrier( + ctx, + warp_view.input_base(ctx), + overlay_view.base(ctx), + "epilogue warp slice must use the base produced by the epilogue overlay", + )?; + require_same_epilogue_carrier( + ctx, + fragment.base(ctx), + warp_view.base(ctx), + "epilogue fragment store must use the base produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.warp_id(ctx), + warp_view.warp_id(ctx), + "epilogue fragment store must use the warp ID produced by the warp slice", + )?; + require_same_epilogue_carrier( + ctx, + fragment.lane(ctx), + warp_view.lane(ctx), + "epilogue fragment store must use the lane produced by the warp slice", + )?; + + let reusable_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::Reusable) + }) + .ok_or_else(|| invalid("epilogue story is missing the Reusable sync"))?; + let ready_sync = syncs + .iter() + .copied() + .find(|operation| { + CuteEpilogueSyncOp::wrap(*operation).phase(ctx) + == Some(CuteEpilogueSyncPhaseAttr::ReadyForTma) + }) + .ok_or_else(|| invalid("epilogue story is missing the ReadyForTma sync"))?; + if reusable_sync == ready_sync { + return Err(invalid( + "epilogue story needs one distinct sync for each phase", + )); + } + for sync in [reusable_sync, ready_sync] { + let sync = CuteEpilogueSyncOp::wrap(sync); + if sync.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue syncs must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + sync.base(ctx), + overlay_view.base(ctx), + "both epilogue syncs must use the base produced by the epilogue overlay", + )?; + } + + let writer_range = require_linear_protocol_order( + ctx, + "epilogue writer protocol", + &[reusable_sync, store_fragment, ready_sync], + )?; + for operation in &writer_range { + let expected = [reusable_sync, store_fragment, ready_sync].contains(operation); + if !expected && !is_epilogue_writer_plumbing(ctx, *operation) { + return Err(invalid(format!( + "epilogue writer protocol has interposed non-carrier operation `{}`", + Operation::get_opid(*operation, ctx) + ))); + } + } + let writer_events: Vec<_> = writer_range + .iter() + .copied() + .filter(|operation| { + let opid = Operation::get_opid(*operation, ctx); + opid.dialect.to_string() == crate::CUTE_DIALECT_NAME + || opid == MirCallOp::get_opid_static() + }) + .collect(); + if writer_events != [reusable_sync, store_fragment, ready_sync] { + return Err(invalid( + "epilogue writer protocol must be exactly Reusable -> fragment store -> ReadyForTma, whose semantic boundary includes publication to the TMA proxy", + )); + } + + let half_zero = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 0) + }) + .ok_or_else(|| invalid("epilogue story is missing half 0"))?; + let half_one = halves + .iter() + .copied() + .find(|operation| { + CuteEpilogueHalfOp::wrap(*operation) + .half(ctx) + .is_some_and(|half| half.0 == 1) + }) + .ok_or_else(|| invalid("epilogue story is missing half 1"))?; + if half_zero == half_one { + return Err(invalid("epilogue story needs two distinct shared halves")); + } + for half in [half_zero, half_one] { + let half = CuteEpilogueHalfOp::wrap(half); + if half.tile(ctx).as_ref() != Some(&overlay_tile) { + return Err(invalid( + "both epilogue halves must carry the overlay tile plan", + )); + } + require_same_epilogue_carrier( + ctx, + half.full_base(ctx), + overlay_view.base(ctx), + "both epilogue halves must use the base produced by the epilogue overlay", + )?; + } + + let store_zero = epilogue_store_facts(ctx, stores[0], all_ops)?; + let store_one = epilogue_store_facts(ctx, stores[1], all_ops)?; + for (half, store) in [(half_zero, store_zero), (half_one, store_one)] { + let half = CuteEpilogueHalfOp::wrap(half); + let source = CuteTmaSmemViewOp::wrap(store.source_view); + require_same_epilogue_carrier( + ctx, + source.base(ctx), + half.half_base(ctx), + "each semantic TMA store source base must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + require_same_epilogue_carrier( + ctx, + source.capacity(ctx), + half.capacity(ctx), + "each semantic TMA store source capacity must come from the matching epilogue half through only the compiler's SSA carrier adapter", + )?; + } + if store_zero.operation.deref(ctx).get_operand(1).get_type(ctx) + != store_one.operation.deref(ctx).get_operand(1).get_type(ctx) + { + return Err(invalid( + "both epilogue stores must target the same typed output tile", + )); + } + let destination_zero = CuteTmaGmemViewOp::wrap(store_zero.destination_view); + let destination_one = CuteTmaGmemViewOp::wrap(store_one.destination_view); + require_same_epilogue_carrier( + ctx, + destination_one.descriptor(ctx), + destination_zero.descriptor(ctx), + "both epilogue halves must target the same TMA destination descriptor", + )?; + let semantic_zero = CuteTmaStore2dSemanticOp::wrap(store_zero.operation); + let semantic_one = CuteTmaStore2dSemanticOp::wrap(store_one.operation); + require_same_epilogue_carrier( + ctx, + semantic_one.tile_row(ctx), + semantic_zero.tile_row(ctx), + "both epilogue halves must target the same output tile row", + )?; + require_next_epilogue_column( + ctx, + semantic_zero.tile_column(ctx), + semantic_one.tile_column(ctx), + )?; + + let issuer_range = require_linear_protocol_order( + ctx, + "epilogue TMA issuer protocol", + &[ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ], + )?; + for operation in issuer_range { + let opid = Operation::get_opid(operation, ctx); + let allowed_tma_view = opid == CuteTmaGmemViewOp::get_opid_static() + || opid == CuteTmaSmemViewOp::get_opid_static(); + let expected = [ + half_zero, + store_zero.operation, + half_one, + store_one.operation, + commit, + ] + .contains(&operation); + if !allowed_tma_view && !expected && !is_epilogue_issuer_plumbing(ctx, operation) { + return Err(invalid(format!( + "epilogue TMA issuer protocol has interposed operation `{opid}`" + ))); + } + } + + let acquire_pipeline = CuteTmaStoreAcquireOp::wrap(acquire).pipeline(ctx); + let commit_pipeline = CuteTmaStoreCommitOp::wrap(commit).pipeline(ctx); + let tail_pipeline = CuteTmaStoreTailOp::wrap(tail).pipeline(ctx); + if acquire_pipeline != commit_pipeline || acquire_pipeline != tail_pipeline { + return Err(invalid( + "epilogue acquire, commit, and tail must carry the same TMA store-pipeline facts", + )); + } + + require_epilogue_cfg_links( + ctx, + acquire, + reusable_sync, + ready_sync, + half_zero, + commit, + tail, + ) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/gemv.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/gemv.rs new file mode 100644 index 0000000000..47a2111330 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/gemv.rs @@ -0,0 +1,212 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Packed-value and scale-view contracts for the block-scaled GEMV flow. + +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{Result, ghosts::gemv_semantic_ids, invalid}; +use crate::{ + attributes::CuteTensorRoleAttr, + ops::gemv::{ + CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, CuteScaledViewRowOp, + CuteTensorMake2DOp, + }, + types::CuteTensorViewType, +}; + +/// Runtime extents and coordinate role recovered from one tensor constructor. +#[derive(Clone, Copy)] +struct GemvTensorState { + /// Runtime row-count SSA value of the original tensor. + rows: Value, + /// Runtime K-extent SSA value of the original tensor. + k: Value, + /// Whether rows select matrix M or vector N coordinates. + role: CuteTensorRoleAttr, +} + +/// The original value/scale tensors and the selections applied to their joint view. +#[derive(Clone, Copy)] +struct GemvSelectionState { + /// Origin and extents of the packed value tensor. + values: GemvTensorState, + /// Origin and extents of the scale tensor. + scales: GemvTensorState, + /// Selected batch coordinate, absent on the full view. + batch: Option, + /// Selected row coordinate, absent before row selection. + row: Option, + /// Selected K-tile coordinate, absent before tiling. + tile_index: Option, +} + +/// Require a direct two-dimensional tensor constructor and recover its extents and role. +fn resolve_gemv_tensor(ctx: &Context, value: Value) -> Result { + let Some(producer) = value.defining_op() else { + return Err(invalid( + "GEMV tensor reached a block argument; direct semantic SSA is required", + )); + }; + let opid = Operation::get_opid(producer, ctx); + if opid != CuteTensorMake2DOp::get_opid_static() { + return Err(invalid(format!( + "GEMV tensor is produced by `{opid}` instead of `cute.tensor_make_2d`" + ))); + } + let op = producer.deref(ctx); + let ty = value.get_type(ctx); + let ty = ty.deref(ctx); + let view = ty + .downcast_ref::() + .ok_or_else(|| invalid("cute.tensor_make_2d result is not a tensor view"))?; + Ok(GemvTensorState { + rows: op.get_operand(2), + k: op.get_operand(3), + role: view.role, + }) +} + +/// Walk the make/row/K-tile chain, bounding recursion and preserving selection identity. +fn resolve_gemv_selection(ctx: &Context, value: Value, depth: usize) -> Result { + if depth > 16 { + return Err(invalid( + "GEMV scaled-view producer chain is unexpectedly deep", + )); + } + let Some(producer) = value.defining_op() else { + return Err(invalid( + "GEMV scaled view reached a block argument; direct semantic SSA is required", + )); + }; + let opid = Operation::get_opid(producer, ctx); + let op = producer.deref(ctx); + if opid == CuteScaledViewMakeOp::get_opid_static() { + let values = resolve_gemv_tensor(ctx, op.get_operand(0))?; + let scales = resolve_gemv_tensor(ctx, op.get_operand(1))?; + if values.rows != scales.rows || values.k != scales.k { + return Err(invalid( + "cute.scaled_view_make values and scales must use the same runtime rows and K operands", + )); + } + return Ok(GemvSelectionState { + values, + scales, + batch: None, + row: None, + tile_index: None, + }); + } + if opid == CuteScaledViewRowOp::get_opid_static() { + let mut state = resolve_gemv_selection(ctx, op.get_operand(0), depth + 1)?; + if state.batch.is_some() || state.row.is_some() || state.tile_index.is_some() { + return Err(invalid( + "cute.scaled_view_row must select one row before selecting a K tile", + )); + } + state.batch = Some(op.get_operand(1)); + state.row = Some(op.get_operand(2)); + return Ok(state); + } + if opid == CuteScaledViewKTileOp::get_opid_static() { + let mut state = resolve_gemv_selection(ctx, op.get_operand(0), depth + 1)?; + if state.batch.is_none() || state.row.is_none() || state.tile_index.is_some() { + return Err(invalid( + "cute.scaled_view_k_tile must follow exactly one row selection", + )); + } + state.tile_index = Some(op.get_operand(1)); + return Ok(state); + } + Err(invalid(format!( + "scaled GEMV view is produced by unsupported operation `{opid}`" + ))) +} + +/// Require a load from one complete, consistently selected block-scaled K tile. +fn resolve_gemv_fragment(ctx: &Context, value: Value) -> Result { + let Some(load) = value.defining_op() else { + return Err(invalid( + "GEMV fragment reached a block argument; a direct load -> dot edge is required", + )); + }; + let opid = Operation::get_opid(load, ctx); + if opid != CuteScaledViewLoadOp::get_opid_static() { + return Err(invalid(format!( + "GEMV fragment is produced by `{opid}` instead of `cute.scaled_view_load`" + ))); + } + let typed = CuteScaledViewLoadOp::wrap(load); + let value_alignment = typed + .promised_value_alignment(ctx) + .ok_or_else(|| invalid("cute.scaled_view_load is missing its value alignment promise"))?; + let scale_alignment = typed + .promised_scale_alignment(ctx) + .ok_or_else(|| invalid("cute.scaled_view_load is missing its scale alignment promise"))?; + if value_alignment < 16 || scale_alignment < 4 { + return Err(invalid( + "GEMV K=64 load needs value alignment >= 16 and scale alignment >= 4", + )); + } + let state = resolve_gemv_selection(ctx, load.deref(ctx).get_operand(0), 0)?; + if state.batch.is_none() || state.row.is_none() || state.tile_index.is_none() { + return Err(invalid("GEMV load needs a selected row and K tile")); + } + if state.values.role != state.scales.role { + return Err(invalid("GEMV values and scales must keep the same role")); + } + Ok(()) +} + +/// Require the producer's ghost result to have exactly one user of the expected operation kind. +fn require_one_user(ctx: &Context, producer: Ptr, expected: &OpId) -> Result { + let result = producer.deref(ctx).get_result(0); + let uses = result.uses(ctx); + if uses.len() != 1 + || Operation::get_opid(uses[0].user_op(), ctx) != *expected + || uses[0].find_index(ctx) >= uses[0].user_op().deref(ctx).get_num_operands() + { + return Err(invalid(format!( + "GEMV semantic value from `{}` must have exactly one `{expected}` consumer", + Operation::get_opid(producer, ctx) + ))); + } + Ok(()) +} + +/// Check direct ghost chains and compatible matrix/vector selections for each dot product. +pub(super) fn verify_gemv_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = gemv_semantic_ids(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + let expected = if opid == ids[0] { + Some(&ids[1]) + } else if opid == ids[1] { + Some(&ids[2]) + } else if opid == ids[2] { + Some(&ids[3]) + } else if opid == ids[3] { + Some(&ids[4]) + } else if opid == ids[4] { + Some(&ids[5]) + } else { + None + }; + if let Some(expected) = expected { + require_one_user(ctx, *operation, expected)?; + } else if opid == ids[5] { + let op = operation.deref(ctx); + resolve_gemv_fragment(ctx, op.get_operand(0))?; + resolve_gemv_fragment(ctx, op.get_operand(1))?; + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/ghosts.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/ghosts.rs new file mode 100644 index 0000000000..405c9c311e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/ghosts.rs @@ -0,0 +1,304 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Restrictions preventing compiler-only CuTe descriptions from entering runtime storage. + +use std::collections::HashSet; + +use dialect_mir::types::{ + MirArrayType, MirDisjointSliceType, MirEnumType, MirPtrType, MirSliceType, MirStructType, + MirTupleType, MirUnionType, +}; +use pliron::{ + attribute::attr_cast, + builtin::{attr_interfaces::TypedAttrInterface, type_interfaces::FunctionTypeInterface}, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::{Op, OpId}, + operation::Operation, + r#type::{TypeHandle, Typed, type_cast}, +}; + +use super::{Result, invalid}; +use crate::{ + ops::epilogue::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + ops::gemm_tma::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + ops::gemv::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, + }, + ops::mma::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, + }, + ops::pipeline::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, + }, + ops::scheduler::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + ops::tensor::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, + }, + types::{ + CuteEpilogueTileType, CuteFragmentType, CuteScaledViewType, CuteSmemTensorType, + CuteTensorViewType, CuteTmaLoadPipelineType, CuteTmaViewType, CuteWorkTileType, + }, +}; + +/// List the tensor operations allowed to consume or produce direct tensor ghosts. +pub(super) fn tensor_semantic_ids() -> [OpId; 8] { + [ + CuteTensorMakeOp::get_opid_static(), + CuteTensorZippedDivideOp::get_opid_static(), + CuteTensorSliceOp::get_opid_static(), + CuteTensorIsFullOp::get_opid_static(), + CuteTensorBaseOp::get_opid_static(), + CuteTensorLoadIntoOp::get_opid_static(), + CuteTensorStoreFromOp::get_opid_static(), + CuteTensorStoreElementAbsOp::get_opid_static(), + ] +} + +/// List the packed-value view and fragment operations covered by the GEMV audit. +pub(super) fn gemv_semantic_ids() -> [OpId; 6] { + [ + CuteTensorMake2DOp::get_opid_static(), + CuteScaledViewMakeOp::get_opid_static(), + CuteScaledViewRowOp::get_opid_static(), + CuteScaledViewKTileOp::get_opid_static(), + CuteScaledViewLoadOp::get_opid_static(), + CuteDotOp::get_opid_static(), + ] +} + +/// List scheduler operations participating in work-tile provenance checks. +pub(super) fn scheduler_semantic_ids() -> [OpId; 5] { + [ + CuteSchedulerNew1dOp::get_opid_static(), + CuteSchedulerHasWorkOp::get_opid_static(), + CuteSchedulerCurrentOp::get_opid_static(), + CuteWorkTileCoordinatesOp::get_opid_static(), + CuteSchedulerAdvanceOp::get_opid_static(), + ] +} + +/// List load-pipeline and cursor operations allowed by the lifecycle audit. +pub(super) fn pipeline_semantic_ids() -> [OpId; 10] { + [ + CuteTmaLoadPipelineMakeOp::get_opid_static(), + CuteTmaLoadPipelineInitOp::get_opid_static(), + CutePipelineStateNewOp::get_opid_static(), + CutePipelineStateSlotOp::get_opid_static(), + CutePipelineStateAdvanceOp::get_opid_static(), + CutePipelineProducerAcquireOp::get_opid_static(), + CutePipelineProducerExpectTxOp::get_opid_static(), + CutePipelineConsumerWaitOp::get_opid_static(), + CutePipelineConsumerReleaseOp::get_opid_static(), + CutePipelineProducerTailOp::get_opid_static(), + ] +} + +/// List direct carrier-view constructors and their semantic TMA copy consumer. +fn tma_semantic_ids() -> [OpId; 3] { + [ + CuteTmaGmemViewOp::get_opid_static(), + CuteTmaSmemViewOp::get_opid_static(), + CuteTmaCopy2dOp::get_opid_static(), + ] +} + +/// List operations whose ordinary carriers carry shared-memory MMA meaning. +pub(super) fn smem_mma_semantic_ids() -> [OpId; 8] { + [ + CuteSmemTensorOverlayOp::get_opid_static(), + CuteTiledMmaSliceOp::get_opid_static(), + CuteFragmentFillOp::get_opid_static(), + CuteMmaLoadScalesOp::get_opid_static(), + CuteFragmentSliceKOp::get_opid_static(), + CuteMmaLoadAOp::get_opid_static(), + CuteMmaPartitionBOp::get_opid_static(), + CuteTiledGemmOp::get_opid_static(), + ] +} + +/// List output-tile, synchronization, and asynchronous store protocol operations. +pub(super) fn epilogue_semantic_ids() -> [OpId; 9] { + [ + CuteEpilogueSmemOverlayOp::get_opid_static(), + CuteEpilogueWarpSliceOp::get_opid_static(), + CuteEpilogueStoreFragmentOp::get_opid_static(), + CuteEpilogueSyncOp::get_opid_static(), + CuteEpilogueHalfOp::get_opid_static(), + CuteTmaStoreAcquireOp::get_opid_static(), + CuteTmaStoreCommitOp::get_opid_static(), + CuteTmaStoreTailOp::get_opid_static(), + CuteTmaStore2dSemanticOp::get_opid_static(), + ] +} + +/// Collect the closed set of operations allowed to carry semantic ghosts or metadata. +fn all_high_level_semantic_ids() -> Vec { + let mut ids = tensor_semantic_ids().to_vec(); + ids.push(crate::ops::copy::CuteStoreGlobalNoAllocateOp::get_opid_static()); + ids.extend(gemv_semantic_ids()); + ids.extend(scheduler_semantic_ids()); + ids.extend(pipeline_semantic_ids()); + ids.extend(tma_semantic_ids()); + ids.extend(smem_mma_semantic_ids()); + ids.extend(epilogue_semantic_ids()); + ids.extend(crate::ops::sm100::semantic_ids()); + ids.extend(crate::ops::sm100_attention::semantic_ids()); + ids +} + +/// Recognize ghost SSA types, excluding types permitted only as static metadata. +fn is_direct_cute_ghost_type(ctx: &Context, ty: TypeHandle) -> bool { + let ty = ty.deref(ctx); + ty.downcast_ref::().is_some() + || ty.downcast_ref::().is_some() + || ty.downcast_ref::().is_some() + || ty.downcast_ref::().is_some() + || ty.downcast_ref::().is_some() + || ty.downcast_ref::().is_some() +} + +/// `true` when `ty` is, or structurally contains, a CuTe compiler-only +/// (ghost) type. Ghost values must remain direct semantic SSA; preparation +/// passes use this to decide which slots must be dissolved before this +/// module's semantic verification runs. +pub fn type_contains_cute_ghost(ctx: &Context, ty: TypeHandle) -> bool { + /// Recursively inspect nested type structure for ghosts, using a visited set to break recursive cycles. + fn walk(ctx: &Context, ty: TypeHandle, seen: &mut HashSet) -> bool { + if !seen.insert(ty) { + return false; + } + let ty_ref = ty.deref(ctx); + if ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + || ty_ref.downcast_ref::().is_some() + { + return true; + } + if let Some(function) = type_cast::(&*ty_ref) { + return function + .arg_types() + .into_iter() + .chain(function.res_types()) + .any(|child| walk(ctx, child, seen)); + } + if let Some(pointer) = ty_ref.downcast_ref::() { + return walk(ctx, pointer.pointee, seen); + } + if let Some(tuple) = ty_ref.downcast_ref::() { + return tuple.types.iter().any(|child| walk(ctx, *child, seen)); + } + if let Some(structure) = ty_ref.downcast_ref::() { + return structure + .field_types + .iter() + .any(|child| walk(ctx, *child, seen)); + } + if let Some(union) = ty_ref.downcast_ref::() { + return union + .field_types + .iter() + .any(|child| walk(ctx, *child, seen)); + } + if let Some(array) = ty_ref.downcast_ref::() { + return walk(ctx, array.element_ty, seen); + } + if let Some(slice) = ty_ref.downcast_ref::() { + return walk(ctx, slice.element_ty, seen); + } + if let Some(slice) = ty_ref.downcast_ref::() { + return walk(ctx, slice.element_ty, seen) + || slice.space_tys.iter().any(|child| walk(ctx, *child, seen)); + } + if let Some(enumeration) = ty_ref.downcast_ref::() { + return walk(ctx, enumeration.discriminant_ty, seen) + || enumeration + .all_field_types + .iter() + .any(|child| walk(ctx, *child, seen)); + } + false + } + + walk(ctx, ty, &mut HashSet::new()) +} + +/// Reject ghost values nested in runtime carriers, signatures, block arguments, or unrelated operations. +pub(super) fn audit_cute_ghosts(ctx: &Context, all_ops: &[Ptr]) -> Result { + let semantic_ids = all_high_level_semantic_ids(); + let escape = |where_: String| { + invalid(format!( + "CuTe compiler-only values must remain direct semantic SSA or typed semantic metadata; a ghost CuTe value escaped in {where_}" + )) + }; + + for op_ptr in all_ops { + let opid = Operation::get_opid(*op_ptr, ctx); + let operation = op_ptr.deref(ctx); + for value in operation.operands().chain(operation.results()) { + let ty = value.get_type(ctx); + if !type_contains_cute_ghost(ctx, ty) { + continue; + } + if !semantic_ids.contains(&opid) || !is_direct_cute_ghost_type(ctx, ty) { + return Err(escape(format!("an operand or result of `{opid}`"))); + } + } + for (name, attribute) in &operation.attributes.0 { + let Some(typed) = attr_cast::(&**attribute) else { + continue; + }; + let ty = typed.get_type(ctx); + if !type_contains_cute_ghost(ctx, ty) { + continue; + } + let direct_metadata = ty.deref(ctx).downcast_ref::().is_some() + || ty + .deref(ctx) + .downcast_ref::() + .is_some(); + if !semantic_ids.contains(&opid) || !direct_metadata { + return Err(escape(format!("attribute `{name}` on `{opid}`"))); + } + } + for region in operation.regions() { + for block in region.deref(ctx).iter(ctx) { + for argument in block.deref(ctx).arguments() { + if type_contains_cute_ghost(ctx, argument.get_type(ctx)) { + return Err(escape(format!("a block argument below `{opid}`"))); + } + } + for (name, attribute) in &block.deref(ctx).attributes.0 { + if let Some(typed) = attr_cast::(&**attribute) + && type_contains_cute_ghost(ctx, typed.get_type(ctx)) + { + return Err(escape(format!("block attribute `{name}` below `{opid}`"))); + } + } + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs new file mode 100644 index 0000000000..ef3389c4cf --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -0,0 +1,782 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared tensor placement, register carriers, and tiled matrix-instruction contracts. + +use std::collections::HashSet; + +use dialect_mir::ops::{ + MirCastOp, MirConstructStructOp, MirConstructTupleOp, MirExtractFieldOp, MirFieldAddrOp, + MirInsertFieldOp, MirLoadOp, MirStoreOp, control_flow::MirCondBranchOp, +}; +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + value::Value, +}; + +use super::{ + Result, + ghosts::smem_mma_semantic_ids, + invalid, + provenance::{ + CarrierKey, OrdinaryFacts, branch_successor_operands, is_semantic_pointer_carrier_cast, + local_cell, local_cell_is_closed, poison_missing_forwarding_sources, + propagate_ordinary_plumbing, seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::CuteTiledMmaPlanAttr, + ops::epilogue::CuteEpilogueStoreFragmentOp, + ops::mma::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, + }, +}; + +/// Read the tiled-MMA plan from a recognized compute operation, or return `None`. +fn mma_plan_for_op( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Option { + if *opid == CuteTiledMmaSliceOp::get_opid_static() { + CuteTiledMmaSliceOp::wrap(operation).plan(ctx) + } else if *opid == CuteFragmentFillOp::get_opid_static() { + CuteFragmentFillOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaLoadScalesOp::get_opid_static() { + CuteMmaLoadScalesOp::wrap(operation).plan(ctx) + } else if *opid == CuteFragmentSliceKOp::get_opid_static() { + CuteFragmentSliceKOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaLoadAOp::get_opid_static() { + CuteMmaLoadAOp::wrap(operation).plan(ctx) + } else if *opid == CuteMmaPartitionBOp::get_opid_static() { + CuteMmaPartitionBOp::wrap(operation).plan(ctx) + } else if *opid == CuteTiledGemmOp::get_opid_static() { + CuteTiledGemmOp::wrap(operation).plan(ctx) + } else { + None + } +} + +/// One possible semantic origin of an ordinary MMA carrier or aggregate field. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub(super) enum MmaFact { + /// Lane selection produced by a tiled-MMA slice. + Lane(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Shared allocation base produced by a tensor overlay. + SmemBase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Element capacity paired with a tensor overlay's shared base. + SmemCapacity( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// A complete stage of register scales produced by a scale load. + StageScales( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// A K=64 scale selection produced from one stage. + SlicedScales( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// Packed A registers produced by a semantic shared load. + AFragment(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Shared B base preserved by lazy partitioning. + BBase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Element capacity paired with a lazy B partition. + BCapacity(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Warp-N coordinate of a lazy B partition. + BWarpN(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// K-half coordinate of a lazy B partition. + BKHalf(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// FP32 accumulator initialized or updated by the same compute chain. + Accumulator( + #[doc = "Operation identifying the semantic origin of this carrier."] Ptr, + ), + /// An ordinary scalar whose exact SSA identity must remain consistent. + Scalar(#[doc = "Original SSA value lacking semantic provenance."] Value), + /// An input without a proven semantic origin; merging it poisons provenance. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Return the sole carrier fact, rejecting absent or conflicting facts; callers validate its kind. +pub(super) fn one_mma_fact( + facts: &OrdinaryFacts, + value: Value, + what: &str, +) -> Result { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no shared-MMA provenance after ordinary carrier forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges carriers from different shared-MMA stories" + ))); + } + Ok(*found.iter().next().expect("one fact")) +} + +/// Require the lane carrier to originate from the expected tiled-MMA slice. +fn require_mma_lane( + facts: &OrdinaryFacts, + value: Value, + what: &str, +) -> Result> { + match one_mma_fact(facts, value, what)? { + MmaFact::Lane(root) => Ok(root), + _ => Err(invalid(format!( + "{what} must come from cute.tiled_mma_slice" + ))), + } +} + +/// Require the shared base and capacity to come from the same typed overlay. +pub(super) fn require_smem_overlay_pair( + facts: &OrdinaryFacts, + base: Value, + capacity: Value, + what: &str, +) -> Result> { + let base_root = match one_mma_fact(facts, base, &format!("{what} base"))? { + MmaFact::SmemBase(root) => root, + _ => { + return Err(invalid(format!( + "{what} base must come from cute.smem_tensor_overlay" + ))); + } + }; + let capacity_root = match one_mma_fact(facts, capacity, &format!("{what} capacity"))? { + MmaFact::SmemCapacity(root) => root, + _ => { + return Err(invalid(format!( + "{what} capacity must come from cute.smem_tensor_overlay" + ))); + } + }; + if base_root != capacity_root { + return Err(invalid(format!( + "{what} mixes base and capacity from different shared overlays" + ))); + } + Ok(base_root) +} + +/// Require two scalar carriers to preserve the same original SSA identity. +fn require_same_mma_scalar( + facts: &OrdinaryFacts, + left: Value, + right: Value, + what: &str, +) -> Result { + if left == right { + return Ok(()); + } + match ( + one_mma_fact(facts, left, what)?, + one_mma_fact(facts, right, what)?, + ) { + (MmaFact::Scalar(left), MmaFact::Scalar(right)) if left == right => Ok(()), + _ => Err(invalid(format!( + "{what} must be one identical ordinary SSA carrier" + ))), + } +} + +/// Propagate compute origins through semantic operations and ordinary carriers, then poison unknown merges. +pub(super) fn build_mma_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + let opid = Operation::get_opid(*operation, ctx); + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + let overlay = CuteSmemTensorOverlayOp::wrap(*operation); + facts.insert(overlay.base(ctx), MmaFact::SmemBase(*operation)); + facts.insert(overlay.capacity(ctx), MmaFact::SmemCapacity(*operation)); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + facts.insert( + CuteTiledMmaSliceOp::wrap(*operation).sliced_lane(ctx), + MmaFact::Lane(*operation), + ); + } else if opid == CuteFragmentFillOp::get_opid_static() { + facts.insert( + CuteFragmentFillOp::wrap(*operation).fragment(ctx), + MmaFact::Accumulator(*operation), + ); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + facts.insert( + CuteMmaLoadScalesOp::wrap(*operation).scales(ctx), + MmaFact::StageScales(*operation), + ); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + facts.insert( + CuteFragmentSliceKOp::wrap(*operation).scales(ctx), + MmaFact::SlicedScales(*operation), + ); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + facts.insert( + CuteMmaLoadAOp::wrap(*operation).fragment(ctx), + MmaFact::AFragment(*operation), + ); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + facts.insert(operation_ref.get_result(0), MmaFact::BBase(*operation)); + facts.insert(operation_ref.get_result(1), MmaFact::BCapacity(*operation)); + facts.insert(operation_ref.get_result(2), MmaFact::BWarpN(*operation)); + facts.insert(operation_ref.get_result(3), MmaFact::BKHalf(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteTiledGemmOp::get_opid_static() { + continue; + } + let operation_ref = operation.deref(ctx); + let Some(accumulator_facts) = facts.at(operation_ref.get_operand(7)).cloned() else { + continue; + }; + for accumulator in accumulator_facts { + if let MmaFact::Accumulator(root) = accumulator { + changed |= facts.insert( + CuteTiledGemmOp::wrap(*operation).accumulator(ctx), + MmaFact::Accumulator(root), + ); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, MmaFact::Scalar); + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, MmaFact::Scalar); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, MmaFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Look up an exact semantic origin inside a carrier's aggregate field path. +fn mma_fact_at_path( + facts: &OrdinaryFacts, + value: Value, + path: &[u32], + fact: MmaFact, +) -> bool { + facts + .at_path(value, path) + .is_some_and(|found| found.contains(&fact)) +} + +/// Check whether an operation consumes the fact in its designated semantic operand position. +fn mma_semantic_consumer( + ctx: &Context, + user: Ptr, + operand_index: usize, + fact: MmaFact, +) -> bool { + let opid = Operation::get_opid(user, ctx); + match fact { + MmaFact::Lane(_) => { + operand_index == 0 + && (opid == CuteMmaLoadScalesOp::get_opid_static() + || opid == CuteMmaLoadAOp::get_opid_static() + || opid == CuteTiledGemmOp::get_opid_static()) + } + MmaFact::SmemBase(_) => { + (opid == CuteMmaLoadScalesOp::get_opid_static() && matches!(operand_index, 1 | 3)) + || (opid == CuteMmaLoadAOp::get_opid_static() && operand_index == 1) + || (opid == CuteMmaPartitionBOp::get_opid_static() && operand_index == 0) + } + MmaFact::SmemCapacity(_) => { + (opid == CuteMmaLoadScalesOp::get_opid_static() && matches!(operand_index, 2 | 4)) + || (opid == CuteMmaLoadAOp::get_opid_static() && operand_index == 2) + || (opid == CuteMmaPartitionBOp::get_opid_static() && operand_index == 1) + } + MmaFact::StageScales(_) => { + opid == CuteFragmentSliceKOp::get_opid_static() && operand_index == 0 + } + MmaFact::SlicedScales(_) => { + opid == CuteTiledGemmOp::get_opid_static() && operand_index == 6 + } + MmaFact::AFragment(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 1, + MmaFact::BBase(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 2, + MmaFact::BCapacity(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 3, + MmaFact::BWarpN(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 4, + MmaFact::BKHalf(_) => opid == CuteTiledGemmOp::get_opid_static() && operand_index == 5, + MmaFact::Accumulator(_) => { + (opid == CuteTiledGemmOp::get_opid_static() && operand_index == 7) + || (opid == CuteEpilogueStoreFragmentOp::get_opid_static() && operand_index == 3) + } + MmaFact::Scalar(_) | MmaFact::Unknown(_) => true, + } +} + +/// Check that a branch forwards the carrier fact to the matching successor argument. +fn mma_branch_forwards_fact( + ctx: &Context, + facts: &OrdinaryFacts, + user: Ptr, + value: Value, + path: &[u32], + fact: MmaFact, +) -> bool { + let operation = user.deref(ctx); + (0..operation.get_num_successors()).any(|successor_index| { + let successor = operation.get_successor(successor_index); + branch_successor_operands(ctx, user, successor_index).is_some_and(|operands| { + operands + .iter() + .copied() + .enumerate() + .any(|(index, incoming)| { + incoming == value + && index < successor.deref(ctx).get_num_arguments() + && mma_fact_at_path( + facts, + successor.deref(ctx).get_argument(index), + path, + fact, + ) + }) + }) + }) +} + +/// Check that ordinary casts, aggregates, or closed local cells forward the tracked fact without escape. +fn mma_ordinary_use_forwards_fact( + ctx: &Context, + facts: &OrdinaryFacts, + key: &CarrierKey, + fact: MmaFact, + user: Ptr, + operand_index: usize, +) -> bool { + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if opid == MirCastOp::get_opid_static() + && operand_index == 0 + && is_semantic_pointer_carrier_cast(ctx, user) + && operation.get_num_results() == 1 + { + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation.get_num_results() == 1 + { + let mut destination_path = vec![operand_index as u32]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + if opid == MirInsertFieldOp::get_opid_static() + && operation.get_num_operands() == 2 + && operation.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(user).get_attr_insert_index(ctx) else { + return false; + }; + if operand_index == 0 { + if key.path.first().copied() == Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path, fact); + } + if operand_index == 1 { + let mut destination_path = vec![index.0]; + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, operation.get_result(0), &destination_path, fact); + } + return false; + } + if opid == MirExtractFieldOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(user).get_attr_index(ctx) else { + return false; + }; + if key.path.first().copied() != Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &key.path[1..], fact); + } + if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { + let store = MirStoreOp::new(user); + if operand_index == 0 { + return local_cell(ctx, store.address_opd(ctx)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if operand_index == 1 + && let Some((root, mut destination_path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + destination_path.extend_from_slice(&key.path); + return mma_fact_at_path(facts, root, &destination_path, fact); + } + return false; + } + if opid == MirLoadOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + { + let load = MirLoadOp::new(user); + let Some((root, cell_path)) = local_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + if !local_cell_is_closed(ctx, root) { + return false; + } + let Some(result_path) = key.path.strip_prefix(cell_path.as_slice()) else { + return true; + }; + return mma_fact_at_path(facts, operation.get_result(0), result_path, fact); + } + if opid == MirFieldAddrOp::get_opid_static() + && operand_index == 0 + && operation.get_num_results() == 1 + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some() + { + return local_cell(ctx, operation.get_result(0)) + .is_some_and(|(root, _)| local_cell_is_closed(ctx, root)); + } + if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { + return false; + } + mma_branch_forwards_fact(ctx, facts, user, key.value, &key.path, fact) +} + +/// Reject ordinary uses that discard or expose semantic MMA carriers outside recognized forwarding. +pub(super) fn verify_mma_use_closure( + ctx: &Context, + operations: &[Ptr], + facts: &OrdinaryFacts, +) -> Result { + for (key, found) in &facts.facts { + for fact in found.iter().copied() { + if matches!(fact, MmaFact::Scalar(_) | MmaFact::Unknown(_)) { + continue; + } + for r#use in key.value.uses(ctx) { + let user = r#use.user_op(); + let operand_index = r#use.find_index(ctx); + if operations.contains(&user) + && (mma_semantic_consumer(ctx, user, operand_index, fact) + || mma_ordinary_use_forwards_fact( + ctx, + facts, + key, + fact, + user, + operand_index, + )) + { + continue; + } + return Err(invalid(format!( + "shared-MMA carrier has an unsupported outgoing use by `{}` operand {operand_index}", + Operation::get_opid(user, ctx) + ))); + } + } + } + Ok(()) +} + +/// Require a compute operation's static plan to match the chain's tiled-MMA root. +fn require_same_mma_plan( + ctx: &Context, + left: Ptr, + right: Ptr, + what: &str, +) -> Result { + let left_id = Operation::get_opid(left, ctx); + let right_id = Operation::get_opid(right, ctx); + if mma_plan_for_op(ctx, left, &left_id) != mma_plan_for_op(ctx, right, &right_id) { + return Err(invalid(format!( + "{what} connects shared-MMA operations with different tiled-MMA plans" + ))); + } + Ok(()) +} + +/// Validate shared overlays, lazy B partitions, scale/A fragments, and accumulator provenance together. +pub(super) fn verify_smem_mma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = smem_mma_semantic_ids(); + let facts = build_mma_facts(ctx, all_ops); + let mut overlays = HashSet::new(); + let mut used_overlays = HashSet::new(); + let mut slices = HashSet::new(); + let mut used_slices = HashSet::new(); + let mut fills = HashSet::new(); + let mut used_fills = HashSet::new(); + let mut scale_loads = HashSet::new(); + let mut used_scale_loads = HashSet::new(); + let mut scale_slices = HashSet::new(); + let mut used_scale_slices = HashSet::new(); + let mut a_loads = HashSet::new(); + let mut used_a_loads = HashSet::new(); + let mut b_partitions = HashSet::new(); + let mut used_b_partitions = HashSet::new(); + let mut gemms = 0usize; + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "shared/MMA operation `{opid}` is not attached to a block" + ))); + } + let operation_ref = operation.deref(ctx); + if opid == CuteSmemTensorOverlayOp::get_opid_static() { + overlays.insert(*operation); + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + slices.insert(*operation); + } else if opid == CuteFragmentFillOp::get_opid_static() { + fills.insert(*operation); + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + scale_loads.insert(*operation); + let semantic = CuteMmaLoadScalesOp::wrap(*operation); + let lane = require_mma_lane( + &facts, + operation_ref.get_operand(0), + "cute.mma_load_scales lane", + )?; + require_same_mma_plan(ctx, lane, *operation, "mma scale load")?; + used_slices.insert(lane); + let scale_a = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_scales A", + )?; + let scale_b = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(3), + operation_ref.get_operand(4), + "cute.mma_load_scales B", + )?; + if semantic.scale_a_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_a).view(ctx) + || semantic.scale_b_view(ctx) != CuteSmemTensorOverlayOp::wrap(scale_b).view(ctx) + { + return Err(invalid( + "cute.mma_load_scales typed views do not match its exact shared overlays", + )); + } + used_overlays.extend([scale_a, scale_b]); + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + scale_slices.insert(*operation); + let load = match one_mma_fact( + &facts, + operation_ref.get_operand(0), + "cute.fragment_slice_k stage scales", + )? { + MmaFact::StageScales(load) => load, + _ => { + return Err(invalid( + "cute.fragment_slice_k must consume cute.mma_load_scales", + )); + } + }; + require_same_mma_plan(ctx, load, *operation, "scale K slice")?; + used_scale_loads.insert(load); + } else if opid == CuteMmaLoadAOp::get_opid_static() { + a_loads.insert(*operation); + let semantic = CuteMmaLoadAOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.mma_load_a lane")?; + require_same_mma_plan(ctx, lane, *operation, "MMA A load")?; + used_slices.insert(lane); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(1), + operation_ref.get_operand(2), + "cute.mma_load_a", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_load_a typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + b_partitions.insert(*operation); + let semantic = CuteMmaPartitionBOp::wrap(*operation); + let overlay = require_smem_overlay_pair( + &facts, + operation_ref.get_operand(0), + operation_ref.get_operand(1), + "cute.mma_partition_b", + )?; + if semantic.view(ctx) != CuteSmemTensorOverlayOp::wrap(overlay).view(ctx) { + return Err(invalid( + "cute.mma_partition_b typed view does not match its exact shared overlay", + )); + } + used_overlays.insert(overlay); + } else if opid == CuteTiledGemmOp::get_opid_static() { + gemms += 1; + let semantic = CuteTiledGemmOp::wrap(*operation); + let lane = + require_mma_lane(&facts, operation_ref.get_operand(0), "cute.tiled_gemm lane")?; + let a = match one_mma_fact( + &facts, + operation_ref.get_operand(1), + "cute.tiled_gemm A fragment", + )? { + MmaFact::AFragment(load) => load, + _ => return Err(invalid("cute.tiled_gemm A must come from cute.mma_load_a")), + }; + let b = match ( + one_mma_fact(&facts, operation_ref.get_operand(2), "tiled_gemm B base")?, + one_mma_fact( + &facts, + operation_ref.get_operand(3), + "tiled_gemm B capacity", + )?, + one_mma_fact(&facts, operation_ref.get_operand(4), "tiled_gemm B warp-N")?, + one_mma_fact(&facts, operation_ref.get_operand(5), "tiled_gemm B K-half")?, + ) { + ( + MmaFact::BBase(base), + MmaFact::BCapacity(capacity), + MmaFact::BWarpN(warp_n), + MmaFact::BKHalf(k_half), + ) if base == capacity && base == warp_n && base == k_half => base, + _ => { + return Err(invalid( + "cute.tiled_gemm B carriers must be the four matching results of one mma_partition_b", + )); + } + }; + let scales = match one_mma_fact( + &facts, + operation_ref.get_operand(6), + "cute.tiled_gemm scales", + )? { + MmaFact::SlicedScales(slice) => slice, + _ => { + return Err(invalid( + "cute.tiled_gemm scales must come from cute.fragment_slice_k", + )); + } + }; + let fill = match one_mma_fact( + &facts, + operation_ref.get_operand(7), + "cute.tiled_gemm accumulator", + )? { + MmaFact::Accumulator(fill) => fill, + _ => { + return Err(invalid( + "cute.tiled_gemm accumulator must be the closed recurrence seeded by fragment_fill", + )); + } + }; + for producer in [lane, a, b, scales, fill] { + require_same_mma_plan(ctx, producer, *operation, "tiled GEMM spine")?; + } + let scale_load = match one_mma_fact( + &facts, + scales.deref(ctx).get_operand(0), + "tiled GEMM scale-load origin", + )? { + MmaFact::StageScales(load) => load, + _ => unreachable!("fragment_slice_k was verified above"), + }; + let scale_load_ref = scale_load.deref(ctx); + let a_ref = a.deref(ctx); + let b_ref = b.deref(ctx); + let scale_lane = require_mma_lane( + &facts, + scale_load_ref.get_operand(0), + "mma_load_scales lane", + )?; + let a_lane = require_mma_lane(&facts, a_ref.get_operand(0), "mma_load_a lane")?; + if scale_lane != lane || a_lane != lane { + return Err(invalid( + "tiled GEMM, scale load, and A load must use one tiled-MMA lane slice", + )); + } + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(5), + a_ref.get_operand(3), + "shared-MMA warp-M", + )?; + require_same_mma_scalar( + &facts, + scale_load_ref.get_operand(6), + b_ref.get_operand(2), + "shared-MMA warp-N", + )?; + require_same_mma_scalar( + &facts, + scales.deref(ctx).get_operand(1), + a_ref.get_operand(4), + "shared-MMA K-half", + )?; + require_same_mma_scalar( + &facts, + a_ref.get_operand(4), + b_ref.get_operand(3), + "shared-MMA K-half", + )?; + if semantic.b_view(ctx) != CuteMmaPartitionBOp::wrap(b).view(ctx) { + return Err(invalid( + "cute.tiled_gemm B view does not match its mma_partition_b producer", + )); + } + used_slices.insert(lane); + used_a_loads.insert(a); + used_b_partitions.insert(b); + used_scale_slices.insert(scales); + used_fills.insert(fill); + } + } + + let found = overlays.len() + + slices.len() + + fills.len() + + scale_loads.len() + + scale_slices.len() + + a_loads.len() + + b_partitions.len() + + gemms; + if found == 0 { + return Ok(()); + } + verify_mma_use_closure(ctx, all_ops, &facts)?; + if gemms == 0 + || overlays != used_overlays + || slices != used_slices + || fills != used_fills + || scale_loads != used_scale_loads + || scale_slices != used_scale_slices + || a_loads != used_a_loads + || b_partitions != used_b_partitions + { + return Err(invalid( + "shared-MMA story is incomplete or has an unconsumed overlay/slice/load/partition/fill producer", + )); + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mod.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mod.rs new file mode 100644 index 0000000000..a438172afd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mod.rs @@ -0,0 +1,312 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Check CuTe values and operation order before either backend lowers them. +//! +//! This module follows where values come from and how they are used. It checks +//! pipeline ordering and restrictions on compiler-only types without changing +//! the input module. +//! +//! Verification runs local operation checks, the ghost escape audit, then +//! family checks within each MIR function (or the root for standalone IR). +//! The supported graph shapes are intentionally specific: +//! +//! - Each TMA load pipeline has one initialization and one complete producer +//! and consumer state. Each lifecycle event has one static operation site; +//! loop iterations reuse those sites. +//! - The SM120 epilogue has one overlay, warp slice, and fragment store, followed +//! by two output halves in one committed store group per function. Its writer +//! and issuer sequences may cross unconditional blocks with no side entry. +//! +//! These are limits of the current pattern checks, not proofs of arbitrary +//! concurrent programs. Dynamic bounds, participation, and barrier phases +//! remain source obligations. `cfg.rs` distinguishes dominance, reachability, +//! and straight-line order; `provenance.rs` follows ordinary carrier values. + +/// Function scopes, reachability, dominance, and source-order protocol checks. +mod cfg; +/// Shared output placement and asynchronous store synchronization contracts. +mod epilogue; +/// Packed-value and scale-view contracts for the block-scaled GEMV flow. +mod gemv; +/// Restrictions preventing compiler-only CuTe descriptions from entering runtime storage. +mod ghosts; +/// Shared tensor placement, register carriers, and tiled matrix-instruction contracts. +mod mma; +/// Barrier-ring configuration and producer/consumer cursor contracts. +mod pipeline; +mod provenance; +/// Persistent output-tile grids and their ordinary cursor provenance. +mod scheduler; +/// Tensor storage, access, format, and selection metadata. +mod tensor; +/// Descriptor-backed raw transport views and their direct copy contracts. +mod tma; + +/// Regression fixtures covering accepted contracts and rejected malformed IR. +#[cfg(test)] +mod tests; + +use cfg::operations_by_scope; +use epilogue::verify_epilogue_story; +use gemv::verify_gemv_story; +use ghosts::audit_cute_ghosts; +pub use ghosts::type_contains_cute_ghost; +use mma::verify_smem_mma_story; +use pipeline::verify_pipeline_story; +use pliron::{ + common_traits::Verify, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, +}; +use scheduler::verify_scheduler_story; +use tensor::verify_tensor_story; +use tma::verify_tma_story; + +use crate::{ + ops::epilogue::{ + CuteEpilogueHalfOp, CuteEpilogueSmemOverlayOp, CuteEpilogueStoreFragmentOp, + CuteEpilogueSyncOp, CuteEpilogueWarpSliceOp, CuteTmaStore2dSemanticOp, + CuteTmaStoreAcquireOp, CuteTmaStoreCommitOp, CuteTmaStoreTailOp, + }, + ops::gemm_tma::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + ops::gemv::{ + CuteDotOp, CuteScaledViewKTileOp, CuteScaledViewLoadOp, CuteScaledViewMakeOp, + CuteScaledViewRowOp, CuteTensorMake2DOp, + }, + ops::mma::{ + CuteFragmentFillOp, CuteFragmentSliceKOp, CuteMmaLoadAOp, CuteMmaLoadScalesOp, + CuteMmaPartitionBOp, CuteSmemTensorOverlayOp, CuteTiledGemmOp, CuteTiledMmaSliceOp, + }, + ops::pipeline::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, + }, + ops::scheduler::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + ops::tensor::{ + CuteTensorBaseOp, CuteTensorIsFullOp, CuteTensorLoadIntoOp, CuteTensorMakeOp, + CuteTensorSliceOp, CuteTensorStoreElementAbsOp, CuteTensorStoreFromOp, + CuteTensorZippedDivideOp, + }, + ops::{ + CuteAssumeDivOp, CuteCopyG2SOp, CuteCopyOp, CuteLdmatrixOp, CuteTmaLoad2dOp, + CuteTmaStore2dOp, + }, +}; + +/// Failure reported by the shared semantic verifier. +/// +/// Native expansion and MLIR translation can both wrap this type without +/// depending on one another's error or lowering representations. +#[derive(Debug, thiserror::Error)] +pub enum VerifyError { + /// The operation graph violates a supported semantic contract. + #[error("invalid CuTe semantic module: {0}")] + Invalid(#[doc = "Explanation of the violated semantic contract."] String), + /// The graph contains a CuTe operation outside the verifier's supported set. + #[error("unsupported CuTe semantic operation: {0}")] + Unsupported(#[doc = "Unsupported operation or feature description."] String), +} + +/// Result type shared by backend-neutral verification and backend adapters. +pub type VerifyResult = std::result::Result; + +/// Internal shorthand for the backend-neutral semantic verification result. +type Result = VerifyResult; + +/// Wrap a semantic contract failure in the shared invalid-module error variant. +fn invalid(message: impl Into) -> VerifyError { + VerifyError::Invalid(message.into()) +} + +/// Visit the root and all nested operations in region/block order without changing the IR. +fn collect_ops(ctx: &Context, root: Ptr, output: &mut Vec>) { + output.push(root); + let regions: Vec<_> = root.deref(ctx).regions().collect(); + for region in regions { + let blocks: Vec<_> = region.deref(ctx).iter(ctx).collect(); + for block in blocks { + let children: Vec<_> = block.deref(ctx).iter(ctx).collect(); + for child in children { + collect_ops(ctx, child, output); + } + } + } +} + +/// Dispatch each known CuTe operation's local verifier and reject unknown CuTe operation IDs. +fn verify_local_op(ctx: &Context, operation: Ptr) -> Result { + let opid = Operation::get_opid(operation, ctx); + /// Wrap the operation as its recognized concrete type and run local verification. + macro_rules! verify_as { + ($ty:ty) => { + <$ty>::wrap(operation).verify(ctx) + }; + } + + let result = if opid == crate::ops::copy::CuteStoreGlobalNoAllocateOp::get_opid_static() { + Some(verify_as!(crate::ops::copy::CuteStoreGlobalNoAllocateOp)) + } else if let Some(result) = crate::ops::sm100_attention::verify_local(ctx, operation) { + Some(result) + } else if let Some(result) = crate::ops::sm100::verify_local(ctx, operation) { + Some(result) + } else if opid == CuteTensorMakeOp::get_opid_static() { + Some(verify_as!(CuteTensorMakeOp)) + } else if opid == CuteTensorZippedDivideOp::get_opid_static() { + Some(verify_as!(CuteTensorZippedDivideOp)) + } else if opid == CuteTensorSliceOp::get_opid_static() { + Some(verify_as!(CuteTensorSliceOp)) + } else if opid == CuteTensorIsFullOp::get_opid_static() { + Some(verify_as!(CuteTensorIsFullOp)) + } else if opid == CuteTensorBaseOp::get_opid_static() { + Some(verify_as!(CuteTensorBaseOp)) + } else if opid == CuteTensorLoadIntoOp::get_opid_static() { + Some(verify_as!(CuteTensorLoadIntoOp)) + } else if opid == CuteTensorStoreFromOp::get_opid_static() { + Some(verify_as!(CuteTensorStoreFromOp)) + } else if opid == CuteTensorStoreElementAbsOp::get_opid_static() { + Some(verify_as!(CuteTensorStoreElementAbsOp)) + } else if opid == CuteTensorMake2DOp::get_opid_static() { + Some(verify_as!(CuteTensorMake2DOp)) + } else if opid == CuteScaledViewMakeOp::get_opid_static() { + Some(verify_as!(CuteScaledViewMakeOp)) + } else if opid == CuteScaledViewRowOp::get_opid_static() { + Some(verify_as!(CuteScaledViewRowOp)) + } else if opid == CuteScaledViewKTileOp::get_opid_static() { + Some(verify_as!(CuteScaledViewKTileOp)) + } else if opid == CuteScaledViewLoadOp::get_opid_static() { + Some(verify_as!(CuteScaledViewLoadOp)) + } else if opid == CuteDotOp::get_opid_static() { + Some(verify_as!(CuteDotOp)) + } else if opid == CuteSchedulerNew1dOp::get_opid_static() { + Some(verify_as!(CuteSchedulerNew1dOp)) + } else if opid == CuteSchedulerHasWorkOp::get_opid_static() { + Some(verify_as!(CuteSchedulerHasWorkOp)) + } else if opid == CuteSchedulerCurrentOp::get_opid_static() { + Some(verify_as!(CuteSchedulerCurrentOp)) + } else if opid == CuteWorkTileCoordinatesOp::get_opid_static() { + Some(verify_as!(CuteWorkTileCoordinatesOp)) + } else if opid == CuteSchedulerAdvanceOp::get_opid_static() { + Some(verify_as!(CuteSchedulerAdvanceOp)) + } else if opid == CuteTmaLoadPipelineMakeOp::get_opid_static() { + Some(verify_as!(CuteTmaLoadPipelineMakeOp)) + } else if opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + Some(verify_as!(CuteTmaLoadPipelineInitOp)) + } else if opid == CutePipelineStateNewOp::get_opid_static() { + Some(verify_as!(CutePipelineStateNewOp)) + } else if opid == CutePipelineStateSlotOp::get_opid_static() { + Some(verify_as!(CutePipelineStateSlotOp)) + } else if opid == CutePipelineStateAdvanceOp::get_opid_static() { + Some(verify_as!(CutePipelineStateAdvanceOp)) + } else if opid == CutePipelineProducerAcquireOp::get_opid_static() { + Some(verify_as!(CutePipelineProducerAcquireOp)) + } else if opid == CutePipelineProducerExpectTxOp::get_opid_static() { + Some(verify_as!(CutePipelineProducerExpectTxOp)) + } else if opid == CutePipelineConsumerWaitOp::get_opid_static() { + Some(verify_as!(CutePipelineConsumerWaitOp)) + } else if opid == CutePipelineConsumerReleaseOp::get_opid_static() { + Some(verify_as!(CutePipelineConsumerReleaseOp)) + } else if opid == CutePipelineProducerTailOp::get_opid_static() { + Some(verify_as!(CutePipelineProducerTailOp)) + } else if opid == CuteTmaGmemViewOp::get_opid_static() { + Some(verify_as!(CuteTmaGmemViewOp)) + } else if opid == CuteTmaSmemViewOp::get_opid_static() { + Some(verify_as!(CuteTmaSmemViewOp)) + } else if opid == CuteTmaCopy2dOp::get_opid_static() { + Some(verify_as!(CuteTmaCopy2dOp)) + } else if opid == CuteSmemTensorOverlayOp::get_opid_static() { + Some(verify_as!(CuteSmemTensorOverlayOp)) + } else if opid == CuteTiledMmaSliceOp::get_opid_static() { + Some(verify_as!(CuteTiledMmaSliceOp)) + } else if opid == CuteFragmentFillOp::get_opid_static() { + Some(verify_as!(CuteFragmentFillOp)) + } else if opid == CuteMmaLoadScalesOp::get_opid_static() { + Some(verify_as!(CuteMmaLoadScalesOp)) + } else if opid == CuteFragmentSliceKOp::get_opid_static() { + Some(verify_as!(CuteFragmentSliceKOp)) + } else if opid == CuteMmaLoadAOp::get_opid_static() { + Some(verify_as!(CuteMmaLoadAOp)) + } else if opid == CuteMmaPartitionBOp::get_opid_static() { + Some(verify_as!(CuteMmaPartitionBOp)) + } else if opid == CuteTiledGemmOp::get_opid_static() { + Some(verify_as!(CuteTiledGemmOp)) + } else if opid == CuteEpilogueSmemOverlayOp::get_opid_static() { + Some(verify_as!(CuteEpilogueSmemOverlayOp)) + } else if opid == CuteEpilogueWarpSliceOp::get_opid_static() { + Some(verify_as!(CuteEpilogueWarpSliceOp)) + } else if opid == CuteEpilogueStoreFragmentOp::get_opid_static() { + Some(verify_as!(CuteEpilogueStoreFragmentOp)) + } else if opid == CuteEpilogueSyncOp::get_opid_static() { + Some(verify_as!(CuteEpilogueSyncOp)) + } else if opid == CuteEpilogueHalfOp::get_opid_static() { + Some(verify_as!(CuteEpilogueHalfOp)) + } else if opid == CuteTmaStoreAcquireOp::get_opid_static() { + Some(verify_as!(CuteTmaStoreAcquireOp)) + } else if opid == CuteTmaStoreCommitOp::get_opid_static() { + Some(verify_as!(CuteTmaStoreCommitOp)) + } else if opid == CuteTmaStoreTailOp::get_opid_static() { + Some(verify_as!(CuteTmaStoreTailOp)) + } else if opid == CuteTmaStore2dSemanticOp::get_opid_static() { + Some(verify_as!(CuteTmaStore2dSemanticOp)) + } else if opid == CuteCopyOp::get_opid_static() { + Some(verify_as!(CuteCopyOp)) + } else if opid == CuteCopyG2SOp::get_opid_static() { + Some(verify_as!(CuteCopyG2SOp)) + } else if opid == CuteAssumeDivOp::get_opid_static() { + Some(verify_as!(CuteAssumeDivOp)) + } else if opid == CuteLdmatrixOp::get_opid_static() { + Some(verify_as!(CuteLdmatrixOp)) + } else if opid == CuteTmaLoad2dOp::get_opid_static() { + Some(verify_as!(CuteTmaLoad2dOp)) + } else if opid == CuteTmaStore2dOp::get_opid_static() { + Some(verify_as!(CuteTmaStore2dOp)) + } else { + None + }; + + if let Some(result) = result { + result.map_err(|error| invalid(format!("`{opid}` failed local verification: {error}")))?; + } else if opid.dialect.to_string() == crate::CUTE_DIALECT_NAME { + return Err(VerifyError::Unsupported(format!( + "unhandled operation `{opid}`" + ))); + } + Ok(()) +} + +/// Validate the complete high-level CuTe module without cloning or mutating +/// the IR, allocating replacement operations, or invoking either backend. +/// +/// Unsafe runtime promises such as bounds, thread participation, and TMA +/// descriptor validity remain caller contracts. This verifier checks the +/// static semantic graph and rejects any compiler-only value that escapes it. +pub fn verify_cute_semantics(ctx: &Context, module: Ptr) -> Result { + let mut all_ops = Vec::new(); + collect_ops(ctx, module, &mut all_ops); + + for operation in &all_ops { + verify_local_op(ctx, *operation)?; + } + audit_cute_ghosts(ctx, &all_ops)?; + for scope in operations_by_scope(ctx, module, &all_ops) { + verify_scheduler_story(ctx, &scope)?; + verify_pipeline_story(ctx, &scope)?; + verify_tma_story(ctx, &scope)?; + verify_smem_mma_story(ctx, &scope)?; + verify_gemv_story(ctx, &scope)?; + verify_tensor_story(ctx, &scope)?; + verify_epilogue_story(ctx, &scope)?; + crate::ops::sm100::verify_story(ctx, &scope).map_err(invalid)?; + crate::ops::sm100_attention::verify_story(ctx, &scope).map_err(invalid)?; + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs new file mode 100644 index 0000000000..355f94342e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs @@ -0,0 +1,607 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Barrier-ring configuration and producer/consumer cursor contracts. + +use std::collections::{HashMap, HashSet}; + +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{operation_dominates, operation_reaches}, + ghosts::pipeline_semantic_ids, + invalid, + provenance::{ + OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::{CutePipelineRoleAttr, CutePipelineStateAttr}, + ops::gemm_tma::CuteTmaCopy2dOp, + ops::pipeline::{ + CutePipelineConsumerReleaseOp, CutePipelineConsumerWaitOp, CutePipelineProducerAcquireOp, + CutePipelineProducerExpectTxOp, CutePipelineProducerTailOp, CutePipelineStateAdvanceOp, + CutePipelineStateNewOp, CutePipelineStateSlotOp, CuteTmaLoadPipelineInitOp, + CuteTmaLoadPipelineMakeOp, + }, + types::CuteTmaViewType, +}; + +/// Static ring facts and the constructor identifying one load pipeline. +#[derive(Clone, Copy)] +struct LoadPipelineFacts { + /// Direct constructor that defines the compiler-only pipeline handle. + make: Ptr, + /// Number of full and empty barrier slots in each ring. + stages: u64, + /// Total bytes expected by one full barrier for its stage. + transaction_bytes: u32, +} + +/// Require a direct pipeline constructor and recover its ring configuration. +fn resolve_load_pipeline(ctx: &Context, value: Value) -> Result { + let Some(make) = value.defining_op() else { + return Err(invalid( + "TMA load-pipeline handle reached a block argument; a direct tma_load_pipeline_make result is required", + )); + }; + let producer_id = Operation::get_opid(make, ctx); + if producer_id != CuteTmaLoadPipelineMakeOp::get_opid_static() { + return Err(invalid(format!( + "TMA load-pipeline handle is produced by `{producer_id}` instead of `cute.tma_load_pipeline_make`" + ))); + } + let make_op = CuteTmaLoadPipelineMakeOp::wrap(make); + if make_op.pipeline(ctx) != value { + return Err(invalid( + "TMA load-pipeline operand is not the make operation's pipeline result", + )); + } + let pipeline = make_op + .pipeline_type(ctx) + .ok_or_else(|| invalid("TMA load-pipeline make has no pipeline type"))?; + Ok(LoadPipelineFacts { + make, + stages: pipeline.stages, + transaction_bytes: pipeline.transaction_bytes, + }) +} + +/// Read the producer/consumer role and ring size from a recognized cursor operation. +fn pipeline_state_for_op( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Option { + if *opid == CutePipelineStateNewOp::get_opid_static() { + CutePipelineStateNewOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineStateSlotOp::get_opid_static() { + CutePipelineStateSlotOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineStateAdvanceOp::get_opid_static() { + CutePipelineStateAdvanceOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerAcquireOp::get_opid_static() { + CutePipelineProducerAcquireOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerExpectTxOp::get_opid_static() { + CutePipelineProducerExpectTxOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineConsumerWaitOp::get_opid_static() { + CutePipelineConsumerWaitOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineConsumerReleaseOp::get_opid_static() { + CutePipelineConsumerReleaseOp::wrap(operation).state(ctx) + } else if *opid == CutePipelineProducerTailOp::get_opid_static() { + CutePipelineProducerTailOp::wrap(operation).state(ctx) + } else { + None + } +} + +/// Find a lifecycle operation's pipeline handle; cursor-only operations return `None`. +fn pipeline_operand(ctx: &Context, operation: Ptr, opid: &OpId) -> Option { + if *opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + Some(CuteTmaLoadPipelineInitOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerAcquireOp::get_opid_static() { + Some(CutePipelineProducerAcquireOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerExpectTxOp::get_opid_static() { + Some(CutePipelineProducerExpectTxOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineConsumerWaitOp::get_opid_static() { + Some(CutePipelineConsumerWaitOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineConsumerReleaseOp::get_opid_static() { + Some(CutePipelineConsumerReleaseOp::wrap(operation).pipeline(ctx)) + } else if *opid == CutePipelineProducerTailOp::get_opid_static() { + Some(CutePipelineProducerTailOp::wrap(operation).pipeline(ctx)) + } else { + None + } +} + +/// Sum the static byte sizes of every direct TMA copy completed by one +/// `expect_tx` result. The result may not escape through any other operand. +fn expected_tma_transaction_bytes(ctx: &Context, operation: Ptr) -> Result { + let barrier = CutePipelineProducerExpectTxOp::wrap(operation).completion_barrier(ctx); + let uses = barrier.uses(ctx); + if uses.is_empty() { + return Err(invalid( + "cute.pipeline_producer_expect_tx needs at least one direct cute.tma_copy_2d user", + )); + } + let mut copies = HashSet::new(); + let mut bytes = 0u64; + for r#use in uses { + let user = r#use.user_op(); + if Operation::get_opid(user, ctx) != CuteTmaCopy2dOp::get_opid_static() + || r#use.find_index(ctx) != 4 + || CuteTmaCopy2dOp::wrap(user).completion_barrier(ctx) != barrier + { + return Err(invalid( + "cute.pipeline_producer_expect_tx result may only be the completion-barrier operand of cute.tma_copy_2d", + )); + } + if !copies.insert(user) { + continue; + } + let copy = CuteTmaCopy2dOp::wrap(user); + let source_ty = copy.source(ctx).get_type(ctx); + let source_ty = source_ty.deref(ctx); + let view = source_ty + .downcast_ref::() + .ok_or_else(|| invalid("TMA copy attached to expect_tx has no TMA view type"))?; + let tile_bytes = view.tile_bytes(ctx).ok_or_else(|| { + invalid("TMA copy attached to expect_tx has no static tile byte count") + })?; + bytes = bytes + .checked_add(tile_bytes) + .ok_or_else(|| invalid("TMA transaction byte total overflows u64"))?; + } + Ok(bytes) +} + +/// Origin of one ordinary cursor component, including untrusted incoming values. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum PipelineFact { + /// Ring slot derived from this cursor-creation operation. + Slot(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Barrier parity derived from this cursor-creation operation. + Phase(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Scalar without a proven cursor origin; it cannot satisfy lifecycle checks. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Require an unambiguous slot or phase fact rooted at cursor creation. +fn pipeline_state_root( + facts: &OrdinaryFacts, + value: Value, + slot: bool, + what: &str, +) -> Result> { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no pipeline-state provenance after ordinary SSA forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges values from different pipeline-state stories" + ))); + } + match found.iter().next().copied() { + Some(PipelineFact::Slot(root)) if slot => Ok(root), + Some(PipelineFact::Phase(root)) if !slot => Ok(root), + _ => Err(invalid(format!( + "{what} is not derived from the required cute.pipeline_state_new result" + ))), + } +} + +/// Track slot/phase origins through advances and forwarding, preserving unknown inputs for consumer checks. +fn build_pipeline_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CutePipelineStateNewOp::get_opid_static() { + let state = CutePipelineStateNewOp::wrap(*operation); + facts.insert(state.slot(ctx), PipelineFact::Slot(*operation)); + facts.insert(state.phase(ctx), PipelineFact::Phase(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CutePipelineStateAdvanceOp::get_opid_static() + { + continue; + } + let advance = CutePipelineStateAdvanceOp::wrap(*operation); + let Some(slot_facts) = facts.at(advance.slot(ctx)).cloned() else { + continue; + }; + let Some(phase_facts) = facts.at(advance.phase(ctx)).cloned() else { + continue; + }; + for slot in &slot_facts { + let PipelineFact::Slot(root) = slot else { + continue; + }; + if phase_facts.contains(&PipelineFact::Phase(*root)) + && CutePipelineStateNewOp::wrap(*root).state(ctx) == advance.state(ctx) + { + changed |= facts.insert(advance.next_slot(ctx), PipelineFact::Slot(*root)); + changed |= facts.insert(advance.next_phase(ctx), PipelineFact::Phase(*root)); + } + } + } + if !changed { + break; + } + } + + seed_unknown_integer_carriers(ctx, operations, &mut facts, 32, PipelineFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, PipelineFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Static lifecycle sites collected for one producer or consumer cursor root. +#[derive(Default)] +struct PipelineStateEvents { + /// Sites that widen the cursor slot for addressing. + slots: Vec>, + /// Sites that advance this cursor's slot and phase. + advances: Vec>, + /// Producer sites waiting for an empty stage. + acquires: Vec>, + /// Producer sites announcing the stage's transaction bytes. + expects: Vec>, + /// Consumer sites waiting for a full stage. + waits: Vec>, + /// Consumer sites releasing the stage after reading. + releases: Vec>, + /// Producer exit sites draining the empty-barrier ring. + tails: Vec>, +} + +/// Require dominance from cursor creation through acquire/expect/advance and a reachable exit tail. +pub(super) fn require_producer_pipeline_order( + ctx: &Context, + root: Ptr, + acquire: Ptr, + expect: Ptr, + advance: Ptr, + tail: Ptr, +) -> Result { + if operation_dominates(ctx, root, acquire) + && operation_dominates(ctx, acquire, expect) + && operation_dominates(ctx, expect, advance) + && operation_reaches(ctx, advance, tail) + { + Ok(()) + } else { + Err(invalid( + "producer pipeline CFG must order state_new -> acquire -> expect_tx -> advance, with tail reachable only after the producer loop", + )) + } +} + +/// Validate handle ownership, complete cursor lifecycles, CFG order, and expected TMA byte totals. +pub(super) fn verify_pipeline_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = pipeline_semantic_ids(); + let facts = build_pipeline_facts(ctx, all_ops); + let makes: Vec<_> = all_ops + .iter() + .copied() + .filter(|operation| { + Operation::get_opid(*operation, ctx) == CuteTmaLoadPipelineMakeOp::get_opid_static() + }) + .collect(); + let mut init_by_make: HashMap, Vec>> = HashMap::new(); + let mut events: HashMap, PipelineStateEvents> = HashMap::new(); + let mut binding: HashMap, Ptr> = HashMap::new(); + + for make in &makes { + if make.deref(ctx).get_parent_block().is_none() { + return Err(invalid("TMA load-pipeline make is not attached to a block")); + } + let handle = CuteTmaLoadPipelineMakeOp::wrap(*make).pipeline(ctx); + let uses = handle.uses(ctx); + if uses.is_empty() { + return Err(invalid( + "cute.tma_load_pipeline_make has no lifecycle consumer", + )); + } + for r#use in uses { + let user = r#use.user_op(); + let user_id = Operation::get_opid(user, ctx); + if !ids.contains(&user_id) + || r#use.find_index(ctx) != 0 + || pipeline_operand(ctx, user, &user_id) != Some(handle) + { + return Err(invalid(format!( + "TMA load-pipeline handle has an invalid use by `{user_id}`" + ))); + } + } + } + + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "pipeline operation `{opid}` is not attached to a block" + ))); + } + if opid == CutePipelineStateNewOp::get_opid_static() { + events.entry(*operation).or_default(); + continue; + } + + let state_root = if opid == CutePipelineStateSlotOp::get_opid_static() { + let semantic = CutePipelineStateSlotOp::wrap(*operation); + let root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_slot slot", + )?; + if semantic.state(ctx) != CutePipelineStateNewOp::wrap(root).state(ctx) { + return Err(invalid( + "cute.pipeline_state_slot facts differ from its pipeline_state_new root", + )); + } + events.entry(root).or_default().slots.push(*operation); + Some(root) + } else if opid == CutePipelineStateAdvanceOp::get_opid_static() { + let semantic = CutePipelineStateAdvanceOp::wrap(*operation); + let slot_root = pipeline_state_root( + &facts, + semantic.slot(ctx), + true, + "cute.pipeline_state_advance slot", + )?; + let phase_root = pipeline_state_root( + &facts, + semantic.phase(ctx), + false, + "cute.pipeline_state_advance phase", + )?; + if slot_root != phase_root + || semantic.state(ctx) != CutePipelineStateNewOp::wrap(slot_root).state(ctx) + { + return Err(invalid( + "cute.pipeline_state_advance must consume slot and phase from one matching pipeline_state_new", + )); + } + events + .entry(slot_root) + .or_default() + .advances + .push(*operation); + Some(slot_root) + } else if opid == CuteTmaLoadPipelineMakeOp::get_opid_static() { + None + } else if opid == CuteTmaLoadPipelineInitOp::get_opid_static() { + let pipeline = resolve_load_pipeline( + ctx, + CuteTmaLoadPipelineInitOp::wrap(*operation).pipeline(ctx), + )?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline init refers to a make outside its semantic scope", + )); + } + init_by_make + .entry(pipeline.make) + .or_default() + .push(*operation); + None + } else { + let state = pipeline_state_for_op(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("pipeline operation `{opid}` has no state")))?; + let operation_ref = operation.deref(ctx); + let slot_root = pipeline_state_root( + &facts, + operation_ref.get_operand(1), + true, + &format!("`{opid}` slot"), + )?; + let needs_phase = opid == CutePipelineProducerAcquireOp::get_opid_static() + || opid == CutePipelineConsumerWaitOp::get_opid_static() + || opid == CutePipelineProducerTailOp::get_opid_static(); + if needs_phase { + let phase_root = pipeline_state_root( + &facts, + operation_ref.get_operand(2), + false, + &format!("`{opid}` phase"), + )?; + if phase_root != slot_root { + return Err(invalid(format!( + "`{opid}` mixes slot and phase from different pipeline states" + ))); + } + } + if CutePipelineStateNewOp::wrap(slot_root).state(ctx) != Some(state) { + return Err(invalid(format!( + "`{opid}` state facts differ from its pipeline_state_new root" + ))); + } + + let handle = pipeline_operand(ctx, *operation, &opid) + .ok_or_else(|| invalid(format!("`{opid}` has no load-pipeline handle")))?; + let pipeline = resolve_load_pipeline(ctx, handle)?; + if !makes.contains(&pipeline.make) { + return Err(invalid( + "pipeline lifecycle operation refers to a make outside its semantic scope", + )); + } + if state.stages != pipeline.stages { + return Err(invalid(format!( + "pipeline state on `{opid}` has {} stages but its handle has {}", + state.stages, pipeline.stages + ))); + } + if let Some(previous) = binding.insert(slot_root, pipeline.make) + && previous != pipeline.make + { + return Err(invalid( + "one ordinary pipeline state is used with multiple load-pipeline handles", + )); + } + + let root_events = events.entry(slot_root).or_default(); + if opid == CutePipelineProducerAcquireOp::get_opid_static() { + root_events.acquires.push(*operation); + } else if opid == CutePipelineProducerExpectTxOp::get_opid_static() { + root_events.expects.push(*operation); + let copied = expected_tma_transaction_bytes(ctx, *operation)?; + if copied != u64::from(pipeline.transaction_bytes) { + return Err(invalid(format!( + "cute.pipeline_producer_expect_tx promises {} bytes, but its direct TMA copies move {copied} bytes", + pipeline.transaction_bytes + ))); + } + } else if opid == CutePipelineConsumerWaitOp::get_opid_static() { + root_events.waits.push(*operation); + } else if opid == CutePipelineConsumerReleaseOp::get_opid_static() { + root_events.releases.push(*operation); + } else if opid == CutePipelineProducerTailOp::get_opid_static() { + root_events.tails.push(*operation); + } + Some(slot_root) + }; + + if let Some(root) = state_root { + let root_state = CutePipelineStateNewOp::wrap(root) + .state(ctx) + .expect("locally verified pipeline state root"); + if let Some(state) = pipeline_state_for_op(ctx, *operation, &opid) + && state != root_state + { + return Err(invalid(format!( + "`{opid}` does not carry its pipeline-state root's role and stage count" + ))); + } + } + } + + let mut roles_by_make: HashMap, (usize, usize)> = HashMap::new(); + for (root, root_events) in &events { + let state = CutePipelineStateNewOp::wrap(*root) + .state(ctx) + .expect("locally verified pipeline state root"); + let Some(make) = binding.get(root).copied() else { + return Err(invalid( + "cute.pipeline_state_new is not connected to any load-pipeline lifecycle", + )); + }; + if root_events.slots.is_empty() || root_events.advances.is_empty() { + return Err(invalid( + "each pipeline state needs state_slot and state_advance operations", + )); + } + let roles = roles_by_make.entry(make).or_default(); + match state.role { + CutePipelineRoleAttr::Producer => { + roles.0 += 1; + if root_events.acquires.len() != 1 + || root_events.expects.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || root_events.tails.len() != 1 + || !root_events.waits.is_empty() + || !root_events.releases.is_empty() + { + return Err(invalid( + "producer pipeline state needs acquire, expect_tx, advance, and exactly one tail, with no consumer events", + )); + } + let acquire = root_events.acquires[0]; + let expect = root_events.expects[0]; + let advance = root_events.advances[0]; + let tail = root_events.tails[0]; + require_producer_pipeline_order(ctx, *root, acquire, expect, advance, tail)?; + let barrier = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(ctx); + for r#use in barrier.uses(ctx) { + let copy = r#use.user_op(); + if !operation_dominates(ctx, expect, copy) + || !operation_dominates(ctx, copy, advance) + { + return Err(invalid( + "every TMA copy must be CFG-ordered after expect_tx and before producer state advance", + )); + } + } + } + CutePipelineRoleAttr::Consumer => { + roles.1 += 1; + if root_events.waits.len() != 1 + || root_events.releases.len() != 1 + || root_events.advances.len() != 1 + || root_events.slots.len() != 1 + || !root_events.acquires.is_empty() + || !root_events.expects.is_empty() + || !root_events.tails.is_empty() + { + return Err(invalid( + "consumer pipeline state needs wait, release, and advance, with no producer events", + )); + } + let wait = root_events.waits[0]; + let release = root_events.releases[0]; + let advance = root_events.advances[0]; + if !operation_dominates(ctx, *root, wait) + || !operation_dominates(ctx, wait, release) + || !operation_dominates(ctx, release, advance) + || !operation_reaches(ctx, advance, wait) + { + return Err(invalid( + "consumer pipeline CFG must order state_new -> wait -> release -> advance -> next wait", + )); + } + } + } + } + for make in makes { + if init_by_make + .get(&make) + .is_none_or(|initializers| initializers.len() != 1) + { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one initialization", + )); + } + if roles_by_make.get(&make).copied() != Some((1, 1)) { + return Err(invalid( + "each tma_load_pipeline_make needs exactly one complete producer state and one complete consumer state", + )); + } + let initializer = init_by_make[&make][0]; + for (root, bound_make) in &binding { + if *bound_make == make && !operation_dominates(ctx, initializer, *root) { + return Err(invalid( + "load-pipeline initialization must dominate both producer and consumer state creation", + )); + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs new file mode 100644 index 0000000000..7bd4ae9a83 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs @@ -0,0 +1,533 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Follow semantic origins through ordinary MIR carriers. +//! +//! Facts are sets of possible origins, including an explicit unknown origin. +//! Family analyses first propagate known facts to a fixed point, allowing loop +//! backedges to inherit their entry origin. They then seed unrelated values, +//! poison missing forwarding inputs, and propagate again before requiring one +//! unambiguous origin. Seeding unknowns earlier would reject valid loop state. +//! +//! A merge of an entry state and its own advanced state retains one origin; +//! merging either with an arbitrary scalar adds `Unknown` and is rejected. +//! Closed local cells conservatively union all stores, irrespective of order. +//! This tracks identity, not dynamic values or general memory aliasing. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAllocaOp, MirCastOp, MirConstructStructOp, MirConstructTupleOp, MirExtractFieldOp, + MirFieldAddrOp, MirInsertFieldOp, MirLoadOp, MirStoreOp, + }, + types::{MirArrayType, MirPtrType, MirStructType, MirTupleType, address_space}, +}; +use pliron::{ + builtin::{op_interfaces::BranchOpInterface, types::IntegerType}, + context::{Context, Ptr}, + op::{Op, op_cast}, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +/// Identity of a whole SSA carrier or one nested aggregate field. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(super) struct CarrierKey { + /// The whole SSA carrier before following any aggregate field path. + pub(super) value: Value, + /// Empty means the value itself. Non-empty paths describe a field inside + /// an ordinary aggregate before a matching `mir.extract_field` exposes it. + pub(super) path: Vec, +} + +/// Monotone sets of possible semantic origins, indexed by carrier and field path. +#[derive(Clone)] +pub(super) struct OrdinaryFacts { + /// Accumulated origins; multiple entries at one key represent a merge that consumers must resolve. + pub(super) facts: HashMap>, +} + +impl Default for OrdinaryFacts { + fn default() -> Self { + Self { + facts: HashMap::new(), + } + } +} + +impl OrdinaryFacts +where + F: Clone + Eq + std::hash::Hash, +{ + /// Add one origin to a field path and report whether the fixed-point state changed. + pub(super) fn insert_path(&mut self, value: Value, path: Vec, fact: F) -> bool { + self.facts + .entry(CarrierKey { value, path }) + .or_default() + .insert(fact) + } + + /// Add an origin for the whole SSA carrier, with an empty field path. + pub(super) fn insert(&mut self, value: Value, fact: F) -> bool { + self.insert_path(value, Vec::new(), fact) + } + + /// Borrow the possible origins at an exact carrier field path. + pub(super) fn at_path(&self, value: Value, path: &[u32]) -> Option<&HashSet> { + self.facts.get(&CarrierKey { + value, + path: path.to_vec(), + }) + } + + /// Borrow origins attached directly to a whole SSA carrier. + pub(super) fn at(&self, value: Value) -> Option<&HashSet> { + self.at_path(value, &[]) + } + + /// Collect all nested field origins carried by one SSA value. + pub(super) fn entries(&self, value: Value) -> Vec<(Vec, F)> { + self.facts + .iter() + .filter(|(key, _)| key.value == value) + .flat_map(|(key, facts)| facts.iter().cloned().map(|fact| (key.path.clone(), fact))) + .collect() + } + + /// Forward origins below a source field prefix into the corresponding destination prefix. + pub(super) fn copy_entries( + &mut self, + source: Value, + destination: Value, + source_prefix: &[u32], + destination_prefix: &[u32], + ) -> bool { + let mut changed = false; + for (path, fact) in self.entries(source) { + let Some(suffix) = path.strip_prefix(source_prefix) else { + continue; + }; + let mut destination_path = destination_prefix.to_vec(); + destination_path.extend_from_slice(suffix); + changed |= self.insert_path(destination, destination_path, fact); + } + changed + } +} + +/// Read values passed to one successor through Pliron's branch interface. +pub(super) fn branch_successor_operands( + ctx: &Context, + terminator: Ptr, + successor_index: usize, +) -> Option> { + let dynamic = Operation::get_op_dyn(terminator, ctx); + let branch = op_cast::(dynamic.as_ref())?; + Some(branch.successor_operands(ctx, successor_index)) +} + +/// A provenance-preserving cast keeps the pointee and never strengthens +/// mutability. Local-cell aliases must also keep their address space; semantic +/// shared-memory carriers may use the compiler's exact generic/shared +/// round-trip, but no other address-space transition is accepted. +fn is_provenance_preserving_pointer_cast( + ctx: &Context, + operation: Ptr, + allow_shared_generic_transition: bool, +) -> bool { + if Operation::get_opid(operation, ctx) != MirCastOp::get_opid_static() + || operation.deref(ctx).get_num_operands() != 1 + || operation.deref(ctx).get_num_results() != 1 + || !MirCastOp::new(operation) + .get_attr_cast_kind(ctx) + .is_some_and(|kind| matches!(*kind, MirCastKindAttr::PtrToPtr)) + { + return false; + } + let operation_ref = operation.deref(ctx); + let source_type = operation_ref.get_operand(0).get_type(ctx); + let result_type = operation_ref.get_result(0).get_type(ctx); + let source_type_ref = source_type.deref(ctx); + let result_type_ref = result_type.deref(ctx); + let (Some(source), Some(result)) = ( + source_type_ref.downcast_ref::(), + result_type_ref.downcast_ref::(), + ) else { + return false; + }; + let address_space_is_preserved = source.address_space == result.address_space + || (allow_shared_generic_transition + && matches!( + (source.address_space, result.address_space), + (address_space::GENERIC, address_space::SHARED) + | (address_space::SHARED, address_space::GENERIC) + )); + source.pointee == result.pointee + && address_space_is_preserved + && (source.is_mutable || !result.is_mutable) +} + +/// Accept pointer casts preserving pointee, address space, and non-strengthened mutability. +pub(super) fn is_local_cell_pointer_alias_cast(ctx: &Context, operation: Ptr) -> bool { + is_provenance_preserving_pointer_cast(ctx, operation, false) +} + +/// Also permit exact generic/shared pointer round-trips for semantic shared carriers. +pub(super) fn is_semantic_pointer_carrier_cast(ctx: &Context, operation: Ptr) -> bool { + is_provenance_preserving_pointer_cast(ctx, operation, true) +} + +/// Resolve a pointer to one compiler-created local cell and an exact constant +/// field path. No pointer arithmetic, references, calls, or unknown casts are +/// accepted here. +pub(super) fn local_cell(ctx: &Context, mut pointer: Value) -> Option<(Value, Vec)> { + let mut reversed_path = Vec::new(); + let mut visited = HashSet::new(); + loop { + if !visited.insert(pointer) { + return None; + } + let definer = pointer.defining_op()?; + let opid = Operation::get_opid(definer, ctx); + let operation = definer.deref(ctx); + if opid == MirAllocaOp::get_opid_static() { + return (operation.get_num_results() == 1 && operation.get_result(0) == pointer).then( + || { + reversed_path.reverse(); + (pointer, reversed_path) + }, + ); + } + if opid == MirCastOp::get_opid_static() + && operation.get_num_operands() == 1 + && operation.get_num_results() == 1 + && operation.get_result(0) == pointer + && is_local_cell_pointer_alias_cast(ctx, definer) + { + pointer = operation.get_operand(0); + continue; + } + if opid == MirFieldAddrOp::get_opid_static() + && operation.get_num_operands() == 1 + && operation.get_num_results() == 1 + && operation.get_result(0) == pointer + { + let index = MirFieldAddrOp::new(definer).get_attr_field_index(ctx)?; + reversed_path.push(index.0); + pointer = operation.get_operand(0); + continue; + } + return None; + } +} + +/// Every alias of an accepted cell must stay within the same small +/// cast/field/load/store vocabulary. This lets the verifier conservatively +/// merge all stores to the cell without guessing about hidden aliases. +pub(super) fn local_cell_is_closed(ctx: &Context, root: Value) -> bool { + let mut worklist = vec![root]; + let mut visited = HashSet::new(); + while let Some(pointer) = worklist.pop() { + if !visited.insert(pointer) { + continue; + } + for r#use in pointer.uses(ctx) { + let user = r#use.user_op(); + let opid = Operation::get_opid(user, ctx); + let operation = user.deref(ctx); + if r#use.find_index(ctx) != 0 { + return false; + } + let forwards_pointer = operation.get_num_results() == 1 + && ((opid == MirCastOp::get_opid_static() + && is_local_cell_pointer_alias_cast(ctx, user)) + || (opid == MirFieldAddrOp::get_opid_static() + && MirFieldAddrOp::new(user) + .get_attr_field_index(ctx) + .is_some())); + if forwards_pointer { + worklist.push(operation.get_result(0)); + } else if opid == MirLoadOp::get_opid_static() { + if MirLoadOp::new(user).is_volatile(ctx) { + return false; + } + } else if opid == MirStoreOp::get_opid_static() { + if MirStoreOp::new(user).is_volatile(ctx) { + return false; + } + } else { + return false; + } + } + } + true +} + +/// Propagate semantic facts through the exact ordinary SSA adapters that can +/// remain after common CuTe preparation. This is deliberately a monotone +/// fixed-point transfer: loop block arguments become known once their entry +/// fact reaches them, then semantic advance operations can feed the same fact +/// around the backedge on a later iteration. +pub(super) fn propagate_ordinary_plumbing( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + if let Some(block) = operation_ref.get_parent_block() { + blocks.insert(block); + } + let opid = Operation::get_opid(*operation, ctx); + if opid == MirCastOp::get_opid_static() + && is_semantic_pointer_carrier_cast(ctx, *operation) + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[], + &[], + ); + } else if (opid == MirConstructStructOp::get_opid_static() + || opid == MirConstructTupleOp::get_opid_static()) + && operation_ref.get_num_results() == 1 + { + let result = operation_ref.get_result(0); + for (index, operand) in operation_ref.operands().enumerate() { + changed |= facts.copy_entries(operand, result, &[], &[index as u32]); + } + } else if opid == MirInsertFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 2 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirInsertFieldOp::new(*operation).get_attr_insert_index(ctx) else { + continue; + }; + let base = operation_ref.get_operand(0); + let inserted = operation_ref.get_operand(1); + let result = operation_ref.get_result(0); + for (path, fact) in facts.entries(base) { + if path.first().copied() != Some(index.0) { + changed |= facts.insert_path(result, path, fact); + } + } + changed |= facts.copy_entries(inserted, result, &[], &[index.0]); + } else if opid == MirExtractFieldOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let Some(index) = MirExtractFieldOp::new(*operation).get_attr_index(ctx) else { + continue; + }; + changed |= facts.copy_entries( + operation_ref.get_operand(0), + operation_ref.get_result(0), + &[index.0], + &[], + ); + } else if opid == MirStoreOp::get_opid_static() && operation_ref.get_num_operands() == 2 { + let store = MirStoreOp::new(*operation); + if !store.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, store.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(store.value_opd(ctx), root, &[], &path); + } + } else if opid == MirLoadOp::get_opid_static() + && operation_ref.get_num_operands() == 1 + && operation_ref.get_num_results() == 1 + { + let load = MirLoadOp::new(*operation); + if !load.is_volatile(ctx) + && let Some((root, path)) = local_cell(ctx, load.address_opd(ctx)) + && local_cell_is_closed(ctx, root) + { + changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + if let Some(incoming) = operands.get(argument_index).copied() { + changed |= facts.copy_entries(incoming, argument, &[], &[]); + } + } + } + } + changed +} + +/// Enumerate integer fields of the requested bit width inside tuples, structs, and arrays. +fn integer_leaf_paths( + ctx: &Context, + ty: TypeHandle, + width: u32, + prefix: &mut Vec, + output: &mut Vec>, +) { + let ty_ref = ty.deref(ctx); + if ty_ref + .downcast_ref::() + .is_some_and(|integer| integer.width() == width) + { + output.push(prefix.clone()); + } else if let Some(structure) = ty_ref.downcast_ref::() { + for (index, field) in structure.field_types.iter().copied().enumerate() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, field, width, prefix, output); + prefix.pop(); + } + } else if let Some(tuple) = ty_ref.downcast_ref::() { + for (index, field) in tuple.get_types().iter().copied().enumerate() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, field, width, prefix, output); + prefix.pop(); + } + } else if let Some(array) = ty_ref.downcast_ref::() { + for index in 0..array.size() { + let Ok(index) = u32::try_from(index) else { + continue; + }; + prefix.push(index); + integer_leaf_paths(ctx, array.element_type(), width, prefix, output); + prefix.pop(); + } + } +} + +/// Seed factless integer leaves only after known semantic facts stabilize. +/// Aggregates are visited by field path, so an unrelated scalar cannot hide +/// inside a carrier that also contains known fields. +pub(super) fn seed_unknown_integer_carriers( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + width: u32, + unknown: impl Fn(Value) -> F, +) where + F: Clone + Eq + std::hash::Hash, +{ + let mut values = Vec::new(); + let mut blocks = HashSet::new(); + for operation in operations { + let operation_ref = operation.deref(ctx); + values.extend(operation_ref.results()); + if let Some(block) = operation_ref.get_parent_block() + && blocks.insert(block) + { + let block_ref = block.deref(ctx); + values.extend( + (0..block_ref.get_num_arguments()).map(|index| block_ref.get_argument(index)), + ); + } + } + + for value in values { + let mut paths = Vec::new(); + integer_leaf_paths(ctx, value.get_type(ctx), width, &mut Vec::new(), &mut paths); + for path in paths { + if facts.at_path(value, &path).is_none() { + facts.insert_path(value, path, unknown(value)); + } + } + } +} + +/// Mark missing inputs to a known block argument or closed local cell unknown. +/// Repeat this and ordinary propagation to a fixed point: one known incoming +/// edge must not conceal another edge or store with no semantic origin. +pub(super) fn poison_missing_forwarding_sources( + ctx: &Context, + operations: &[Ptr], + facts: &mut OrdinaryFacts, + unknown: impl Fn(Value) -> F + Copy, +) -> bool +where + F: Clone + Eq + std::hash::Hash, +{ + let mut changed = false; + let mut blocks = HashSet::new(); + for operation in operations { + if let Some(block) = operation.deref(ctx).get_parent_block() { + blocks.insert(block); + } + if Operation::get_opid(*operation, ctx) != MirStoreOp::get_opid_static() { + continue; + } + let store = MirStoreOp::new(*operation); + let Some((root, cell_path)) = local_cell(ctx, store.address_opd(ctx)) else { + continue; + }; + if !local_cell_is_closed(ctx, root) { + continue; + } + for (path, _) in facts.entries(root) { + let Some(source_path) = path.strip_prefix(cell_path.as_slice()) else { + continue; + }; + if facts.at_path(store.value_opd(ctx), source_path).is_none() { + changed |= facts.insert_path( + store.value_opd(ctx), + source_path.to_vec(), + unknown(store.value_opd(ctx)), + ); + } + } + } + + for block in blocks { + let block_ref = block.deref(ctx); + for argument_index in 0..block_ref.get_num_arguments() { + let argument = block_ref.get_argument(argument_index); + let argument_entries = facts.entries(argument); + if argument_entries.is_empty() { + continue; + } + for edge in block.uses(ctx) { + let terminator = edge.user_op(); + let Some(operands) = + branch_successor_operands(ctx, terminator, edge.find_index(ctx)) + else { + continue; + }; + let Some(incoming) = operands.get(argument_index).copied() else { + continue; + }; + for (path, _) in &argument_entries { + if facts.at_path(incoming, path).is_none() { + changed |= facts.insert_path(incoming, path.clone(), unknown(incoming)); + } + } + } + } + } + changed +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs new file mode 100644 index 0000000000..0d4e052952 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs @@ -0,0 +1,382 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Persistent output-tile grids and their ordinary cursor provenance. + +use std::collections::{HashMap, HashSet}; + +use dialect_mir::ops::{MirNotOp, control_flow::MirCondBranchOp}; +use pliron::{ + basic_block::BasicBlock, + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{ + Result, + cfg::{block_reaches, operation_reaches}, + ghosts::scheduler_semantic_ids, + invalid, + provenance::{ + OrdinaryFacts, poison_missing_forwarding_sources, propagate_ordinary_plumbing, + seed_unknown_integer_carriers, + }, +}; +use crate::{ + attributes::CuteTileGridAttr, + ops::scheduler::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + types::CuteWorkTileType, +}; + +/// Read the static output grid from a recognized scheduler operation. +fn scheduler_grid( + ctx: &Context, + operation: Ptr, + opid: &OpId, +) -> Result { + let grid = if *opid == CuteSchedulerNew1dOp::get_opid_static() { + CuteSchedulerNew1dOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerHasWorkOp::get_opid_static() { + CuteSchedulerHasWorkOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerCurrentOp::get_opid_static() { + CuteSchedulerCurrentOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteSchedulerAdvanceOp::get_opid_static() { + CuteSchedulerAdvanceOp::wrap(operation).tile_grid(ctx) + } else if *opid == CuteWorkTileCoordinatesOp::get_opid_static() { + CuteWorkTileCoordinatesOp::wrap(operation) + .work_tile(ctx) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|tile| tile.grid) + } else { + None + }; + grid.ok_or_else(|| invalid(format!("`{opid}` is missing its scheduler tile grid"))) +} + +/// Origin of an ordinary persistent-scheduler state component. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub(super) enum SchedulerFact { + /// Linear tile cursor derived from one scheduler constructor. + Current(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Physical-grid stride derived from the same scheduler constructor. + Stride(#[doc = "Operation identifying the semantic origin of this carrier."] Ptr), + /// Scalar without scheduler provenance, invalid in a scheduler-state merge. + Unknown(#[doc = "Original SSA value lacking semantic provenance."] Value), +} + +/// Require exactly one current-index or stride origin of the requested kind. +pub(super) fn scheduler_root( + facts: &OrdinaryFacts, + value: Value, + current: bool, + what: &str, +) -> Result> { + let Some(found) = facts.at(value) else { + return Err(invalid(format!( + "{what} has no scheduler provenance after ordinary SSA forwarding" + ))); + }; + if found.len() != 1 { + return Err(invalid(format!( + "{what} merges values from different scheduler stories" + ))); + } + match found.iter().next().copied() { + Some(SchedulerFact::Current(root)) if current => Ok(root), + Some(SchedulerFact::Stride(root)) if !current => Ok(root), + _ => Err(invalid(format!( + "{what} is not derived from the required scheduler_new_1d result" + ))), + } +} + +/// Propagate scheduler origins through advances and ordinary carriers before poisoning unknown edges. +pub(super) fn build_scheduler_facts( + ctx: &Context, + operations: &[Ptr], +) -> OrdinaryFacts { + let mut facts = OrdinaryFacts::default(); + for operation in operations { + if Operation::get_opid(*operation, ctx) == CuteSchedulerNew1dOp::get_opid_static() { + let scheduler = CuteSchedulerNew1dOp::wrap(*operation); + facts.insert(scheduler.current(ctx), SchedulerFact::Current(*operation)); + facts.insert(scheduler.stride(ctx), SchedulerFact::Stride(*operation)); + } + } + + loop { + let mut changed = propagate_ordinary_plumbing(ctx, operations, &mut facts); + for operation in operations { + if Operation::get_opid(*operation, ctx) != CuteSchedulerAdvanceOp::get_opid_static() { + continue; + } + let advance = CuteSchedulerAdvanceOp::wrap(*operation); + let Some(current_facts) = facts.at(advance.current(ctx)).cloned() else { + continue; + }; + let Some(stride_facts) = facts.at(advance.stride(ctx)).cloned() else { + continue; + }; + for current in ¤t_facts { + let SchedulerFact::Current(root) = current else { + continue; + }; + if stride_facts.contains(&SchedulerFact::Stride(*root)) { + changed |= facts.insert(advance.next(ctx), SchedulerFact::Current(*root)); + } + } + } + if !changed { + break; + } + } + + // A factless scalar is an ordinary, unrelated value. Seeding it only + // after the semantic fixed point lets a valid loop recurrence establish + // itself first, while still poisoning a merge with an arbitrary edge. + seed_unknown_integer_carriers(ctx, operations, &mut facts, 64, SchedulerFact::Unknown); + loop { + let poisoned = + poison_missing_forwarding_sources(ctx, operations, &mut facts, SchedulerFact::Unknown); + let propagated = propagate_ordinary_plumbing(ctx, operations, &mut facts); + if !poisoned && !propagated { + break; + } + } + facts +} + +/// Find a branch whose has-work condition guards the selected scheduler cursor. +fn scheduler_guard( + ctx: &Context, + has_work: Value, +) -> Option<(Ptr, Ptr, Ptr)> { + let mut value = has_work; + let mut negated = false; + loop { + let uses = value.uses(ctx); + if uses.len() != 1 || uses[0].find_index(ctx) != 0 { + return None; + } + let user = uses[0].user_op(); + let opid = Operation::get_opid(user, ctx); + if opid == MirNotOp::get_opid_static() { + let operation = user.deref(ctx); + if operation.get_num_operands() != 1 || operation.get_num_results() != 1 { + return None; + } + value = operation.get_result(0); + negated = !negated; + continue; + } + if opid != MirCondBranchOp::get_opid_static() { + return None; + } + let operation = user.deref(ctx); + if operation.get_num_successors() != 2 || operation.get_operand(0) != value { + return None; + } + let work_index = usize::from(negated); + return Some(( + user, + operation.get_successor(work_index), + operation.get_successor(1 - work_index), + )); + } +} + +/// Check one scheduler's state origins, loop advance, work guard, and coordinate consumers. +pub(super) fn verify_scheduler_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = scheduler_semantic_ids(); + let facts = build_scheduler_facts(ctx, all_ops); + let mut roots = HashSet::new(); + let mut has_work_roots = HashSet::new(); + let mut current_roots = HashSet::new(); + let mut advance_roots = HashSet::new(); + let mut coordinate_roots = HashSet::new(); + let mut has_work_by_root: HashMap, Vec>> = HashMap::new(); + let mut current_by_root: HashMap, Vec>> = HashMap::new(); + let mut advance_by_root: HashMap, Vec>> = HashMap::new(); + let mut coordinates_by_root: HashMap, Vec>> = HashMap::new(); + let mut currents = Vec::new(); + let mut coordinates = Vec::new(); + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + if operation.deref(ctx).get_parent_block().is_none() { + return Err(invalid(format!( + "scheduler operation `{opid}` is not attached to a block" + ))); + } + let grid = scheduler_grid(ctx, *operation, &opid)?; + if opid == CuteSchedulerNew1dOp::get_opid_static() { + roots.insert(*operation); + } else if opid == CuteSchedulerHasWorkOp::get_opid_static() { + let semantic = CuteSchedulerHasWorkOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_has_work current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_has_work grid differs from its scheduler_new_1d root", + )); + } + if !semantic.has_work(ctx).is_used(ctx) { + return Err(invalid( + "cute.scheduler_has_work result is unused; it must guard the scheduler body", + )); + } + has_work_roots.insert(root); + has_work_by_root.entry(root).or_default().push(*operation); + } else if opid == CuteSchedulerCurrentOp::get_opid_static() { + let semantic = CuteSchedulerCurrentOp::wrap(*operation); + let root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_current current", + )?; + if CuteSchedulerNew1dOp::wrap(root).tile_grid(ctx) != Some(grid) { + return Err(invalid( + "cute.scheduler_current grid differs from its scheduler_new_1d root", + )); + } + current_roots.insert(root); + current_by_root.entry(root).or_default().push(*operation); + currents.push(*operation); + } else if opid == CuteWorkTileCoordinatesOp::get_opid_static() { + coordinates.push(*operation); + } else if opid == CuteSchedulerAdvanceOp::get_opid_static() { + let semantic = CuteSchedulerAdvanceOp::wrap(*operation); + let current_root = scheduler_root( + &facts, + semantic.current(ctx), + true, + "cute.scheduler_advance current", + )?; + let stride_root = scheduler_root( + &facts, + semantic.stride(ctx), + false, + "cute.scheduler_advance stride", + )?; + if current_root != stride_root + || CuteSchedulerNew1dOp::wrap(current_root).tile_grid(ctx) != Some(grid) + { + return Err(invalid( + "cute.scheduler_advance must use current, stride, and grid from one scheduler_new_1d", + )); + } + advance_roots.insert(current_root); + advance_by_root + .entry(current_root) + .or_default() + .push(*operation); + } + } + + for current in ¤ts { + let tile = CuteSchedulerCurrentOp::wrap(*current).work_tile(ctx); + let uses = tile.uses(ctx); + if uses.is_empty() + || uses.iter().any(|r#use| { + Operation::get_opid(r#use.user_op(), ctx) + != CuteWorkTileCoordinatesOp::get_opid_static() + || r#use.find_index(ctx) != 0 + }) + { + return Err(invalid( + "cute.scheduler_current work tile must be used only by at least one cute.work_tile_coordinates", + )); + } + } + for coordinates in coordinates { + let tile = CuteWorkTileCoordinatesOp::wrap(coordinates).work_tile(ctx); + let Some(producer) = tile.defining_op() else { + return Err(invalid( + "cute.work_tile_coordinates reached a block argument; a direct scheduler_current result is required", + )); + }; + if Operation::get_opid(producer, ctx) != CuteSchedulerCurrentOp::get_opid_static() + || !currents.contains(&producer) + { + return Err(invalid( + "cute.work_tile_coordinates must use a scheduler_current attached to this module", + )); + } + let root = scheduler_root( + &facts, + CuteSchedulerCurrentOp::wrap(producer).current(ctx), + true, + "cute.work_tile_coordinates scheduler current", + )?; + coordinate_roots.insert(root); + coordinates_by_root + .entry(root) + .or_default() + .push(coordinates); + } + + for root in roots { + if !has_work_roots.contains(&root) + || !current_roots.contains(&root) + || !coordinate_roots.contains(&root) + || !advance_roots.contains(&root) + { + return Err(invalid( + "each cute.scheduler_new_1d needs a complete has_work -> current -> coordinates -> advance story", + )); + } + let has_work = &has_work_by_root[&root]; + let current = ¤t_by_root[&root]; + let coordinates = &coordinates_by_root[&root]; + let advances = &advance_by_root[&root]; + if has_work.len() != 1 + || current.len() != 1 + || coordinates.len() != 1 + || advances.len() != 1 + { + return Err(invalid( + "each scheduler root needs exactly one static has_work/current/coordinates/advance protocol", + )); + } + let Some((guard, work_successor, exit_successor)) = + scheduler_guard(ctx, CuteSchedulerHasWorkOp::wrap(has_work[0]).has_work(ctx)) + else { + return Err(invalid( + "cute.scheduler_has_work must exclusively guard one conditional branch (optionally through mir.not)", + )); + }; + let current_block = current[0] + .deref(ctx) + .get_parent_block() + .expect("attached scheduler current"); + if !operation_reaches(ctx, root, has_work[0]) + || !block_reaches(ctx, work_successor, current_block) + || block_reaches(ctx, exit_successor, current_block) + || !operation_reaches(ctx, current[0], coordinates[0]) + || !operation_reaches(ctx, coordinates[0], advances[0]) + || !operation_reaches(ctx, advances[0], guard) + { + return Err(invalid( + "scheduler CFG must be new -> guarded work current -> coordinates -> advance -> next guard, with the exit edge outside the work body", + )); + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tensor.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tensor.rs new file mode 100644 index 0000000000..00546d1d92 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tensor.rs @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tensor storage, access, format, and selection metadata. + +use pliron::{ + context::{Context, Ptr}, + op::Op, + operation::Operation, + r#type::Typed, + value::Value, +}; + +use super::{Result, ghosts::tensor_semantic_ids, invalid}; +use crate::{ + ops::tensor::{ + CuteTensorMakeOp, CuteTensorSliceOp, CuteTensorStoreFromOp, CuteTensorZippedDivideOp, + }, + types::CuteTensorViewType, +}; + +/// Tile width and selection recovered from a direct tensor-view chain. +#[derive(Clone, Copy)] +struct TensorViewState { + /// Static tile width in elements, absent before zipped division. + tile_size: Option, + /// Runtime selected tile number, absent before slicing. + tile_index: Option, +} + +/// Walk direct make/divide/slice producers and reject unsupported or excessively deep chains. +fn resolve_tensor_view(ctx: &Context, value: Value, depth: usize) -> Result { + if depth > 16 { + return Err(invalid( + "CuTe tensor-view producer chain is unexpectedly deep", + )); + } + let Some(defining_op) = value.defining_op() else { + return Err(invalid( + "CuTe tensor view reached a block argument; a direct make/divide/slice chain is required", + )); + }; + let opid = Operation::get_opid(defining_op, ctx); + if opid == CuteTensorMakeOp::get_opid_static() { + return Ok(TensorViewState { + tile_size: None, + tile_index: None, + }); + } + if opid == CuteTensorZippedDivideOp::get_opid_static() { + let mut state = resolve_tensor_view(ctx, defining_op.deref(ctx).get_operand(0), depth + 1)?; + if state.tile_size.is_some() || state.tile_index.is_some() { + return Err(invalid( + "cute.tensor_zipped_divide must occur exactly once before tensor_slice", + )); + } + let result_ty = value.get_type(ctx); + let result_ty = result_ty.deref(ctx); + let view = result_ty + .downcast_ref::() + .ok_or_else(|| invalid("cute.tensor_zipped_divide result is not a tensor view"))?; + state.tile_size = view.layout.tile_size(); + return Ok(state); + } + if opid == CuteTensorSliceOp::get_opid_static() { + let mut state = resolve_tensor_view(ctx, defining_op.deref(ctx).get_operand(0), depth + 1)?; + if state.tile_size.is_none() || state.tile_index.is_some() { + return Err(invalid( + "cute.tensor_slice must follow exactly one zipped divide and may select only once", + )); + } + state.tile_index = Some(defining_op.deref(ctx).get_operand(1)); + return Ok(state); + } + Err(invalid(format!( + "tensor view is produced by unsupported operation `{opid}`" + ))) +} + +/// Require direct view provenance and valid full-tile or scalar-tail consumers. +pub(super) fn verify_tensor_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + let ids = tensor_semantic_ids(); + let consumer_ids = &ids[3..]; + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if !ids.contains(&opid) { + continue; + } + let op = operation.deref(ctx); + if opid == ids[0] || opid == ids[1] || opid == ids[2] { + let result = op.get_result(0); + let uses = result.uses(ctx); + if uses.is_empty() { + return Err(invalid(format!( + "tensor view from `{opid}` has no semantic consumer" + ))); + } + for r#use in uses { + let user = r#use.user_op(); + let user_id = Operation::get_opid(user, ctx); + let allowed = if opid == ids[0] { + user_id == ids[1] && r#use.find_index(ctx) == 0 + } else if opid == ids[1] { + user_id == ids[2] && r#use.find_index(ctx) == 0 + } else { + consumer_ids.contains(&user_id) + }; + if !allowed { + return Err(invalid(format!( + "tensor view from `{opid}` has unsupported consumer `{user_id}`" + ))); + } + } + } else { + let tensor = if opid == CuteTensorStoreFromOp::get_opid_static() { + op.get_operand(1) + } else { + op.get_operand(0) + }; + let state = resolve_tensor_view(ctx, tensor, 0)?; + if state.tile_size.is_none() || state.tile_index.is_none() { + return Err(invalid(format!( + "`{opid}` needs a complete make -> zipped_divide -> slice tensor story" + ))); + } + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs new file mode 100644 index 0000000000..30d90d2765 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs @@ -0,0 +1,754 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe operation, type, and semantic verification contracts. + +use dialect_mir::{ + attributes::MirCastKindAttr, + ops::{ + MirAllocaOp, MirCastOp, MirGotoOp, MirLoadOp, MirNotOp, MirReturnOp, MirStoreOp, + MirUndefOp, control_flow::MirCondBranchOp, function::MirFuncOp, + }, + types::{MirPtrType, address_space}, +}; +use pliron::{ + basic_block::BasicBlock, + builtin::{ + attributes::TypeAttr, + op_interfaces::{OperandSegmentInterface, SymbolOpInterface}, + ops::ModuleOp, + types::{FP32Type, FunctionType, IntegerType, Signedness}, + }, + context::{Context, Ptr}, + linked_list::ContainsLinkedList, + op::Op, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +use super::{ + epilogue::require_epilogue_cfg_links, + mma::{build_mma_facts, require_smem_overlay_pair, verify_mma_use_closure}, + pipeline::require_producer_pipeline_order, + provenance::{ + is_local_cell_pointer_alias_cast, is_semantic_pointer_carrier_cast, local_cell_is_closed, + }, + scheduler::{build_scheduler_facts, scheduler_root}, + *, +}; +use crate::{ + attributes::{ + CutePipelineStateAttr, CuteTensorAccessAttr, CuteTensorAddressSpaceAttr, + CuteTensorLayoutAttr, CuteTileGridAttr, CuteTiledMmaPlanAttr, + }, + layout::{ComposedLayout, OffsetUnit}, + ops::gemm_tma::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + ops::mma::{CuteMmaPartitionBOp, CuteSmemTensorOverlayOp}, + ops::pipeline::{ + CutePipelineConsumerWaitOp, CutePipelineProducerExpectTxOp, CutePipelineStateNewOp, + CuteTmaLoadPipelineMakeOp, + }, + ops::scheduler::{ + CuteSchedulerAdvanceOp, CuteSchedulerCurrentOp, CuteSchedulerHasWorkOp, + CuteSchedulerNew1dOp, CuteWorkTileCoordinatesOp, + }, + types::CuteTensorViewType, +}; + +/// Create a test module and return its entry block for constructing semantic IR directly. +fn module_top(ctx: &mut Context) -> (Ptr, Ptr) { + dialect_mir::register(ctx); + crate::register(ctx); + let module = ModuleOp::new(ctx, "cute_verify_test".try_into().unwrap()); + let module = module.get_operation(); + let region = module.deref(ctx).get_region(0); + let block = region.deref(ctx).iter(ctx).next().unwrap(); + (module, block) +} + +/// Create a typed placeholder SSA result so a verifier test needs no runtime input. +fn undef(ctx: &mut Context, block: Ptr, ty: TypeHandle) -> Value { + let operation = MirUndefOp::new(ctx, ty).get_operation(); + operation.insert_at_back(block, ctx); + operation.deref(ctx).get_result(0) +} + +/// Construct a valid compiler-only contiguous tensor type for escape-audit tests. +fn ghost_tensor_type(ctx: &mut Context) -> TypeHandle { + let element: TypeHandle = FP32Type::get(ctx).into(); + CuteTensorViewType::get( + ctx, + element, + element, + CuteTensorAddressSpaceAttr::Gmem, + CuteTensorAccessAttr::ReadOnly, + 4, + CuteTensorLayoutAttr::Contiguous1D, + ) + .into() +} + +/// Append a compiler-owned local allocation used to test closed carrier storage. +fn local_cell(ctx: &mut Context, block: Ptr, ty: TypeHandle) -> Value { + let pointer: TypeHandle = MirPtrType::get_generic(ctx, ty, true).into(); + let alloca = Operation::new( + ctx, + MirAllocaOp::get_concrete_op_info(), + vec![pointer], + vec![], + vec![], + 0, + ); + alloca.insert_at_back(block, ctx); + alloca.deref(ctx).get_result(0) +} + +/// Append a non-volatile MIR store into the fixture's local cell. +fn store(ctx: &mut Context, block: Ptr, pointer: Value, value: Value) { + Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![pointer, value], + vec![], + 0, + ) + .insert_at_back(block, ctx); +} + +/// Append a non-volatile MIR load and return its SSA result. +fn load(ctx: &mut Context, block: Ptr, pointer: Value, ty: TypeHandle) -> Value { + let load = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![ty], + vec![pointer], + vec![], + 0, + ); + load.insert_at_back(block, ctx); + load.deref(ctx).get_result(0) +} + +/// Append a pointer cast with controlled pointee, mutability, and address-space facts. +fn pointer_cast( + ctx: &mut Context, + block: Ptr, + pointer: Value, + result_type: TypeHandle, +) -> Ptr { + let cast = Operation::new( + ctx, + MirCastOp::get_concrete_op_info(), + vec![result_type], + vec![pointer], + vec![], + 0, + ); + MirCastOp::new(cast).set_attr_cast_kind(ctx, MirCastKindAttr::PtrToPtr); + cast.insert_at_back(block, ctx); + cast +} + +/// Append an unconditional successor edge with explicit block-argument operands. +fn goto(ctx: &mut Context, block: Ptr, target: Ptr, operands: Vec) { + Operation::new( + ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + operands, + vec![target], + 0, + ) + .insert_at_back(block, ctx); +} + +/// Append a conditional branch with distinct successor operands for CFG/provenance tests. +fn cond_branch( + ctx: &mut Context, + block: Ptr, + condition: Value, + true_target: Ptr, + false_target: Ptr, +) { + let (operands, sizes) = + MirCondBranchOp::compute_segment_sizes(vec![vec![condition], Vec::new(), Vec::new()]); + let branch = Operation::new( + ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + operands, + vec![true_target, false_target], + 0, + ); + MirCondBranchOp::new(branch).set_operand_segment_sizes(ctx, sizes); + branch.insert_at_back(block, ctx); +} + +/// Construct one valid persistent-scheduler function with configurable body/exit placement. +fn append_complete_scheduler_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + grid: CuteTileGridAttr, + misroute_work_edge: bool, +) { + let function_type = FunctionType::get(ctx, vec![], vec![]); + let function = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function_op = MirFuncOp::new(ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(ctx, name.try_into().unwrap()); + function.insert_at_back(module_block, ctx); + let region = function.deref(ctx).get_region(0); + let u64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Unsigned).into(); + let entry = BasicBlock::new(ctx, None, vec![]); + entry.insert_at_back(region, ctx); + let header = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); + header.insert_at_back(region, ctx); + let body = BasicBlock::new(ctx, None, vec![]); + body.insert_at_back(region, ctx); + let exit = BasicBlock::new(ctx, None, vec![]); + exit.insert_at_back(region, ctx); + + let start = CuteSchedulerNew1dOp::new(ctx, grid).get_operation(); + start.insert_at_back(entry, ctx); + let start = CuteSchedulerNew1dOp::wrap(start); + goto( + ctx, + entry, + header, + vec![start.current(ctx), start.stride(ctx)], + ); + + let current = header.deref(ctx).get_argument(0); + let stride = header.deref(ctx).get_argument(1); + let has_work = CuteSchedulerHasWorkOp::new(ctx, current, grid).get_operation(); + has_work.insert_at_back(header, ctx); + let not = Operation::new( + ctx, + MirNotOp::get_concrete_op_info(), + vec![ + CuteSchedulerHasWorkOp::wrap(has_work) + .has_work(ctx) + .get_type(ctx), + ], + vec![CuteSchedulerHasWorkOp::wrap(has_work).has_work(ctx)], + vec![], + 0, + ); + not.insert_at_back(header, ctx); + let finished = not.deref(ctx).get_result(0); + let (true_target, false_target) = if misroute_work_edge { + (body, exit) + } else { + (exit, body) + }; + cond_branch(ctx, header, finished, true_target, false_target); + + let selected = CuteSchedulerCurrentOp::new(ctx, current, grid).get_operation(); + selected.insert_at_back(body, ctx); + CuteWorkTileCoordinatesOp::new(ctx, CuteSchedulerCurrentOp::wrap(selected).work_tile(ctx)) + .get_operation() + .insert_at_back(body, ctx); + let advance = CuteSchedulerAdvanceOp::new(ctx, current, stride, grid).get_operation(); + advance.insert_at_back(body, ctx); + goto( + ctx, + body, + header, + vec![CuteSchedulerAdvanceOp::wrap(advance).next(ctx), stride], + ); + + Operation::new( + ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(exit, ctx); +} + +/// The escape audit must inspect ordinary IR types even when no semantic CuTe operation exposes the ghost. +#[test] +fn rejects_a_ghost_nested_in_an_ordinary_pointer_without_any_cute_op() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let ghost = ghost_tensor_type(&mut ctx); + let inner: TypeHandle = MirPtrType::get_generic(&mut ctx, ghost, false).into(); + let outer: TypeHandle = MirPtrType::get_generic(&mut ctx, inner, false).into(); + undef(&mut ctx, block, outer); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("ghost CuTe value"), "{error}"); + assert!(error.contains("operand or result"), "{error}"); +} + +/// Function scopes isolate scheduler origins, and verification must leave the input module unchanged. +#[test] +fn accepts_independent_scheduler_stories_in_two_functions_without_mutation() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "first", + CuteTileGridAttr::new(4, 4, 1), + false, + ); + append_complete_scheduler_function( + &mut ctx, + block, + "second", + CuteTileGridAttr::new(5, 4, 1), + false, + ); + let mut before = Vec::new(); + collect_ops(&ctx, module, &mut before); + + verify_cute_semantics(&ctx, module).unwrap(); + + let mut after = Vec::new(); + collect_ops(&ctx, module, &mut after); + assert_eq!(after, before); +} + +/// CFG work guards, rather than textual traversal order, determine whether a selected tile is valid. +#[test] +fn rejects_a_scheduler_body_on_the_exit_edge_even_when_preorder_looks_ordered() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "misrouted", + CuteTileGridAttr::new(4, 4, 1), + true, + ); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("scheduler CFG"), "{error}"); +} + +/// Non-escaping compiler-owned loads/stores may forward a scheduler origin without requiring full SSA promotion. +#[test] +fn accepts_a_closed_local_cell_carrying_scheduler_provenance() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let grid = CuteTileGridAttr::new(4, 4, 1); + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid).get_operation(); + start.insert_at_back(block, &ctx); + let pointer = local_cell(&mut ctx, block, u64_ty); + let current = CuteSchedulerNew1dOp::wrap(start).current(&ctx); + store(&mut ctx, block, pointer, current); + let carried = load(&mut ctx, block, pointer, u64_ty); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + + let facts = build_scheduler_facts(&ctx, &operations); + assert_eq!( + scheduler_root(&facts, carried, true, "test carrier").unwrap(), + start + ); +} + +/// One unrelated incoming store must invalidate a cell even when another store has valid scheduler provenance. +#[test] +fn factless_block_argument_store_poisons_a_valid_scheduler_cell() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![u64_ty]); + let arbitrary = block.deref(&ctx).get_argument(0); + let grid = CuteTileGridAttr::new(4, 4, 1); + let start = CuteSchedulerNew1dOp::new(&mut ctx, grid).get_operation(); + start.insert_at_back(block, &ctx); + let pointer = local_cell(&mut ctx, block, u64_ty); + let current = CuteSchedulerNew1dOp::wrap(start).current(&ctx); + store(&mut ctx, block, pointer, current); + store(&mut ctx, block, pointer, arbitrary); + let poisoned = load(&mut ctx, block, pointer, u64_ty); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + + let facts = build_scheduler_facts(&ctx, &operations); + let error = scheduler_root(&facts, poisoned, true, "test carrier") + .unwrap_err() + .to_string(); + assert!(error.contains("merges values"), "{error}"); +} + +/// Casts that change pointee, address space, or write capability cannot establish closed-cell aliases. +#[test] +fn non_identity_pointer_casts_cannot_alias_a_local_semantic_cell() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let pointer = local_cell(&mut ctx, block, u64_ty); + let shared_pointer_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let different_pointee_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); + for result_type in [shared_pointer_ty, different_pointee_ty] { + let cast = pointer_cast(&mut ctx, block, pointer, result_type); + let alias = cast.deref(&ctx).get_result(0); + assert!(!is_local_cell_pointer_alias_cast(&ctx, cast)); + assert!(super::provenance::local_cell(&ctx, alias).is_none()); + } + let read_only_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, false).into(); + let read_only = pointer_cast(&mut ctx, block, pointer, read_only_type); + let read_only_alias = read_only.deref(&ctx).get_result(0); + assert!(is_local_cell_pointer_alias_cast(&ctx, read_only)); + assert!(super::provenance::local_cell(&ctx, read_only_alias).is_some()); + + let mutable_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, true).into(); + let escalation = pointer_cast(&mut ctx, block, read_only_alias, mutable_type); + let escalated_alias = escalation.deref(&ctx).get_result(0); + assert!(!is_local_cell_pointer_alias_cast(&ctx, escalation)); + assert!(super::provenance::local_cell(&ctx, escalated_alias).is_none()); + assert!(!local_cell_is_closed(&ctx, pointer)); +} + +/// Semantic shared carriers may survive the exact generic/shared round-trip while other address-space casts are rejected. +#[test] +fn semantic_pointer_carriers_allow_only_shared_generic_address_space_roundtrips() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let shared_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let generic_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, true).into(); + let shared = undef(&mut ctx, block, shared_type); + let to_generic = pointer_cast(&mut ctx, block, shared, generic_type); + assert!(is_semantic_pointer_carrier_cast(&ctx, to_generic)); + let generic = to_generic.deref(&ctx).get_result(0); + let to_shared = pointer_cast(&mut ctx, block, generic, shared_type); + assert!(is_semantic_pointer_carrier_cast(&ctx, to_shared)); + + let global_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::GLOBAL).into(); + let global = undef(&mut ctx, block, global_type); + let global_to_generic = pointer_cast(&mut ctx, block, global, generic_type); + assert!(!is_semantic_pointer_carrier_cast(&ctx, global_to_generic)); + + let generic_u32_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, true).into(); + let wrong_pointee = pointer_cast(&mut ctx, block, shared, generic_u32_type); + assert!(!is_semantic_pointer_carrier_cast(&ctx, wrong_pointee)); + + let read_only_shared_type: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, false, address_space::SHARED).into(); + let read_only_shared = undef(&mut ctx, block, read_only_shared_type); + let mutability_escalation = pointer_cast(&mut ctx, block, read_only_shared, generic_type); + assert!(!is_semantic_pointer_carrier_cast( + &ctx, + mutability_escalation + )); +} + +/// Matching scalar types cannot substitute a consumer-origin cursor for producer lifecycle state. +#[test] +fn rejects_pipeline_lifecycle_state_from_the_wrong_role_root() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let barrier_base = undef(&mut ctx, block, barrier_ty); + let make = CuteTmaLoadPipelineMakeOp::new(&mut ctx, barrier_base, 3, 8, 8, 8).get_operation(); + make.insert_at_back(block, &ctx); + let pipeline = CuteTmaLoadPipelineMakeOp::wrap(make).pipeline(&ctx); + let producer = + CutePipelineStateNewOp::new(&mut ctx, CutePipelineStateAttr::producer(3)).get_operation(); + producer.insert_at_back(block, &ctx); + let producer = CutePipelineStateNewOp::wrap(producer); + let producer_slot = producer.slot(&ctx); + let producer_phase = producer.phase(&ctx); + CutePipelineConsumerWaitOp::new( + &mut ctx, + pipeline, + producer_slot, + producer_phase, + CutePipelineStateAttr::consumer(3), + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("state facts differ"), "{error}"); +} + +/// Producer drain must follow the cursor's acquire/expect/advance chain in the CFG. +#[test] +fn rejects_producer_tail_that_precedes_its_lifecycle() { + let mut ctx = Context::new(); + module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let block = BasicBlock::new(&mut ctx, None, vec![]); + let root = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let tail = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let acquire = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let expect = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + let advance = undef(&mut ctx, block, u64_ty).defining_op().unwrap(); + + let error = require_producer_pipeline_order(&ctx, root, acquire, expect, advance, tail) + .unwrap_err() + .to_string(); + assert!(error.contains("producer pipeline CFG"), "{error}"); +} + +/// A correctly typed barrier still needs provenance from the matching transaction-expectation operation. +#[test] +fn rejects_tma_copy_with_an_arbitrary_typed_completion_barrier() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let arbitrary_barrier = undef(&mut ctx, block, barrier_ty); + let layout = + ComposedLayout::from_layout("(1,16):(16,1)".parse().unwrap(), OffsetUnit::Elements); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let descriptor = undef(&mut ctx, block, descriptor_ty); + let shared = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let row = undef(&mut ctx, block, u64_ty); + let column = undef(&mut ctx, block, u64_ty); + let source = + CuteTmaGmemViewOp::new(&mut ctx, descriptor, u8_ty, layout.clone()).get_operation(); + source.insert_at_back(block, &ctx); + let destination = + CuteTmaSmemViewOp::new(&mut ctx, shared, capacity, u8_ty, layout, 16).get_operation(); + destination.insert_at_back(block, &ctx); + let source_view = source.deref(&ctx).get_result(0); + let destination_view = destination.deref(&ctx).get_result(0); + CuteTmaCopy2dOp::new( + &mut ctx, + source_view, + destination_view, + row, + column, + arbitrary_barrier, + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!( + error.contains("completion barrier is not the result of an expect_tx"), + "{error}" + ); +} + +/// Shared base and capacity must share one overlay origin, even when both overlays have identical types. +#[test] +fn rejects_mma_base_and_capacity_from_unrelated_same_typed_overlays() { + let mut ctx = Context::new(); + let (_module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let first_base = undef(&mut ctx, block, shared_ty); + let first_capacity = undef(&mut ctx, block, u64_ty); + let second_base = undef(&mut ctx, block, shared_ty); + let second_capacity = undef(&mut ctx, block, u64_ty); + let first = + CuteSmemTensorOverlayOp::new(&mut ctx, first_base, first_capacity, u64_ty).get_operation(); + first.insert_at_back(block, &ctx); + let second = CuteSmemTensorOverlayOp::new(&mut ctx, second_base, second_capacity, u64_ty) + .get_operation(); + second.insert_at_back(block, &ctx); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let facts = build_mma_facts(&ctx, &operations); + + let error = require_smem_overlay_pair( + &facts, + CuteSmemTensorOverlayOp::wrap(first).base(&ctx), + CuteSmemTensorOverlayOp::wrap(second).capacity(&ctx), + "test MMA carrier", + ) + .unwrap_err() + .to_string(); + assert!(error.contains("different shared overlays"), "{error}"); +} + +/// One valid semantic use does not permit another use to escape the overlay's closed carrier flow. +#[test] +fn rejects_mma_overlay_with_a_valid_partition_edge_and_an_unrelated_load() { + let mut ctx = Context::new(); + let (_module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let base = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let overlay = CuteSmemTensorOverlayOp::new(&mut ctx, base, capacity, u64_ty).get_operation(); + overlay.insert_at_back(block, &ctx); + let overlay = CuteSmemTensorOverlayOp::wrap(overlay); + let overlay_base = overlay.base(&ctx); + let overlay_capacity = overlay.capacity(&ctx); + let warp_n = undef(&mut ctx, block, u64_ty); + let k_half = undef(&mut ctx, block, u64_ty); + let placement = + ComposedLayout::from_layout("(128,32):(32,1)".parse().unwrap(), OffsetUnit::Elements); + CuteMmaPartitionBOp::new( + &mut ctx, + overlay_base, + overlay_capacity, + warp_n, + k_half, + u64_ty, + CuteTiledMmaPlanAttr::mxf4_128x128x128(placement), + ) + .get_operation() + .insert_at_back(block, &ctx); + Operation::new( + &mut ctx, + MirLoadOp::get_concrete_op_info(), + vec![u8_ty], + vec![overlay_base], + vec![], + 0, + ) + .insert_at_back(block, &ctx); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let facts = build_mma_facts(&ctx, &operations); + + let error = verify_mma_use_closure(&ctx, &operations, &facts) + .unwrap_err() + .to_string(); + assert!(error.contains("unsupported outgoing use"), "{error}"); + assert!(error.contains("mir.load"), "{error}"); +} + +/// A plausible traversal order cannot replace a reachable writer-to-issuer synchronization path. +#[test] +fn rejects_epilogue_writer_and_issuer_in_mutually_exclusive_preorder_blocks() { + let mut ctx = Context::new(); + let (_module, module_block) = module_top(&mut ctx); + let function_type = FunctionType::get(&ctx, vec![], vec![]); + let function = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())) + .set_symbol_name(&mut ctx, "epilogue_cfg".try_into().unwrap()); + function.insert_at_back(module_block, &ctx); + let region = function.deref(&ctx).get_region(0); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + let writer = BasicBlock::new(&mut ctx, None, vec![]); + writer.insert_at_back(region, &ctx); + let issuer = BasicBlock::new(&mut ctx, None, vec![]); + issuer.insert_at_back(region, &ctx); + let i1_ty: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let condition = undef(&mut ctx, entry, i1_ty); + cond_branch(&mut ctx, entry, condition, writer, issuer); + let acquire = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + let reusable = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + let ready = undef(&mut ctx, writer, u64_ty).defining_op().unwrap(); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(writer, &ctx); + let half_zero = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + let commit = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + let tail = undef(&mut ctx, issuer, u64_ty).defining_op().unwrap(); + Operation::new( + &mut ctx, + MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ) + .insert_at_back(issuer, &ctx); + + let error = require_epilogue_cfg_links(&ctx, acquire, reusable, ready, half_zero, commit, tail) + .unwrap_err() + .to_string(); + assert!(error.contains("epilogue CFG"), "{error}"); +} + +/// The expectation promise must equal the sum of its direct TMA users' static tile bytes. +#[test] +fn rejects_expect_tx_when_direct_tma_bytes_do_not_match() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let barrier_ty: TypeHandle = + MirPtrType::get(&mut ctx, u64_ty, true, address_space::SHARED).into(); + let barrier_base = undef(&mut ctx, block, barrier_ty); + let make = CuteTmaLoadPipelineMakeOp::new(&mut ctx, barrier_base, 3, 8, 8, 8).get_operation(); + make.insert_at_back(block, &ctx); + let pipeline = CuteTmaLoadPipelineMakeOp::wrap(make).pipeline(&ctx); + let state = CutePipelineStateAttr::producer(3); + let state_root = CutePipelineStateNewOp::new(&mut ctx, state).get_operation(); + state_root.insert_at_back(block, &ctx); + let slot = CutePipelineStateNewOp::wrap(state_root).slot(&ctx); + let expect = + CutePipelineProducerExpectTxOp::new(&mut ctx, pipeline, slot, state).get_operation(); + expect.insert_at_back(block, &ctx); + + let layout = + ComposedLayout::from_layout("(1,16):(16,1)".parse().unwrap(), OffsetUnit::Elements); + let descriptor_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, u8_ty, false).into(); + let shared_ty: TypeHandle = + MirPtrType::get(&mut ctx, u8_ty, true, address_space::SHARED).into(); + let descriptor = undef(&mut ctx, block, descriptor_ty); + let shared = undef(&mut ctx, block, shared_ty); + let capacity = undef(&mut ctx, block, u64_ty); + let row = undef(&mut ctx, block, u64_ty); + let column = undef(&mut ctx, block, u64_ty); + let source = + CuteTmaGmemViewOp::new(&mut ctx, descriptor, u8_ty, layout.clone()).get_operation(); + source.insert_at_back(block, &ctx); + let destination = + CuteTmaSmemViewOp::new(&mut ctx, shared, capacity, u8_ty, layout, 16).get_operation(); + destination.insert_at_back(block, &ctx); + let source_view = source.deref(&ctx).get_result(0); + let destination_view = destination.deref(&ctx).get_result(0); + let completion = CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(&ctx); + CuteTmaCopy2dOp::new( + &mut ctx, + source_view, + destination_view, + row, + column, + completion, + ) + .get_operation() + .insert_at_back(block, &ctx); + + let error = verify_cute_semantics(&ctx, module).unwrap_err().to_string(); + assert!(error.contains("promises 8 bytes"), "{error}"); + assert!(error.contains("move 16 bytes"), "{error}"); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tma.rs b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tma.rs new file mode 100644 index 0000000000..e92b5b0029 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tma.rs @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Descriptor-backed raw transport views and their direct copy contracts. + +use pliron::{ + context::{Context, Ptr}, + op::{Op, OpId}, + operation::Operation, + value::Value, +}; + +use super::{Result, invalid}; +use crate::{ + ops::epilogue::CuteTmaStore2dSemanticOp, + ops::gemm_tma::{CuteTmaCopy2dOp, CuteTmaGmemViewOp, CuteTmaSmemViewOp}, + ops::pipeline::CutePipelineProducerExpectTxOp, +}; + +/// Require a view operand to come directly from the expected constructor in this function scope. +pub(super) fn direct_tma_view_producer( + ctx: &Context, + value: Value, + expected: &OpId, + what: &str, + scope: &[Ptr], +) -> Result> { + let Some(producer) = value.defining_op() else { + return Err(invalid(format!( + "semantic TMA {what} reached a block argument; a direct view producer is required" + ))); + }; + let actual = Operation::get_opid(producer, ctx); + if actual != *expected { + return Err(invalid(format!( + "semantic TMA {what} is produced by `{actual}` instead of `{expected}`" + ))); + } + if !scope.contains(&producer) { + return Err(invalid(format!( + "semantic TMA {what} is produced outside its function/story scope" + ))); + } + Ok(producer) +} + +/// Restrict global/shared view ghosts to the matching TMA copy operand positions. +fn verify_tma_view_uses(ctx: &Context, producer: Ptr, gmem: bool) -> Result { + let value = producer.deref(ctx).get_result(0); + let uses = value.uses(ctx); + if uses.is_empty() { + return Err(invalid(format!( + "semantic TMA view from `{}` has no copy or store consumer", + Operation::get_opid(producer, ctx) + ))); + } + for r#use in uses { + let user = r#use.user_op(); + let user_id = Operation::get_opid(user, ctx); + let index = r#use.find_index(ctx); + let allowed = if gmem { + (user_id == CuteTmaCopy2dOp::get_opid_static() && index == 0) + || (user_id == CuteTmaStore2dSemanticOp::get_opid_static() && index == 1) + } else { + (user_id == CuteTmaCopy2dOp::get_opid_static() && index == 1) + || (user_id == CuteTmaStore2dSemanticOp::get_opid_static() && index == 0) + }; + if !allowed { + return Err(invalid(format!( + "semantic TMA view from `{}` has unsupported use by `{user_id}` operand {index}", + Operation::get_opid(producer, ctx) + ))); + } + } + Ok(()) +} + +/// Check direct global/shared view pairing and reject unused or escaping transport views. +pub(super) fn verify_tma_story(ctx: &Context, all_ops: &[Ptr]) -> Result { + for operation in all_ops { + let opid = Operation::get_opid(*operation, ctx); + if opid == CuteTmaGmemViewOp::get_opid_static() { + verify_tma_view_uses(ctx, *operation, true)?; + } else if opid == CuteTmaSmemViewOp::get_opid_static() { + verify_tma_view_uses(ctx, *operation, false)?; + } else if opid == CuteTmaCopy2dOp::get_opid_static() { + let copy = CuteTmaCopy2dOp::wrap(*operation); + direct_tma_view_producer( + ctx, + copy.source(ctx), + &CuteTmaGmemViewOp::get_opid_static(), + "load source", + all_ops, + )?; + direct_tma_view_producer( + ctx, + copy.destination(ctx), + &CuteTmaSmemViewOp::get_opid_static(), + "load destination", + all_ops, + )?; + let barrier = copy.completion_barrier(ctx); + let Some(expect) = barrier.defining_op() else { + return Err(invalid( + "cute.tma_copy_2d completion barrier must come directly from pipeline_producer_expect_tx", + )); + }; + if Operation::get_opid(expect, ctx) != CutePipelineProducerExpectTxOp::get_opid_static() + || !all_ops.contains(&expect) + || CutePipelineProducerExpectTxOp::wrap(expect).completion_barrier(ctx) != barrier + { + return Err(invalid( + "cute.tma_copy_2d completion barrier is not the result of an expect_tx in this semantic scope", + )); + } + } else if opid == CuteTmaStore2dSemanticOp::get_opid_static() { + let store = CuteTmaStore2dSemanticOp::wrap(*operation); + direct_tma_view_producer( + ctx, + store.source(ctx), + &CuteTmaSmemViewOp::get_opid_static(), + "store source", + all_ops, + )?; + direct_tma_view_producer( + ctx, + store.destination(ctx), + &CuteTmaGmemViewOp::get_opid_static(), + "store destination", + all_ops, + )?; + } + } + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/tests/global_store.rs b/cuda-oxide/cute-rs-preview/dialect-cute/tests/global_store.rs new file mode 100644 index 0000000000..595d497519 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/tests/global_store.rs @@ -0,0 +1,56 @@ +//! Regression tests for CuTe cache-aware global-store width and pointer contracts. +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dialect_cute::ops::copy::CuteStoreGlobalNoAllocateOp; +use dialect_mir::types::{MirPtrType, address_space}; +use pliron::{ + basic_block::BasicBlock, + builtin::types::{IntegerType, Signedness}, + common_traits::Verify, + context::Context, + r#type::TypeHandle, +}; + +/// Cache-aware stores require supported vector widths, matching storage, and a writable global destination. +#[test] +fn global_store_checks_width_storage_and_address_roles() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let word: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let half: TypeHandle = IntegerType::get(&ctx, 16, Signedness::Unsigned).into(); + let mut ptr = |elem, mutable, space| -> TypeHandle { + MirPtrType::get(&mut ctx, elem, mutable, space).into() + }; + let types = vec![ + ptr(word, true, address_space::GLOBAL), + ptr(word, false, address_space::LOCAL), + ptr(word, false, address_space::GLOBAL), + ptr(word, true, address_space::SHARED), + ptr(half, false, address_space::LOCAL), + ]; + let block = BasicBlock::new(&mut ctx, None, types); + let a: Vec<_> = block.deref(&ctx).arguments().collect(); + for n in [1, 2, 4, 8] { + assert!( + CuteStoreGlobalNoAllocateOp::new(&mut ctx, a[0], a[1], n) + .verify(&ctx) + .is_ok() + ); + } + for n in [0, 3, 16, u64::MAX] { + assert!( + CuteStoreGlobalNoAllocateOp::new(&mut ctx, a[0], a[1], n) + .verify(&ctx) + .is_err() + ); + } + for (dst, src) in [(a[2], a[1]), (a[3], a[1]), (a[0], a[2]), (a[0], a[4])] { + assert!( + CuteStoreGlobalNoAllocateOp::new(&mut ctx, dst, src, 8) + .verify(&ctx) + .is_err() + ); + } +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/tests/public_paths.rs b/cuda-oxide/cute-rs-preview/dialect-cute/tests/public_paths.rs new file mode 100644 index 0000000000..0929506ead --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/tests/public_paths.rs @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Public operation families keep their compatibility imports as identical types. + +use std::any::TypeId; + +use dialect_cute::{ + copy_ops, epilogue_ops, gemm_tma_ops, gemv_ops, ops, pipeline_ops, scheduler_ops, + sm100_attention_ops, sm100_ops, smem_mma_ops, tensor_ops, +}; + +#[test] +fn operation_family_aliases_preserve_type_identity() { + macro_rules! same_type { + ($old:path, $grouped:path) => { + assert_eq!(TypeId::of::<$old>(), TypeId::of::<$grouped>()); + }; + } + + same_type!( + copy_ops::CuteStoreGlobalNoAllocateOp, + ops::copy::CuteStoreGlobalNoAllocateOp + ); + same_type!(tensor_ops::CuteTensorMakeOp, ops::tensor::CuteTensorMakeOp); + same_type!(gemv_ops::CuteDotOp, ops::gemv::CuteDotOp); + same_type!( + gemm_tma_ops::CuteTmaCopy2dOp, + ops::gemm_tma::CuteTmaCopy2dOp + ); + same_type!(smem_mma_ops::CuteTiledGemmOp, ops::mma::CuteTiledGemmOp); + same_type!( + epilogue_ops::CuteEpilogueHalfOp, + ops::epilogue::CuteEpilogueHalfOp + ); + same_type!( + pipeline_ops::CutePipelineStateNewOp, + ops::pipeline::CutePipelineStateNewOp + ); + same_type!( + scheduler_ops::CuteSchedulerNew1dOp, + ops::scheduler::CuteSchedulerNew1dOp + ); + same_type!( + sm100_ops::CuteSm100TiledMmaOp, + ops::sm100::CuteSm100TiledMmaOp + ); + same_type!( + sm100_attention_ops::CuteSm100MmaSsOp, + ops::sm100_attention::CuteSm100MmaSsOp + ); + same_type!(ops::CuteCopyOp, ops::copy::CuteCopyOp); + same_type!(ops::CuteCopyG2SOp, ops::copy::CuteCopyG2SOp); + same_type!(ops::CuteLdmatrixOp, ops::mma::CuteLdmatrixOp); +} diff --git a/cuda-oxide/cute-rs-preview/dialect-cute/tests/sm100_attention.rs b/cuda-oxide/cute-rs-preview/dialect-cute/tests/sm100_attention.rs new file mode 100644 index 0000000000..810fc09ebe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/tests/sm100_attention.rs @@ -0,0 +1,175 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Regression tests for CuTe independent SM100 MMA and TMA contracts. + +use dialect_cute::ops::sm100_attention::{CuteSm100MmaAttr, shared_layout}; +use pliron::{common_traits::Verify, context::Context}; +/// Construct the group-two M128/N128/K64 query-key plan used by attention tests. +fn qk() -> CuteSm100MmaAttr { + CuteSm100MmaAttr { + m: 128, + n: 128, + k: 64, + cta_group: 2, + a_mn: false, + b_mn: false, + a_layout: shared_layout(64, 64, false).unwrap(), + b_layout: shared_layout(64, 64, false).unwrap(), + } +} +/// QK and PV may use distinct canonical layouts without a kernel-wide MMA shape constraint. +#[test] +fn independent_qk_pv_layouts_verify() { + let ctx = Context::new(); + assert!(qk().verify(&ctx).is_ok()); + let pv = CuteSm100MmaAttr { + m: 128, + n: 256, + k: 32, + cta_group: 2, + a_mn: false, + b_mn: true, + a_layout: shared_layout(64, 32, false).unwrap(), + b_layout: shared_layout(128, 32, true).unwrap(), + }; + assert!(pv.verify(&ctx).is_ok()); +} +/// Changing collective size, matrix extents, or major order must invalidate the matching layout proof. +#[test] +fn rejects_wrong_group_shape_and_major() { + let ctx = Context::new(); + let mut p = qk(); + p.cta_group = 3; + assert!(p.verify(&ctx).is_err()); + let mut p = qk(); + p.n = 256; + assert!(p.verify(&ctx).is_err()); + let mut p = qk(); + p.b_mn = true; + assert!(p.verify(&ctx).is_err()); +} +/// Canonical tiled layouts reject incomplete swizzle spans and unsupported outer extents. +#[test] +fn rejects_invalid_shared_layout_extents() { + assert!(shared_layout(63, 64, false).is_err()); + assert!(shared_layout(64, 48, false).is_err()); + assert!(shared_layout(0, 64, false).is_err()); +} + +/// A shape-compatible tile is insufficient when its swizzle or physical ordering differs from the canonical layout. +#[test] +fn tma_rejects_noncanonical_layout() { + use dialect_cute::attributes::CuteComposedLayoutAttr; + use dialect_cute::layout::{ComposedLayout, OffsetUnit, Swizzle}; + use dialect_cute::ops::sm100_attention::CuteSm100TmaAttr; + let ctx = Context::new(); + for mn in [false, true] { + assert!( + CuteSm100TmaAttr { + layout: shared_layout(128, 128, mn).unwrap(), + cta_group: 1 + } + .verify(&ctx) + .is_ok() + ); + } + let layout = CuteComposedLayoutAttr( + ComposedLayout::new( + Swizzle::new(3, 3, 3), + 0, + "(128,128):(128,1)".parse().unwrap(), + OffsetUnit::Elements, + ) + .unwrap(), + ); + assert!( + CuteSm100TmaAttr { + layout, + cta_group: 1 + } + .verify(&ctx) + .is_err() + ); +} + +/// A global-to-shared transfer must retain writable destination capability. +#[test] +fn tma_load_rejects_readonly_shared_destination() { + use dialect_cute::ops::sm100_attention::{CuteSm100TmaAttr, CuteSm100TmaLoadOp}; + use dialect_mir::types::{MirFP16Type, MirPtrType, address_space}; + use pliron::{ + basic_block::BasicBlock, + builtin::types::{IntegerType, Signedness}, + r#type::TypeHandle, + }; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let half: TypeHandle = MirFP16Type::get(&ctx).into(); + let u32t: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let i32t: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signed).into(); + let u64t: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let desc = MirPtrType::get(&mut ctx, u64t, false, address_space::GENERIC).into(); + let barrier = MirPtrType::get(&mut ctx, u64t, true, address_space::SHARED).into(); + for mutable in [false, true] { + let destination = MirPtrType::get(&mut ctx, half, mutable, address_space::SHARED).into(); + let block = BasicBlock::new( + &mut ctx, + None, + vec![desc, destination, i32t, i32t, i32t, barrier, u32t], + ); + let arguments = block.deref(&ctx).arguments().collect(); + let op = CuteSm100TmaLoadOp::new( + &mut ctx, + arguments, + CuteSm100TmaAttr { + layout: shared_layout(128, 128, false).unwrap(), + cta_group: 1, + }, + ); + assert_eq!(op.verify(&ctx).is_ok(), mutable); + } +} + +/// MMA cannot silently omit byte offsets or accept offsets of the wrong scalar width. +#[test] +fn mma_tile_offsets_require_explicit_u32_operands() { + use dialect_cute::ops::sm100_attention::{CuteSm100MmaSsOp, CuteSm100MmaTsOp}; + use dialect_mir::types::{MirFP16Type, MirPtrType}; + use pliron::{ + basic_block::BasicBlock, + builtin::types::{IntegerType, Signedness}, + r#type::TypeHandle, + }; + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + dialect_cute::register(&mut ctx); + let half: TypeHandle = MirFP16Type::get(&ctx).into(); + let u32t: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let u64t: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let boolt: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get(&mut ctx, half, true, 3).into(); + let block = BasicBlock::new( + &mut ctx, + None, + vec![u32t, shared, shared, boolt, u32t, u64t], + ); + let a = block.deref(&ctx).arguments().collect::>(); + for (offset, valid) in [(a[4], true), (a[5], false), (a[1], false)] { + let ss = CuteSm100MmaSsOp::new(&mut ctx, vec![a[0], a[1], a[2], a[3], offset, a[4]], qk()); + assert_eq!(ss.verify(&ctx).is_ok(), valid); + let ss = CuteSm100MmaSsOp::new(&mut ctx, vec![a[0], a[1], a[2], a[3], a[4], offset], qk()); + assert_eq!(ss.verify(&ctx).is_ok(), valid); + let mut plan = qk(); + plan.cta_group = 1; + plan.a_layout = shared_layout(128, 64, false).unwrap(); + plan.b_layout = plan.a_layout.clone(); + let ts = CuteSm100MmaTsOp::new(&mut ctx, vec![a[0], a[4], a[2], a[3], offset], plan); + assert_eq!(ts.verify(&ctx).is_ok(), valid); + } + let missing = CuteSm100MmaSsOp::new(&mut ctx, vec![a[0], a[1], a[2], a[3]], qk()); + assert!(missing.verify(&ctx).is_err()); +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock similarity index 95% rename from cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock rename to cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock index c46db5e9e1..e45a7d5d5b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock @@ -55,6 +55,37 @@ dependencies = [ "generic-array", ] +[[package]] +name = "blockscale_gemm_ctm" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "serde_json", +] + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -115,6 +146,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "thiserror", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -128,7 +168,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +180,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +204,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +215,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +236,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -227,15 +276,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "error_kernel_shared_ptr_param" -version = "0.1.0" -dependencies = [ - "cuda-core", - "cuda-device", - "cuda-host", -] - [[package]] name = "foldhash" version = "0.1.5" @@ -264,6 +304,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.toml new file mode 100644 index 0000000000..a02e446e25 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "blockscale_gemm_ctm" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "SM120 MXFP4 GEMM using typed Rust instruction primitives" +[workspace] +[dependencies] +ctm-rs = { path = "../../ctm-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } +serde_json = "1" diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/README.md b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/README.md new file mode 100644 index 0000000000..e1899793a2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/README.md @@ -0,0 +1,65 @@ +# MXFP4 GEMM with Rust CTM (ctm-rs) + +This cuda-oxide example multiplies packed E2M1 matrices with UE8M0 scales per +32 K values, FP32 accumulation and FP16 output. The fixed problem is +**M=N=2048, K=1024** on an **SM120 Blackwell GPU**, such as RTX 5090. + +```text +A/B + scales → TMA shared-memory ring → register fragments + → block-scaled MMA → FP16 epilogue → TMA output +``` + +The 128×128 tiles use three pipeline stages, eight compute warps and one +producer warp. Typed register primitives keep operand roles explicit; +physical lane addresses, barriers and publication remain explicit Rust. +The Rust executable needs no Python runtime. + +From `cuda-oxide/`, with the pinned CUTLASS 4.7 compiler installed using +`cargo oxide toolchain install cutlass`: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ +CUDA_OXIDE_MLIR_PROFILE=cutlass-primitives-mlir22 \ +CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ + cargo oxide run blockscale_gemm_ctm --arch sm_120a -- \ + --ctas 170 --warmup 100 --samples 31 +``` + +Success checks all **4,194,304 FP16 output bit patterns** against a substantial +CPU oracle, checks repeated-launch output, and reports CUDA-event timings. +This runs Rust only. The [Python CTM reference](reference/README.md) runs the +same GEMM through CUTLASS Primitives and needs its Python dependencies; +`compare.py` beside this README times both under Nsight Systems. `validate.py` +covers 170-, 1- and 17-CTA grids plus memory/tensor and synchronization +sanitizers. From `cuda-oxide/`: + +```sh +python3 scripts/check-primitives-api-compile.py --example blockscale \ + --output-dir /tmp/primitives-blockscale-evidence +python3 cute-rs-preview/examples/blockscale_gemm_ctm/validate.py \ + --binary /path/to/blockscale_gemm_ctm --output-dir /tmp/blockscale-results --sanitizer +``` + +The first command builds a fresh executable through the CUTLASS compiler and +records its manifest; use its `run-*/target/release/blockscale_gemm_ctm` as the +binary. + +## Rust CTM versus Python CTM + +RTX 5090 at `a0d5776b2` (2026-10-18), same inputs and 170 CTAs. Rust uses +CUTLASS compiler **4.7.0**; Python uses **4.8.0**; both use optimization level 2. +Both pass all output-bit and replay checks. Times are microseconds: median of +three round medians, with 100 warmups and 31 single-launch samples per round. + +| Implementation | CUDA events (µs) | Nsight kernel duration (µs) | +| --- | ---: | ---: | +| Rust CTM (`ctm-rs`) | 13.536 | 11.616 | +| Python CTM (CUTLASS Primitives) | 19.520 | 10.880 | + +Events include submission gaps; Nsight reports GEMM execution duration, +including device stalls. Python rebuilds tensor-map descriptors per launch; +Rust prepares them before timing. The two metrics come from separate runs, +with rotating order and unlocked clocks. + +- [Python CTM API and validation](reference/README.md) +- [CuTe-RS counterpart](../blockscale_gemm_cute) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/compare.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/compare.py new file mode 100644 index 0000000000..9f94b2da92 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/compare.py @@ -0,0 +1,285 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compare bit-exact GEMM binaries using serial Nsight Systems kernel-active times. + +Each binary must accept --ctas, --warmup, --samples and --json, and launch only +its exact configured kernel name (default: blockscale_gemm): one correctness +launch, then warmups, then measured samples. +CUDA event times remain in the original executable reports and are not used here. +""" + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import re +import signal +import sqlite3 +import subprocess +import tempfile +import time + + +def validate_report(path, ctas, warmup, samples, started_ns): + """Reject stale evidence and require both initial and replay correctness.""" + try: + if path.stat().st_mtime_ns < started_ns: + raise ValueError(f"stale correctness report: {path}") + result = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"cannot read correctness report {path}: {error}") from error + if not isinstance(result, dict): + raise ValueError("correctness report must be a JSON object") + expected = { + "verification": "bit-exact-pass", "checked_outputs": 4194304, + "mnk": [2048, 2048, 1024], "compute_capability": [12, 0], + "ctas": ctas, "warmup": warmup, "samples": samples, + "replay_verified": True, + } + for key, value in expected.items(): + actual = result.get(key) + if (type(actual) is not type(value) or actual != value + or isinstance(actual, list) and any(type(x) is not int for x in actual)): + raise ValueError(f"invalid or missing {key}: expected {value!r}, got {actual!r}") + return result + + +def kernel_samples(path, warmup, samples, expected_name="blockscale_gemm"): + """Return the last N ordered kernel durations in microseconds, failing closed.""" + try: + with sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True) as db: + rows = db.execute( + "SELECT k.start, k.end, s.value " + "FROM CUPTI_ACTIVITY_KIND_KERNEL AS k " + "LEFT JOIN StringIds AS s ON s.id = k.shortName " + "ORDER BY k.start, k.end" + ).fetchall() + except sqlite3.Error as error: + raise ValueError(f"cannot read kernel trace {path}: {error}") from error + expected = 1 + warmup + samples + if len(rows) != expected: + raise ValueError(f"expected {expected} kernel launches, found {len(rows)}") + for start, end, name in rows: + if name != expected_name: + raise ValueError(f"unexpected kernel shortName: {name!r}") + if type(start) is not int or type(end) is not int or end <= start: + raise ValueError(f"invalid kernel timestamps: {start!r}, {end!r}") + return [(end - start) / 1000 for start, end, _ in rows[-samples:]] + + +def summarize_samples(samples): + """Match src/main.rs percentile: rounded index into sorted samples.""" + ordered = sorted(samples) + if not ordered or any(not math.isfinite(x) or x <= 0 for x in ordered): + raise ValueError("kernel durations must be finite and positive") + + def percentile(numerator, denominator): + return ordered[((len(ordered) - 1) * numerator + denominator // 2) // denominator] + + return {"median_us": percentile(1, 2), "p10_us": percentile(1, 10), + "p90_us": percentile(9, 10)} + + +def aggregate_rounds(entries): + summary = {} + for entry in entries: + result = summary.setdefault(entry["label"], {"round_medians_us": [], "samples_us": []}) + result["round_medians_us"].append(entry["median_us"]) + result["samples_us"].extend(entry["samples_us"]) + for result in summary.values(): + result["median_of_round_medians_us"] = summarize_samples(result["round_medians_us"])["median_us"] + return summary + + +def execution_order(labels, rounds): + return [labels[r % len(labels):] + labels[:r % len(labels)] for r in range(rounds)] + + +def run_command(command, timeout, environment, log_path): + """Wait for each process; kill the whole process group on timeout or interrupt.""" + with log_path.open("w") as log: + with subprocess.Popen(command, env=environment, stdout=log, + stderr=subprocess.STDOUT, start_new_session=True) as process: + try: + return process.wait(timeout=timeout) + except (subprocess.TimeoutExpired, KeyboardInterrupt): + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + raise + + +def parse_arguments(argv): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", action="append", required=True, metavar="LABEL=PATH", + help="repeat for at least two implementations with unique labels") + parser.add_argument("--kernel-name", action="append", default=[], metavar="LABEL=NAME", + help="exact Nsight kernel shortName for a binary label (default: blockscale_gemm)") + parser.add_argument("--output-dir", type=Path, required=True, + help="parent for a fresh run directory containing raw traces and reports") + parser.add_argument("--device", default="1", help="single physical CUDA device index or UUID (default: 1)") + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--samples", type=int, default=31) + parser.add_argument("--ctas", type=int, default=170) + parser.add_argument("--timeout", type=float, default=240, + help="maximum seconds for each tool invocation (default: 240)") + args = parser.parse_args(argv) + if args.rounds <= 0 or args.samples <= 0 or args.warmup < 0: + parser.error("rounds and samples must be positive; warmup must be nonnegative") + if not 1 <= args.ctas <= 256: + parser.error("ctas must be between 1 and 256") + if not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("timeout must be finite and positive") + if not args.device or "," in args.device: + parser.error("device must identify exactly one GPU") + binaries = {} + for specification in args.binary: + label, separator, filename = specification.partition("=") + if not separator or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", label): + parser.error("binary must be LABEL=PATH; labels use letters, digits, underscores or hyphens") + if label in binaries: + parser.error(f"duplicate binary label: {label}") + try: + path = Path(filename).resolve(strict=True) + except OSError as error: + parser.error(f"invalid binary {filename!r}: {error}") + if not path.is_file() or not os.access(path, os.X_OK): + parser.error(f"binary must be an executable file: {path}") + binaries[label] = path + if len(binaries) < 2: + parser.error("provide at least two --binary LABEL=PATH arguments") + kernel_names = {} + for specification in args.kernel_name: + label, separator, name = specification.partition("=") + if not separator or not name.strip(): + parser.error("kernel-name must be LABEL=NAME with a nonempty name") + if label not in binaries: + parser.error(f"unknown kernel-name label: {label}") + if label in kernel_names: + parser.error(f"duplicate kernel-name label: {label}") + kernel_names[label] = name + args.kernel_names = {label: kernel_names.get(label, "blockscale_gemm") for label in binaries} + return args, binaries + + +def main(argv=None): + args, binaries = parse_arguments(argv) + args.output_dir.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="run-", dir=args.output_dir.resolve())) + environment = os.environ.copy() + environment.update(CUDA_VISIBLE_DEVICES=args.device, CUDA_DEVICE_ORDER="PCI_BUS_ID") + report = { + "status": "running", "metric": "Nsight Systems kernel-active duration", + "units": "microseconds", "kernel": "blockscale_gemm", "started_ns": time.time_ns(), + "kernel_names": args.kernel_names, "kernel_name_field": "shortName", + "parameters": {"device": args.device, "rounds": args.rounds, "warmup": args.warmup, + "samples": args.samples, "ctas": args.ctas, "timeout_seconds": args.timeout}, + "environment": {key: environment[key] for key in ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER")}, + "percentile_method": "sorted[((count - 1) * numerator + denominator // 2) // denominator]", + "binaries": {}, "execution_order": execution_order(list(binaries), args.rounds), + "commands": [], "runs": [], + } + + def save(): + temporary = directory / "report.json.tmp" + temporary.write_text(json.dumps(report, indent=2) + "\n") + temporary.replace(directory / "report.json") + + def execute(command, log_path): + entry = {"command": command, "log": str(log_path), "status": "running", + "started_ns": time.time_ns()} + report["commands"].append(entry) + save() + try: + entry["exit_code"] = run_command(command, args.timeout, environment, log_path) + if entry["exit_code"] != 0: + raise ValueError(f"command exited {entry['exit_code']}; see {log_path}") + entry["status"] = "passed" + except (OSError, ValueError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + entry.update(status="failed", error=str(error) or "interrupted") + raise + finally: + entry["finished_ns"] = time.time_ns() + save() + + save() + print(f"Evidence: {directory}", flush=True) + active = None + try: + for label, binary in binaries.items(): + report["binaries"][label] = { + "path": str(binary), "sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + } + version_log = directory / "nsys-version.log" + execute(["nsys", "--version"], version_log) + report["nsys_version"] = version_log.read_text().strip() + snapshot = directory / "nvidia-smi.log" + execute(["nvidia-smi", "--id", args.device, "-q"], snapshot) + report["gpu_snapshot"] = {"path": str(snapshot), "text": snapshot.read_text()} + for round_index, order in enumerate(report["execution_order"], start=1): + for label in order: + prefix = directory / f"round-{round_index:02d}-{label}" + result_path = prefix.with_suffix(".json") + trace_path = prefix.with_suffix(".nsys-rep") + database_path = prefix.with_suffix(".sqlite") + profile_log = prefix.with_suffix(".profile.log") + export_log = prefix.with_suffix(".export.log") + binary_command = [str(binaries[label]), "--ctas", str(args.ctas), + "--warmup", str(args.warmup), "--samples", str(args.samples), + "--json", str(result_path)] + profile_command = ["nsys", "profile", "--trace=cuda,nvtx", "--sample=none", + "--cpuctxsw=none", "--output", str(prefix), *binary_command] + export_command = ["nsys", "export", "--type=sqlite", "--output", + str(database_path), str(trace_path)] + active = {"round": round_index, "label": label, "status": "running", + "correctness_report": str(result_path), "raw_trace": str(trace_path), + "sqlite": str(database_path), "profile_log": str(profile_log), + "export_log": str(export_log), "profile_command": profile_command, + "export_command": export_command, "started_ns": time.time_ns()} + report["runs"].append(active) + save() + if any(path.exists() for path in (result_path, trace_path, database_path)): + raise ValueError("refusing to reuse existing benchmark evidence") + print(f"Round {round_index}/{args.rounds}: {label}", flush=True) + execute(profile_command, profile_log) + active["verification"] = validate_report( + result_path, args.ctas, args.warmup, args.samples, active["started_ns"]) + if not trace_path.is_file(): + raise ValueError(f"missing raw Nsight Systems trace: {trace_path}") + execute(export_command, export_log) + measured = kernel_samples(database_path, args.warmup, args.samples, + expected_name=args.kernel_names[label]) + active.update(status="passed", samples_us=measured, + kernel_launches=1 + args.warmup + args.samples, + finished_ns=time.time_ns(), **summarize_samples(measured)) + save() + print(f" kernel-active median={active['median_us']:.3f} us, " + f"p10={active['p10_us']:.3f} us, p90={active['p90_us']:.3f} us", flush=True) + active = None + report.update(status="passed", summary=aggregate_rounds(report["runs"])) + except (OSError, ValueError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + message = str(error) or "interrupted" + if active is not None: + active.update(status="failed", error=message, finished_ns=time.time_ns()) + report.update(status="failed", error=message) + print(f"FAIL: {message}", flush=True) + report["finished_ns"] = time.time_ns() + save() + if report["status"] != "passed": + print(f"Evidence: {directory}", flush=True) + return 1 + for label, summary in report["summary"].items(): + print(f"{label}: median of round medians " + f"{summary['median_of_round_medians_us']:.3f} us", flush=True) + print(f"Evidence: {directory}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/.gitignore b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/.gitignore new file mode 100644 index 0000000000..a348e5040c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/.gitignore @@ -0,0 +1 @@ +/__pycache__/ diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/README.md b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/README.md new file mode 100644 index 0000000000..045c3b6d40 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/README.md @@ -0,0 +1,86 @@ +# CUTLASS CTM blockscale GEMM reference + +This is a port of cuda-oxide's [Rust CTM GEMM](../) +using the actual `cutlass.experimental.primitives` API from CUTLASS 4.8.0. +The kernel is written in Python, as exposed by that upstream API; the Rust +implementation remains in `blockscale_gemm_ctm`. The older [C++ reference](../../blockscale_gemm_cute/reference/cpp/) +reference uses custom CUDA C++ inline PTX. + +The entry README shows the Rust CTM versus Python CTM comparison, with both +CUDA-event and Nsight kernel timings on the same inputs; `compare.py` in the +example directory produces it. + +Both CTM kernels use M=N=2048, K=1024, packed E2M1 inputs, UE8M0 scales per +32 K elements, FP32 accumulation, and FP16 output. Physical offsets, the +128×128×128 tile, three pipeline stages, eight consumer warps plus one producer, +persistent tile traversal, and the STMATRIX/TMA epilogue match. Each launch has +288 threads, five grid-constant tensor maps, 86,016 dynamic shared bytes and +1,024 static shared bytes for alignment. + +[`kernel.py`](kernel.py) calls upstream TMA, barriers, LDMATRIX, block-scaled MMA, +STMATRIX, and conversion operations. Its small MMA adapter packs/unpacks the +upstream API's raw LLVM register carriers. CuTe supplies decorators, tracing, +and launch/allocation infrastructure; there is no CuTe tensor/layout algebra +in this programming model. Compiler IR still contains CuTe launch/allocation +plumbing, so this is not a claim that every intermediate `cute.*` operation +has disappeared. + +## Run and reproduce + +From this directory, with Python 3.12 and an SM120 GPU: + +```sh +python3.12 -m venv /tmp/cutlass-ctm-venv +/tmp/cutlass-ctm-venv/bin/python -m pip install -r requirements.txt +/tmp/cutlass-ctm-venv/bin/python -m unittest discover +CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \ + /tmp/cutlass-ctm-venv/bin/python benchmark.py --ctas 170 \ + --warmup 100 --samples 31 --json /tmp/cutlass-ctm-result.json +``` + +The harness creates all inputs on the CPU, poisons output, checks every FP16 +bit before timing, and checks replay afterward. It launches only GEMM kernels. +The independent original Rust fixture generator and this NumPy implementation +were compared byte for byte for all five complete buffers, including all +4,194,304 expected outputs. CLI JSON includes fixture/source/package hashes. + +For fresh **Nsight kernel-active** measurements, build Rust with +`scripts/check-primitives-api-compile.py --example blockscale`, then use: + +```sh +PATH=/tmp/cutlass-ctm-venv/bin:$PATH \ +CUTE_DSL_KEEP=ir,ptx,cubin CUTE_DSL_DUMP_DIR=/tmp/cutlass-ctm-artifacts \ +python3 ../compare.py \ + --binary rust=/absolute/path/to/blockscale_gemm_ctm \ + --binary cutlass="$PWD/benchmark.py" \ + --kernel-name 'cutlass=kernel_cutlass_blockscale_gemm_1...' \ + --output-dir /tmp/blockscale-ctm-comparison --device 1 \ + --ctas 170 --rounds 3 --warmup 100 --samples 31 +``` + +The ellipsis is part of the literal Nsight 2026.3 `shortName`, not a wildcard. +Other profiler/compiler versions may name it differently; inspect their trace +and supply its exact value. Every process must contain exactly 132 launches of +its expected kernel, plus a fresh passing correctness report. Runs are serial +and rotate order; timing uses `(end-start)/1000` from CUPTI kernel activities. +CUDA event intervals in executable JSON include host submission gaps and are +a separate metric. The Python host rebuilds descriptors on each invocation; +that host work is outside the kernel-active measurement. + +Use `../validate.py --binary "$PWD/benchmark.py"` +with the same Python PATH, GPU selection, `--output-dir`, and `--sanitizer` to +check 170-, 1-, and 17-CTA grids, memcheck with tensor checks, and synccheck. + +## Scope and provenance + +The official 4.8.0 wheel's primitives wrapper matches the local CUTLASS checkout +`0b55a2f691d69981583568fd9eb69687b1f0de8a` byte for byte (SHA-256 +`e11a53bfe7dc6f27ea80d189ae84996147f153660d6d97ac1d8fcaede66c7499`). +This GEMM port is repository-authored, rather than an upstream tuned example. + +The current Rust compiler is pinned to CUTLASS 4.7.0; this primitives API needs +4.8.0. Upstream barrier waits also emit an explicit 10,000,000-nanosecond +maximum suspension hint, while Rust omits that optional operand. Phase and +ordering semantics match; emitted instructions and compiler optimizations can +differ. Results compare these implementations/toolchains and do not isolate +language overhead or establish SOL. No global compiler pin was changed. diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/benchmark.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/benchmark.py new file mode 100755 index 0000000000..ffb4fa2a9b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/benchmark.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Actual CUTLASS primitives GEMM, with the Rust CTM fixture and timing contract. + +Use compare.py for Nsight kernel-active timings. The CUDA event measurements +reported here include host launch overhead and are a separate metric. +""" +import argparse +from contextlib import ExitStack +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path + +import numpy as np + + +def parse_arguments(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--ctas", type=int, default=170) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--samples", type=int, default=31) + parser.add_argument("--json", type=Path) + args = parser.parse_args(argv) + if not 1 <= args.ctas <= 256: + parser.error("ctas must be between 1 and 256") + if args.warmup < 0 or args.samples <= 0: + parser.error("warmup must be nonnegative and samples must be positive") + return args + + +def verify_bits(got, expected): + if (got.shape != expected.shape or got.dtype != np.uint16 + or expected.dtype != np.uint16): + raise ValueError("output shape/storage differs from the FP16 bit oracle") + mismatches = np.flatnonzero(got.reshape(-1) != expected.reshape(-1)) + if mismatches.size: + first = int(mismatches[0]) + raise ValueError(f"{mismatches.size} FP16 bit mismatches; first at {first}: " + f"got {got.flat[first]:#06x}, expected {expected.flat[first]:#06x}") + + +def main(argv=None): + args = parse_arguments(argv) + if args.json: + args.json.unlink(missing_ok=True) + # Select the architecture before initializing the DSL. No torch dependency + # or helper GPU kernels: the trace contains only the GEMM under test. + os.environ["CUTE_DSL_ARCH"] = "sm_120a" + import cuda.bindings.driver as driver + import cutlass + import cutlass.cute as cute + from fixture import build_fixture + from launch import host + + def checked(result): + if result[0] != driver.CUresult.CUDA_SUCCESS: + raise RuntimeError(str(result)) + return result[1] if len(result) == 2 else result[1:] + + checked(driver.cuInit(0)) + device = checked(driver.cuDeviceGet(0)) + capability = list(checked(driver.cuDeviceComputeCapability(device))) + if capability != [12, 0]: + raise RuntimeError(f"SM120 required, got {capability}") + context = checked(driver.cuDevicePrimaryCtxRetain(device)) + with ExitStack() as resources: + resources.callback(lambda: checked(driver.cuDevicePrimaryCtxRelease(device))) + checked(driver.cuCtxSetCurrent(context)) + stream = checked(driver.cuStreamCreate(0)) + resources.callback(lambda: checked(driver.cuStreamDestroy(stream))) + allocations = [] + fixture = build_fixture() + expected = fixture["expected"] + output = np.full(expected.shape, 0x7e00, dtype=np.uint16) + pointers = [] + for data in [fixture[key] for key in ("a", "b", "sa", "sb")] + [output]: + pointer = checked(driver.cuMemAlloc(data.nbytes)) + resources.callback(lambda p=pointer: checked(driver.cuMemFree(p))) + allocations.append(pointer) + checked(driver.cuMemcpyHtoD(pointer, data.ctypes.data, data.nbytes)) + pointers.append(cutlass.Int64(int(pointer))) + parameters = (*pointers, cutlass.Int32(args.ctas), stream) + # ExitStack attempts every release even if a driver cleanup call fails. + resources.callback(lambda: checked(driver.cuStreamSynchronize(stream))) + compiled = cute.compile(host, *parameters) + compiled(*parameters) + checked(driver.cuStreamSynchronize(stream)) + checked(driver.cuMemcpyDtoH(output.ctypes.data, allocations[-1], output.nbytes)) + verify_bits(output, expected) + print(f"PASS: all {expected.size:,} FP16 output bits; {args.ctas} CTAs", flush=True) + for _ in range(args.warmup): + compiled(*parameters) + checked(driver.cuStreamSynchronize(stream)) + times = [] + for _ in range(args.samples): + with ExitStack() as events: + start = checked(driver.cuEventCreate(0)) + events.callback(lambda: checked(driver.cuEventDestroy(start))) + end = checked(driver.cuEventCreate(0)) + events.callback(lambda: checked(driver.cuEventDestroy(end))) + checked(driver.cuEventRecord(start, stream)) + compiled(*parameters) + checked(driver.cuEventRecord(end, stream)) + checked(driver.cuEventSynchronize(end)) + times.append(float(checked(driver.cuEventElapsedTime(start, end)))) + checked(driver.cuMemcpyDtoH(output.ctypes.data, allocations[-1], output.nbytes)) + verify_bits(output, expected) + ordered = sorted(times) + median = ordered[len(ordered) // 2] + report = dict(implementation="CUTLASS experimental primitives", + verification="bit-exact-pass", checked_outputs=int(expected.size), + mnk=[2048, 2048, 1024], compute_capability=capability, + ctas=args.ctas, warmup=args.warmup, samples=args.samples, + replay_verified=True, median_ms=median, + event_samples_ms=times, + event_metric="CUDA events, includes host submission intervals", + packages={name: importlib.metadata.version(name) for name in + ("nvidia-cutlass-dsl", "nvidia-cutlass-dsl-libs-cu13", + "cuda-bindings", "numpy")}, + sources={p.name: hashlib.sha256(p.read_bytes()).hexdigest() + for p in Path(__file__).parent.glob("*.py")}, + fixture_sha256={key: hashlib.sha256(data.tobytes()).hexdigest() + for key, data in fixture.items()}, + gpu_visible_device=os.getenv("CUDA_VISIBLE_DEVICES")) + from cutlass.experimental.primitives import nvvm_wrapper + upstream = Path(nvvm_wrapper.__file__) + report["upstream_primitives"] = dict( + path=str(upstream), sha256=hashlib.sha256(upstream.read_bytes()).hexdigest()) + cubin = compiled.__cubin__ + ptx = compiled.__ptx__ + if cubin is not None: + report["cubin_sha256"] = hashlib.sha256(cubin).hexdigest() + if ptx is not None: + report["ptx_sha256"] = hashlib.sha256(ptx.encode()).hexdigest() + # Publish success only after every CUDA resource was released successfully. + if args.json: + args.json.write_text(json.dumps(report, indent=2) + "\n") + print(f"PASS: replay exact; CUDA event median {median * 1000:.3f} us", flush=True) + + +if __name__ == "__main__": + main() diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/fixture.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/fixture.py new file mode 100644 index 0000000000..d9a3458452 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/fixture.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU fixture shared with the Rust CTM block-scaled GEMM example. + +No GPU framework is used here. Packed inputs match +``blockscale_gemm_ctm/src/main.rs`` and its ``physical::scale_offset`` layout. +""" + +import numpy as np + +M = N = 2048 +K = 1024 +KB = K // 32 + + +def _mix32(value): + """Apply the Rust fixture's wrapping u32 hash elementwise.""" + value = np.asarray(value, dtype=np.uint32) + value = value ^ (value >> 16) + value = np.multiply(value, np.uint32(0x7FEB352D), dtype=np.uint32) + value = value ^ (value >> 15) + value = np.multiply(value, np.uint32(0x846CA68B), dtype=np.uint32) + return value ^ (value >> 16) + + +def _nibbles(rows, multiplier, seed): + row = np.arange(rows, dtype=np.uint32)[:, None] + k = np.arange(K, dtype=np.uint32)[None, :] + return (_mix32(row * np.uint32(multiplier) ^ k ^ np.uint32(seed)) & 15).astype( + np.uint8 + ) + + +def _scales(rows, multiplier, seed): + row = np.arange(rows, dtype=np.uint32)[:, None] + group = np.arange(KB, dtype=np.uint32)[None, :] + return ( + 126 + _mix32(row * np.uint32(multiplier) ^ group ^ np.uint32(seed)) % 3 + ).astype(np.uint8) + + +def _pack_nibbles(nibbles): + # Consecutive low/high nibbles form a byte; pairs of bytes reproduce Rust's + # little-endian f16 transport carriers without interpreting them as floats. + return nibbles[:, 0::2] | (nibbles[:, 1::2] << 4) + + +def _pack_scales(dense): + rows, groups = dense.shape + row = np.arange(rows)[:, None] + group = np.arange(groups)[None, :] + # Exact physical::scale_offset(row, group, groups), in bytes. + scale_word = (row & 31) * 4 + ((row >> 5) & 3) + offset = ( + (row // 128) * (groups // 4) * 512 + + (group // 4) * 512 + + scale_word * 4 + + group % 4 + ) + packed = np.empty(rows * groups, dtype=np.uint8) + packed[offset] = dense + words = packed[0::2].astype(np.uint16) | (packed[1::2].astype(np.uint16) << 8) + return words.reshape((rows // 128) * (groups // 4), 256) + + +def _scaled_values(nibbles, scales): + positive = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=np.float32) + values = np.concatenate((positive, -positive))[nibbles] + powers = np.exp2(scales.astype(np.float32) - 127) + values.reshape(nibbles.shape[0], KB, 32)[:] *= powers[:, :, None] + return values + + +def build_fixture(): + """Return C-contiguous packed A/B, scale atoms, and exact FP16 result bits. + + ``a`` and ``b`` are uint8 arrays of shape (2048, 512), with B's rows + representing output columns. ``sa`` and ``sb`` are uint16 arrays of shape + (128, 256), each row holding one 512-byte canonical scale atom. ``expected`` + is a uint16 (2048, 2048) matrix containing FP16 bit patterns. + + CPU float32 matmul is exact for this fixture: scaled e2m1 values are + multiples of 1/4 with magnitude at most 12. Their products are multiples + of 1/16, and even the sum of absolute products, 1024 * 144, needs fewer + than 24 significant bits in those units. Thus every partial sum is exact, + including the Rust oracle's per-32-element grouping. Only the final FP16 + conversion rounds, using round-to-nearest, ties-to-even. + """ + a = _nibbles(M, 0x9E3779B9, 0xA511E9B3) + b = _nibbles(N, 0x85EBCA6B, 0x63D83595) + sa = _scales(M, 0xC2B2AE35, 0) + sb = _scales(N, 0x27D4EB2F, 0x165667B1) + expected = (_scaled_values(a, sa) @ _scaled_values(b, sb).T).astype(np.float16) + return { + "a": _pack_nibbles(a), + "b": _pack_nibbles(b), + "sa": _pack_scales(sa), + "sb": _pack_scales(sb), + "expected": expected.view(np.uint16), + } diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/kernel.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/kernel.py new file mode 100644 index 0000000000..db80da80ba --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/kernel.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact physical/scheduling port of blockscale_gemm_ctm using CUTLASS CTM. + +CuTe supplies kernel tracing and launch only. All hardware primitives are the +upstream cutlass.experimental.primitives operations; no tensor/layout algebra +or replacement inline PTX is used. Integer offsets match physical.rs. +""" + +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass.experimental import primitives as prims +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass._mlir_helpers import dsl_user_op + +SMEM_BYTES = 86016 +STAGES = 3 +DATA_BYTES = 8192 +SCALE_BYTES = 512 +B_OFFSET = 24576 +SFA_OFFSET = 49152 +SFB_OFFSET = 50688 +BARRIERS = 52224 +OUTPUT = 53248 +TX_BYTES = 17408 +N_OFFSETS = (0, 8, 32, 40, 64, 72, 96, 104) + + +@dsl_user_op +def _mma(c, a, b0, b1, sa, sb, *, loc=None, ip=None): + """Adapt register vectors to the actual upstream raw-shaped CTM MMA API.""" + f32 = cutlass.Float32.mlir_type + result = prims.mma_block_scale( + llvm.StructType.get_literal([f32] * 4), + ir.Attribute.parse("#nvvm.shape"), + prims.ScaleVecSize.X2, + prims.BlockScaleFormat.UE8M0, + prims.MMABlockScaleKind.MXF4, + [a[i].ir_value(loc=loc, ip=ip) for i in range(4)], + [b0.ir_value(loc=loc, ip=ip), b1.ir_value(loc=loc, ip=ip)], + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)], + cutlass.Int32(sa).ir_value(loc=loc, ip=ip), + cutlass.Int16(0).ir_value(loc=loc, ip=ip), + cutlass.Int16(0).ir_value(loc=loc, ip=ip), + cutlass.Int32(sb).ir_value(loc=loc, ip=ip), + cutlass.Int16(0).ir_value(loc=loc, ip=ip), + cutlass.Int16(0).ir_value(loc=loc, ip=ip), + multiplicand_a_ptx_type=ir.Attribute.parse("#nvvm.mma_type"), + multiplicand_b_ptx_type=ir.Attribute.parse("#nvvm.mma_type"), + loc=loc, + ip=ip, + ) + return cutlass.Vector.from_elements( + tuple(cutlass.Float32(llvm.extractvalue(f32, result, [i], loc=loc, ip=ip)) + for i in range(4)), dtype=cutlass.Float32, loc=loc, ip=ip) + + +@cute.jit +def _scale_word(row): + return (row & 31) * 4 + ((row >> 5) & 3) + + +@cute.jit +def _matrix_row(atom, half, lane): + sub = lane // 8 + row = 16 * atom + 8 * (sub % 2) + lane % 8 + column = 16 * half + 8 * (sub // 2) + plain = row * 64 + column * 2 + return plain ^ ((plain & 0x180) >> 3) + + +@cute.jit +def _store(output, c, warp, lane, band: cutlass.Constexpr[int], + slot: cutlass.Constexpr[int]): + row = 16 * (warp % 4) + 64 * band + (lane & 15) + column = 16 * (warp // 4) + N_OFFSETS[slot] + plain = row * 64 + column % 64 + element = (column // 64) * 8192 + (plain ^ ((plain & 0x1C0) >> 3)) + # PTX packs its first input high; Rust's public helper takes low first. + top = prims.cvt_packfloat_f32(c[1], c[0], 0, prims.CVTPackFloat.F16X2, + rnd=prims.FPRoundingMode.RN) + bottom = prims.cvt_packfloat_f32(c[3], c[2], 0, prims.CVTPackFloat.F16X2, + rnd=prims.FPRoundingMode.RN) + prims.stmatrix(output + element, [top, bottom], prims.MMALayout.ROW) + + +@cute.kernel +def blockscale_gemm( + a_map: cutlass.GridConstant[cuda.TensorMap], + b_map: cutlass.GridConstant[cuda.TensorMap], + sa_map: cutlass.GridConstant[cuda.TensorMap], + sb_map: cutlass.GridConstant[cuda.TensorMap], + c_map: cutlass.GridConstant[cuda.TensorMap], +): + tid, _, _ = cute.arch.thread_idx() + tile, _, _ = cute.arch.block_idx() + stride, _, _ = cute.arch.grid_dim() + warp = tid // 32 + lane = tid % 32 + storage = cutlass.Array(dtype=cutlass.Uint8, shape=(SMEM_BYTES,), + space=cutlass.AddressSpace.smem, alignment=1024) + base = storage.data_ptr() + full = cutlass.Pointer(base + BARRIERS, dtype=cutlass.Int64) + empty = full + STAGES + if tid == 0: + for s in cutlass.range_constexpr(STAGES): + prims.mbarrier_init(full + s, 1) + prims.mbarrier_init(empty + s, 8) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0, thread_count=288) + + slot = cutlass.Int32(0) + if warp < 8: + wm = warp % 4 + wn = warp // 4 + phase = cutlass.Int32(0) + while tile < 256: + tile_m = tile % 16 + tile_n = tile // 16 + + c00 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c01 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c02 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c03 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c04 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c05 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c06 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c07 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c10 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c11 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c12 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c13 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c14 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c15 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c16 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + c17 = cutlass.full((4,), 0.0, dtype=cutlass.Float32) + stage = cutlass.Int32(0) + while stage < 8: + while not prims.mbarrier_try_wait_parity(full + slot, phase): + pass + a = base + slot * DATA_BYTES + b = base + B_OFFSET + slot * DATA_BYTES + sa = cutlass.Pointer(base + SFA_OFFSET + slot * SCALE_BYTES, + dtype=cutlass.Uint32) + sb = cutlass.Pointer(base + SFB_OFFSET + slot * SCALE_BYTES, + dtype=cutlass.Uint32) + a_provider = lane // 4 + 8 * (lane & 1) + sa0 = sa[_scale_word(16 * wm + a_provider)] + sa1 = sa[_scale_word(16 * wm + a_provider + 64)] + sb0 = sb[_scale_word(16 * wn + lane // 4 + 0)] + sb1 = sb[_scale_word(16 * wn + lane // 4 + 8)] + sb2 = sb[_scale_word(16 * wn + lane // 4 + 32)] + sb3 = sb[_scale_word(16 * wn + lane // 4 + 40)] + sb4 = sb[_scale_word(16 * wn + lane // 4 + 64)] + sb5 = sb[_scale_word(16 * wn + lane // 4 + 72)] + sb6 = sb[_scale_word(16 * wn + lane // 4 + 96)] + sb7 = sb[_scale_word(16 * wn + lane // 4 + 104)] + half = cutlass.Int32(0) + while half < 2: + a0 = prims.ldmatrix(a + _matrix_row(wm, half, lane), + 4, prims.MMALayout.ROW) + a1 = prims.ldmatrix(a + _matrix_row(wm + 4, half, lane), + 4, prims.MMALayout.ROW) + scale_a0 = (sa0 >> (16 * half)) & 65535 + scale_a1 = (sa1 >> (16 * half)) & 65535 + br = prims.ldmatrix(b + _matrix_row(wn + 0, half, lane), + 4, prims.MMALayout.ROW) + c00 = _mma(c00, a0, br[0], br[2], scale_a0, + (sb0 >> (16 * half)) & 65535) + c01 = _mma(c01, a0, br[1], br[3], scale_a0, + (sb1 >> (16 * half)) & 65535) + c10 = _mma(c10, a1, br[0], br[2], scale_a1, + (sb0 >> (16 * half)) & 65535) + c11 = _mma(c11, a1, br[1], br[3], scale_a1, + (sb1 >> (16 * half)) & 65535) + br = prims.ldmatrix(b + _matrix_row(wn + 2, half, lane), + 4, prims.MMALayout.ROW) + c02 = _mma(c02, a0, br[0], br[2], scale_a0, + (sb2 >> (16 * half)) & 65535) + c03 = _mma(c03, a0, br[1], br[3], scale_a0, + (sb3 >> (16 * half)) & 65535) + c12 = _mma(c12, a1, br[0], br[2], scale_a1, + (sb2 >> (16 * half)) & 65535) + c13 = _mma(c13, a1, br[1], br[3], scale_a1, + (sb3 >> (16 * half)) & 65535) + br = prims.ldmatrix(b + _matrix_row(wn + 4, half, lane), + 4, prims.MMALayout.ROW) + c04 = _mma(c04, a0, br[0], br[2], scale_a0, + (sb4 >> (16 * half)) & 65535) + c05 = _mma(c05, a0, br[1], br[3], scale_a0, + (sb5 >> (16 * half)) & 65535) + c14 = _mma(c14, a1, br[0], br[2], scale_a1, + (sb4 >> (16 * half)) & 65535) + c15 = _mma(c15, a1, br[1], br[3], scale_a1, + (sb5 >> (16 * half)) & 65535) + br = prims.ldmatrix(b + _matrix_row(wn + 6, half, lane), + 4, prims.MMALayout.ROW) + c06 = _mma(c06, a0, br[0], br[2], scale_a0, + (sb6 >> (16 * half)) & 65535) + c07 = _mma(c07, a0, br[1], br[3], scale_a0, + (sb7 >> (16 * half)) & 65535) + c16 = _mma(c16, a1, br[0], br[2], scale_a1, + (sb6 >> (16 * half)) & 65535) + c17 = _mma(c17, a1, br[1], br[3], scale_a1, + (sb7 >> (16 * half)) & 65535) + half += 1 + if lane == 0: + prims.mbarrier_arrive(empty + slot) + slot += 1 + if slot == STAGES: + slot = cutlass.Int32(0) + phase ^= 1 + stage += 1 + if tid == 0: + prims.cp_async_bulk_wait_group(0, read=True) + prims.barrier_cta_sync(2, thread_count=256) + output = cutlass.Pointer(base + OUTPUT, dtype=cutlass.Float16) + _store(output, c00, warp, lane, 0, 0) + _store(output, c01, warp, lane, 0, 1) + _store(output, c02, warp, lane, 0, 2) + _store(output, c03, warp, lane, 0, 3) + _store(output, c04, warp, lane, 0, 4) + _store(output, c05, warp, lane, 0, 5) + _store(output, c06, warp, lane, 0, 6) + _store(output, c07, warp, lane, 0, 7) + _store(output, c10, warp, lane, 1, 0) + _store(output, c11, warp, lane, 1, 1) + _store(output, c12, warp, lane, 1, 2) + _store(output, c13, warp, lane, 1, 3) + _store(output, c14, warp, lane, 1, 4) + _store(output, c15, warp, lane, 1, 5) + _store(output, c16, warp, lane, 1, 6) + _store(output, c17, warp, lane, 1, 7) + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=256) + if tid == 0: + prims.cp_async_bulk_tensor_global_shared_cta( + c_map.get_ptr(), base + OUTPUT, + (tile_n * 128, tile_m * 128)) + prims.cp_async_bulk_tensor_global_shared_cta( + c_map.get_ptr(), base + OUTPUT + 16384, + (tile_n * 128 + 64, tile_m * 128)) + prims.cp_async_bulk_commit_group() + tile += stride + if tid == 0: + prims.cp_async_bulk_wait_group(0, read=True) + elif lane == 0: + phase = cutlass.Int32(1) + while tile < 256: + tile_m = tile % 16 + tile_n = tile // 16 + stage = cutlass.Int32(0) + while stage < 8: + while not prims.mbarrier_try_wait_parity(empty + slot, phase): + pass + prims.mbarrier_arrive_expect_tx(full + slot, TX_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + base + slot * DATA_BYTES, a_map.get_ptr(), + (stage * 64, tile_m * 128), full + slot) + prims.cp_async_bulk_tensor_shared_cta_global( + base + B_OFFSET + slot * DATA_BYTES, b_map.get_ptr(), + (stage * 64, tile_n * 128), full + slot) + prims.cp_async_bulk_tensor_shared_cta_global( + base + SFA_OFFSET + slot * SCALE_BYTES, sa_map.get_ptr(), + (0, tile_m * 8 + stage), full + slot) + prims.cp_async_bulk_tensor_shared_cta_global( + base + SFB_OFFSET + slot * SCALE_BYTES, sb_map.get_ptr(), + (0, tile_n * 8 + stage), full + slot) + slot += 1 + if slot == STAGES: + slot = cutlass.Int32(0) + phase ^= 1 + stage += 1 + tile += stride + tail = cutlass.Int32(0) + while tail < STAGES: + while not prims.mbarrier_try_wait_parity(empty + slot, phase): + pass + slot += 1 + if slot == STAGES: + slot = cutlass.Int32(0) + phase ^= 1 + tail += 1 diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/launch.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/launch.py new file mode 100644 index 0000000000..8fb396c0e9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/launch.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host descriptor creation and launch; no CuTe tensor/layout algebra.""" +import cuda.bindings.driver as driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda + +from kernel import blockscale_gemm + + +@cute.jit +def host(a: cutlass.Int64, b: cutlass.Int64, sa: cutlass.Int64, + sb: cutlass.Int64, c: cutlass.Int64, ctas: cutlass.Int32, + stream: driver.CUstream): + # The experimental CUDA API expresses strides in 16-byte units. + a_map = cuda.create_tensor_map_tiled( + a, cutlass.Uint8, [512, 2048], [32], [64, 128], + swizzle=cuda.TensorMapSwizzle.s64b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b) + b_map = cuda.create_tensor_map_tiled( + b, cutlass.Uint8, [512, 2048], [32], [64, 128], + swizzle=cuda.TensorMapSwizzle.s64b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b) + sa_map = cuda.create_tensor_map_tiled( + sa, cutlass.Uint16, [256, 128], [32], [256, 1], + swizzle=cuda.TensorMapSwizzle.none, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b) + sb_map = cuda.create_tensor_map_tiled( + sb, cutlass.Uint16, [256, 128], [32], [256, 1], + swizzle=cuda.TensorMapSwizzle.none, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b) + c_map = cuda.create_tensor_map_tiled( + c, cutlass.Float16, [2048, 2048], [256], [64, 128], + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b) + blockscale_gemm(a_map, b_map, sa_map, sb_map, c_map).launch( + grid=(ctas, 1, 1), block=(288, 1, 1), smem=86016, stream=stream) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/requirements.txt b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/requirements.txt new file mode 100644 index 0000000000..47083a3598 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/requirements.txt @@ -0,0 +1,3 @@ +nvidia-cutlass-dsl[cu13]==4.8.0 +cuda-bindings==13.4.3 +numpy==2.5.4 diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_benchmark.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_benchmark.py new file mode 100644 index 0000000000..d1adc9eadb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_benchmark.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fail closed on invalid benchmark requests and incomplete GPU output.""" +import contextlib +import io +import unittest + +import numpy as np + +import benchmark + + +class BenchmarkTests(unittest.TestCase): + def test_invalid_launch_counts_are_rejected_before_gpu_work(self): + for args in (["--ctas", "0"], ["--ctas", "257"], + ["--warmup", "-1"], ["--samples", "0"]): + with self.subTest(args=args), contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit): + benchmark.parse_arguments(args) + + def test_bitwise_verification_rejects_poison_and_signed_zero_difference(self): + expected = np.array([[0x0000, 0x3c00]], dtype=np.uint16) + benchmark.verify_bits(expected.copy(), expected) + for got in (np.array([[0x8000, 0x3c00]], dtype=np.uint16), + np.array([[0x0000, 0x7e00]], dtype=np.uint16)): + with self.assertRaisesRegex(ValueError, "mismatch"): + benchmark.verify_bits(got, expected) + + def test_verification_rejects_missing_outputs_and_wrong_storage_type(self): + expected = np.array([[0, 1]], dtype=np.uint16) + for got in (expected[:, :1], expected.astype(np.float16)): + with self.assertRaises(ValueError): + benchmark.verify_bits(got, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_fixture.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_fixture.py new file mode 100644 index 0000000000..68f53b75ae --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/reference/test_fixture.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only byte-layout and independent exact-result checks for the CTM fixture.""" + +import importlib +import importlib.util +import struct +import unittest + +import numpy as np + + +def scalar_mix(value): + """Rust's wrapping u32 arithmetic, evaluated with Python integers.""" + value &= 0xFFFFFFFF + value ^= value >> 16 + value = value * 0x7FEB352D % (1 << 32) + value ^= value >> 15 + value = value * 0x846CA68B % (1 << 32) + return value ^ (value >> 16) + + +def scalar_nibble(row, k, operand): + multiplier, seed = { + "a": (0x9E3779B9, 0xA511E9B3), + "b": (0x85EBCA6B, 0x63D83595), + }[operand] + return scalar_mix((row * multiplier & 0xFFFFFFFF) ^ k ^ seed) % 16 + + +def scalar_scale(row, group, operand): + multiplier, seed = { + "a": (0xC2B2AE35, 0), + "b": (0x27D4EB2F, 0x165667B1), + }[operand] + return 126 + scalar_mix((row * multiplier & 0xFFFFFFFF) ^ group ^ seed) % 3 + + +def scalar_result_bits(row, column): + """Accumulate integer sixteenths, independently of NumPy or BLAS.""" + twice_magnitude = (0, 1, 2, 3, 4, 6, 8, 12) + total = 0 + for group in range(32): + block = 0 + for k in range(group * 32, (group + 1) * 32): + a = scalar_nibble(row, k, "a") + b = scalar_nibble(column, k, "b") + product = twice_magnitude[a % 8] * twice_magnitude[b % 8] + block += -product if (a ^ b) & 8 else product + # Each value and each scale was doubled: products are in sixteenths. + total += block * (1 << (scalar_scale(row, group, "a") - 126)) * ( + 1 << (scalar_scale(column, group, "b") - 126) + ) + return struct.unpack("> 4)) + self.assertEqual(observed, set(range(16))) + for row in (0, 1, 31, 32, 127, 128, 1023, 2047): + want = [ + scalar_nibble(row, k, operand) + | (scalar_nibble(row, k + 1, operand) << 4) + for k in range(0, 1024, 2) + ] + with self.subTest(operand=operand, row=row): + np.testing.assert_array_equal(packed[row], want) + + def test_scales_use_little_endian_words_and_canonical_atoms(self): + fixture = self.fixture() + self.assertEqual(fixture["sa"][0, :2].tolist(), [0x7F7E, 0x7F7E]) + self.assertEqual(fixture["sb"][0, :2].tolist(), [0x7E7E, 0x807F]) + self.assertEqual(int(fixture["sa"][1, 0]), 0x7F7F) + self.assertEqual(int(fixture["sb"][1, 0]), 0x8080) + for operand, name in (("a", "sa"), ("b", "sb")): + raw = fixture[name].astype(" Self { + // q <= 23 and q+8 <= 31, so scale_word reduces to 4*q and 4*(q+8). + // Each start is 16-byte aligned; even the last vector ends at byte 512. + let q = 16 * wn + lane / 4; + let even = unsafe { *stage.add(physical::scale_word(q)).cast::() }; + let odd = unsafe { *stage.add(physical::scale_word(q + 8)).cast::() }; + let mut words = [[0; 2]; 4]; + mma::static_for!(PAIR in [0, 1, 2, 3] { + // Narrow after constant shifts, before the runtime K-half selection. + words[PAIR] = [(even >> (32 * PAIR)) as u32, (odd >> (32 * PAIR)) as u32]; + }); + Self { words } + } + + /// Select the two UE8M0 scale pairs for K64 half 0 or 1. + #[inline(always)] + fn pair(&self, half: usize) -> [ScalePairB; 2] { + const { assert!(PAIR < 4) } + [ + ScalePairB::from_bits((self.words[PAIR][0] >> (16 * half)) as u16), + ScalePairB::from_bits((self.words[PAIR][1] >> (16 * half)) as u16), + ] + } +} + +/// Borrow a typed shared tile; the caller provides live, aligned storage covering its full footprint. +#[inline(always)] +unsafe fn region<'s, F: PhysicalFormat>(pointer: *mut u8) -> SharedRegion<'s, F> { + unsafe { SharedRegion::from_raw_parts_or_trap(pointer.cast(), F::FOOTPRINT_BYTES) } +} +/// Retry CTA parity probes on a live initialized barrier until the phase completes. +#[inline(always)] +unsafe fn wait(bar: &ManagedBarrier, phase: u32) { + unsafe { while !mbarrier::try_wait_parity(bar, phase) {} } +} + +// BAND/COL identify fragments; the example retains its explicit physical mapping. +/// Compute an FP16 element address for BAND<2, COL<8, warp<8, and lane<32. +#[inline(always)] +const fn output_fragment_row( + warp: usize, + lane: usize, +) -> usize { + let row = 16 * (warp % 4) + 64 * BAND + (lane & 15); + // Resolve the table before MIR import: a runtime array promotion would + // otherwise require device global-allocation support. + let column = 16 * (warp / 4) + const { physical::N_OFFSETS[COL] }; + physical::output_element(row, column) +} + +/// Store a converted accumulator through stmatrix; callers provide the complete warp and live aligned output tile. +#[inline(always)] +unsafe fn store_fragment( + output: *mut f16, + warp: usize, + lane: usize, + fragment: Accumulator16x8, +) { + unsafe { + mma::store_f16( + SharedAddress::from_raw( + output + .add(output_fragment_row::(warp, lane)) + .cast(), + ), + fragment, + ); + } +} + +/// SM120 block-scaled GEMM entry point built from typed instruction primitives. +#[cuda_module] +pub mod kernels { + + use super::*; + /// Compute persistent 128x128 tiles from five live tensor maps using three load stages and eight compute warps. + #[launch_bounds(288)] + #[launch_contract(domain=1, block=(288,1,1), dynamic_shared=86016, + dynamic_shared_alignment=1024, min_compute_capability=(12,0))] + #[kernel] + pub fn blockscale_gemm( + #[grid_constant] a_map: &TensorMap, + #[grid_constant] b_map: &TensorMap, + #[grid_constant] sa_map: &TensorMap, + #[grid_constant] sb_map: &TensorMap, + #[grid_constant] c_map: &TensorMap, + ) { + let tid = thread::threadIdx_x(); + let warp = (tid / 32) as usize; + let lane = (tid % 32) as usize; + let base = DynamicSharedArray::::get_raw(); + unsafe { + let full = base.add(BARRIERS).cast::(); + let empty = full.add(STAGES); + if tid == 0 { + let mut s = 0u32; + while s < STAGES as u32 { + ManagedBarrier::::from_static(full.add(s as usize)) + .init_instruction(1); + ManagedBarrier::::from_static(empty.add(s as usize)) + .init_instruction(8); + s += 1; + } + } + fence::fence_mbarrier_init(); + barrier::barrier_cta_sync(0, THREADS as u32); + let mut tile = thread::blockIdx_x(); + let stride = thread::gridDim_x(); + let mut slot = 0u32; + if warp < 8 { + let wm = warp % 4; + let wn = warp / 4; + let mut phase = 0u32; + while tile < TOTAL_TILES as u32 { + let tile_m = tile % (M / 128) as u32; + let tile_n = tile / (M / 128) as u32; + let mut accumulators = AccumulatorTile16x8::<2, 8>::zero(); + let mut stage = 0u32; + while stage < K_STAGES as u32 { + let slot_idx = slot as usize; + let ready = + ManagedBarrier::::adopt_ready(full.add(slot_idx)); + wait(&ready, phase); + let a = base.add(slot_idx * DATA_BYTES); + let b = base.add(B_OFFSET + slot_idx * DATA_BYTES); + let sa = base.add(SFA_OFFSET + slot_idx * SCALE_BYTES).cast::(); + let sb = base.add(SFB_OFFSET + slot_idx * SCALE_BYTES).cast::(); + let a_provider = lane / 4 + 8 * (lane & 1); + let sa0 = *sa.add(physical::scale_word(16 * wm + a_provider)); + let sa1 = *sa.add(physical::scale_word(16 * wm + a_provider + 64)); + let scales_b = BScaleFragments::load(sb, wn, lane); + let mut half = 0u32; + while half < 2 { + let half_idx = half as usize; + let a0 = mma::load_mxf4_a(SharedAddress::from_raw( + a.add(physical::matrix_row(wm, half_idx, lane)).cast(), + )); + let a1 = mma::load_mxf4_a(SharedAddress::from_raw( + a.add(physical::matrix_row(wm + 4, half_idx, lane)).cast(), + )); + let scale_a0 = ScalePairA::from_bits((sa0 >> (16 * half)) as u16); + let scale_a1 = ScalePairA::from_bits((sa1 >> (16 * half)) as u16); + // Finish all four MMAs before loading the next B pair. + mma::static_for!(PAIR in [0, 1, 2, 3] { + let (b0, b1) = mma::load_mxf4_b_pair(SharedAddress::from_raw( + b.add(physical::matrix_row(wn + 2 * PAIR, half_idx, lane)).cast(), + )); + accumulators.mma_mxf4_pair::( + [a0, a1], + [b0, b1], + [scale_a0, scale_a1], + scales_b.pair::(half_idx), + ); + }); + half += 1; + } + if lane == 0 { + mbarrier::arrive( + &ManagedBarrier::::adopt_ready( + empty.add(slot_idx), + ), + ); + } + slot += 1; + if slot == STAGES as u32 { + slot = 0; + phase ^= 1; + } + stage += 1; + } + // The producer warp never joins the 256-thread output barriers. + if tid == 0 { + tma::cp_async_bulk_wait_group_read::<0>(); + } + barrier::barrier_cta_sync(2, 256); + let output = base.add(OUTPUT).cast::(); + mma::static_for!(BAND in [0, 1] { + mma::static_for!(COL in [0, 1, 2, 3, 4, 5, 6, 7] { + store_fragment::( + output, warp, lane, accumulators.fragment::(), + ); + }); + }); + fence::fence_proxy_async_shared_cta(); + barrier::barrier_cta_sync(2, 256); + if tid == 0 { + let left = region::(base.add(OUTPUT)); + let right = region::(base.add(OUTPUT + 16384)); + tma::cp_async_bulk_tensor_global_shared_cta( + &left, + c_map, + Coordinates2D::new((tile_n * 128) as i32, (tile_m * 128) as i32), + ); + tma::cp_async_bulk_tensor_global_shared_cta( + &right, + c_map, + Coordinates2D::new((tile_n * 128 + 64) as i32, (tile_m * 128) as i32), + ); + tma::cp_async_bulk_commit_group(); + } + tile += stride; + } + if tid == 0 { + tma::cp_async_bulk_wait_group_read::<0>(); + } + } else if lane == 0 { + let mut phase = 1u32; + while tile < TOTAL_TILES as u32 { + let tile_m = tile % (M / 128) as u32; + let tile_n = tile / (M / 128) as u32; + let mut stage = 0u32; + while stage < K_STAGES as u32 { + let slot_idx = slot as usize; + wait( + &ManagedBarrier::::adopt_ready( + empty.add(slot_idx), + ), + phase, + ); + let ready = + ManagedBarrier::::adopt_ready(full.add(slot_idx)); + mbarrier::arrive_expect_tx(&ready, TX_BYTES); + let a = region::(base.add(slot_idx * DATA_BYTES)); + let b = region::(base.add(B_OFFSET + slot_idx * DATA_BYTES)); + let sa = + region::(base.add(SFA_OFFSET + slot_idx * SCALE_BYTES)); + let sb = + region::(base.add(SFB_OFFSET + slot_idx * SCALE_BYTES)); + tma::cp_async_bulk_tensor_shared_cta_global( + &a, + a_map, + Coordinates2D::new((stage * 64) as i32, (tile_m * 128) as i32), + ready.shared_address(), + ); + tma::cp_async_bulk_tensor_shared_cta_global( + &b, + b_map, + Coordinates2D::new((stage * 64) as i32, (tile_n * 128) as i32), + ready.shared_address(), + ); + tma::cp_async_bulk_tensor_shared_cta_global( + &sa, + sa_map, + Coordinates2D::new(0, (tile_m * K_STAGES as u32 + stage) as i32), + ready.shared_address(), + ); + tma::cp_async_bulk_tensor_shared_cta_global( + &sb, + sb_map, + Coordinates2D::new(0, (tile_n * K_STAGES as u32 + stage) as i32), + ready.shared_address(), + ); + slot += 1; + if slot == STAGES as u32 { + slot = 0; + phase ^= 1; + } + stage += 1; + } + tile += stride; + } + let mut tail = 0u32; + while tail < STAGES as u32 { + wait( + &ManagedBarrier::::adopt_ready( + empty.add(slot as usize), + ), + phase, + ); + slot += 1; + if slot == STAGES as u32 { + slot = 0; + phase ^= 1; + } + tail += 1; + } + } + } + } +} + +#[cfg(test)] +mod tests { + //! CPU checks for vectorized scale loading and constant-index fragment placement. + + use super::*; + + /// Aligned scale-ring test storage with distinguishable words and K halves. + #[repr(align(16))] + struct ScaleStages( + #[doc = "Three 512-byte scale stages represented as consecutive u32 words."] + [u32; 128 * STAGES], + ); + + /// Fill the scale ring with tags distinguishing each stage, word, and 16-bit half. + fn tagged_scales() -> ScaleStages { + // Distinct low/high halves also distinguish adjacent words and stages. + ScaleStages(core::array::from_fn(|index| { + ((0x8000 | index as u32) << 16) | (0x1200 | index as u32) + })) + } + + /// Compare vectorized scale loads with scalar physical addressing across every legal provider and K half. + #[test] + fn vector_scale_loads_match_scalar_mapping_for_all_stages_warps_lanes_and_halves() { + let storage = tagged_scales(); + for slot in 0..STAGES { + let stage = unsafe { storage.0.as_ptr().add(slot * 128) }; + for warp in 0..8 { + let wn = warp / 4; + for lane in 0..32 { + let q = 16 * wn + lane / 4; + for row in [q, q + 8] { + let offset = 4 * physical::scale_word(row); + assert_eq!((stage as usize + offset) % 16, 0); + assert!(offset + 16 <= SCALE_BYTES); + } + let scales = unsafe { BScaleFragments::load(stage, wn, lane) }; + for half in 0..2 { + mma::static_for!(PAIR in [0, 1, 2, 3] { + let actual = scales.pair::(half); + for (side, scale) in actual.into_iter().enumerate() { + let row = q + physical::N_OFFSETS[2 * PAIR + side]; + let word = storage.0[slot * 128 + physical::scale_word(row)]; + assert_eq!(scale.bits(), (word >> (16 * half)) as u16); + } + }); + } + } + } + } + } + + /// Check the final vector provider reaches the last stage words without reordering either K half. + #[test] + fn last_scale_provider_preserves_vector_word_order_and_k_halves() { + let storage = tagged_scales(); + // wn=1, lane=31: q=23, the two vectors start at words 92 and 124. + let scales = unsafe { BScaleFragments::load(storage.0.as_ptr(), 1, 31) }; + assert_eq!(scales.pair::<0>(0).map(ScalePairB::bits), [0x125c, 0x127c]); + assert_eq!(scales.pair::<0>(1).map(ScalePairB::bits), [0x805c, 0x807c]); + assert_eq!(scales.pair::<3>(0).map(ScalePairB::bits), [0x125f, 0x127f]); + assert_eq!(scales.pair::<3>(1).map(ScalePairB::bits), [0x805f, 0x807f]); + } + + /// Check compile-time fragment addressing agrees with the independent physical map for every warp/lane. + #[test] + fn const_fragment_store_addresses_match_the_physical_mapping_for_every_lane() { + for warp in 0..8 { + for lane in 0..32 { + mma::static_for!(BAND in [0, 1] { + mma::static_for!(COL in [0, 1, 2, 3, 4, 5, 6, 7] { + assert_eq!( + output_fragment_row::(warp, lane), + physical::output_row(warp, lane, BAND, COL), + ); + }); + }); + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/main.rs b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/main.rs new file mode 100644 index 0000000000..771e74a7eb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/main.rs @@ -0,0 +1,412 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Host data packing, CPU verification, and CUDA-event timing for instruction-level MXFP4 GEMM. + +#![feature(f16)] +mod kernel; +mod physical; +use ctm_rs::host::tma::{EncodeOptions, L2Promotion, create_tensor_map_tiled}; +use ctm_rs::{ + memory::{F16Tile128, Mxf4ScaleTile, Mxf4Tile128}, + tma::TensorMap, +}; +use cuda_core::{CudaContext, LaunchConfig1D}; +use kernel::{DYNAMIC_SMEM_BYTES, K, K_STAGES, KB, M, N, THREADS, TOTAL_TILES, kernels}; + +/// The eight non-negative e2m1 values, indexed by nibble 0x0..=0x7; +/// nibbles 0x8..=0xF are their negations (0x8 is -0.0). +const E2M1_VALUES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + +/// Decode the sign and three-bit magnitude of one FP4 nibble for the CPU oracle. +fn e2m1_decode(nibble: u8) -> f32 { + let magnitude = E2M1_VALUES[(nibble & 0x7) as usize]; + if nibble & 0x8 != 0 { + -magnitude + } else { + magnitude + } +} + +/// Small deterministic hash used only to build the host fixture. +/// +/// Mixing both coordinates prevents tile-sized row/column permutations from +/// hiding behind a short periodic input pattern. +fn mix32(mut value: u32) -> u32 { + value ^= value >> 16; + value = value.wrapping_mul(0x7feb_352d); + value ^= value >> 15; + value = value.wrapping_mul(0x846c_a68b); + value ^ (value >> 16) +} + +/// Deterministic data covering all 16 nibbles (zeros, negatives, both -0). +fn a_nibble(m: usize, k: usize) -> u8 { + (mix32((m as u32).wrapping_mul(0x9e37_79b9) ^ k as u32 ^ 0xa511_e9b3) & 0xf) as u8 +} + +/// Generate a deterministic B value from its K/N coordinates, including signed zeros. +fn b_nibble(k: usize, n: usize) -> u8 { + (mix32((n as u32).wrapping_mul(0x85eb_ca6b) ^ k as u32 ^ 0x63d8_3595) & 0xf) as u8 +} + +/// ue8m0 scale bytes: 2^(byte-127). 126..=128 keeps every product a small +/// binary fraction that `f32` represents exactly. +fn sfa_byte(m: usize, kb: usize) -> u8 { + 126 + (mix32((m as u32).wrapping_mul(0xc2b2_ae35) ^ kb as u32) % 3) as u8 +} + +/// Generate a deterministic B scale exponent in 126..=128 from its row and K-group indices. +fn sfb_byte(n: usize, kb: usize) -> u8 { + 126 + (mix32((n as u32).wrapping_mul(0x27d4_eb2f) ^ kb as u32 ^ 0x1656_67b1) % 3) as u8 +} + +/// Decode a finite UE8M0 exponent byte as 2 raised to byte minus 127. +fn ue8m0_decode(byte: u8) -> f32 { + 2f32.powi(byte as i32 - 127) +} + +/// Pack four increasing-K FP4 nibbles into each 16-bit transport carrier. +fn pack_mxf4_view(nibbles: &[u8]) -> Vec { + nibbles + .as_chunks::<4>() + .0 + .iter() + .map(|q| { + f16::from_bits( + (q[0] as u16 & 0xf) + | ((q[1] as u16 & 0xf) << 4) + | ((q[2] as u16 & 0xf) << 8) + | ((q[3] as u16 & 0xf) << 12), + ) + }) + .collect() +} + +/// Reorder dense `[row][K/32]` scales into SM120's 128-row x four-group +/// atoms, then expose the exact bytes as the u16 carrier used by TMA. +fn pack_canonical_scales(dense: &[u8], rows: usize) -> Vec { + let mut packed = vec![0u8; rows * KB]; + for row in 0..rows { + for group in 0..KB { + packed[physical::scale_offset(row, group, KB)] = dense[row * KB + group]; + } + } + packed + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u16::from_le_bytes([pair[0], pair[1]])) + .collect() +} + +/// Select a rounded percentile index from nonempty sorted timings using an integer fraction. +fn percentile(sorted: &[f32], numerator: usize, denominator: usize) -> f32 { + assert!(!sorted.is_empty()); + assert!(denominator > 0 && numerator <= denominator); + let scaled = (sorted.len() - 1) + .checked_mul(numerator) + .expect("percentile index overflow"); + sorted[(scaled + denominator / 2) / denominator] +} + +/// Parse launch/timing options, prepare packed tensors, verify GPU results, and write timing reports. +fn main() { + let mut ctas = None; + let mut warmup = 100usize; + let mut samples = 31usize; + let mut json_path = None; + let mut arguments = std::env::args().skip(1); + while let Some(argument) = arguments.next() { + match argument.as_str() { + "--ctas" => { + ctas = Some( + arguments + .next() + .expect("--ctas value") + .parse::() + .expect("positive CTA count"), + ) + } + "--warmup" => { + warmup = arguments + .next() + .expect("--warmup value") + .parse() + .expect("warmup count") + } + "--samples" => { + samples = arguments + .next() + .expect("--samples value") + .parse() + .expect("sample count") + } + "--json" => json_path = Some(arguments.next().expect("--json path")), + _ => panic!("unknown option {argument}; use --ctas, --warmup, --samples, --json"), + } + } + assert!(samples > 0, "at least one timing sample is required"); + if let Some(ctas) = ctas { + assert!( + ctas > 0 && ctas <= TOTAL_TILES as u32, + "CTA count out of range" + ); + } + if let Some(path) = &json_path { + match std::fs::remove_file(path) { + Ok(()) => (), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => (), + Err(error) => panic!("cannot remove stale report: {error}"), + } + } + println!("=== Block-scaled MXFP4 GEMM with Rust CTM ==="); + println!("C({M}x{N}) = A({M}x{K} e2m1) x B({K}x{N} e2m1), ue8m0 scales per 32-element block\n"); + + let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); + + // Accelerated MXFP4 block-scaled MMA exists on SM120/121 targets. + let (major, minor) = ctx.compute_capability().expect("query compute capability"); + println!("GPU Compute Capability: sm_{major}{minor}"); + if (major, minor) != (12, 0) { + eprintln!("mxf4 primitives require an SM120 GPU"); + std::process::exit(1); + } + + // Host-side packed FP4 payloads; B is stored transposed as N x K. + let a_nibbles: Vec = (0..M * K).map(|i| a_nibble(i / K, i % K)).collect(); + let b_nibbles: Vec = (0..N * K).map(|i| b_nibble(i % K, i / K)).collect(); + let sfa: Vec = (0..M * KB).map(|i| sfa_byte(i / KB, i % KB)).collect(); + let sfb: Vec = (0..N * KB).map(|i| sfb_byte(i / KB, i % KB)).collect(); + + // Exact f32 accumulator reference: per K-block, an exact 32-term dot times the two + // power-of-two scales. Every partial sum is an exact binary fraction well + // inside f32's exact range, so accumulation order cannot matter. + println!("computing exact host reference ({M}x{N}x{K})..."); + let mut expected = vec![0.0f32; M * N]; + for m in 0..M { + let a_row = &a_nibbles[m * K..(m + 1) * K]; + let sfa_row = &sfa[m * KB..(m + 1) * KB]; + for n in 0..N { + let b_row = &b_nibbles[n * K..(n + 1) * K]; + let sfb_row = &sfb[n * KB..(n + 1) * KB]; + let mut total = 0.0f32; + for kb in 0..KB { + let mut block = 0.0f32; + for j in 0..32 { + let k = kb * 32 + j; + block += e2m1_decode(a_row[k]) * e2m1_decode(b_row[k]); + } + total += ue8m0_decode(sfa_row[kb]) * ue8m0_decode(sfb_row[kb]) * block; + } + expected[m * N + n] = total; + } + } + let expected_f16: Vec = expected.iter().copied().map(|value| value as f16).collect(); + + let stream = ctx.default_stream(); + // SAFETY: this executable is the sole owner of the embedded device module. + let module = unsafe { kernels::load(&ctx) }.expect("Failed to load embedded CUDA module"); + + let a_packed = pack_mxf4_view(&a_nibbles); + let b_packed = pack_mxf4_view(&b_nibbles); + let sfa_packed = pack_canonical_scales(&sfa, M); + let sfb_packed = pack_canonical_scales(&sfb, N); + let a_dev = cuda_core::DeviceBuffer::from_host(&stream, &a_packed).unwrap(); + let b_dev = cuda_core::DeviceBuffer::from_host(&stream, &b_packed).unwrap(); + let sfa_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfa_packed).unwrap(); + let sfb_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfb_packed).unwrap(); + + // The descriptor types are the host/kernel agreement: A/B land in the + // B64 byte layout, while each flattened scale row is one canonical atom. + let a_desc: TensorMap = create_tensor_map_tiled( + a_dev.cu_deviceptr() as *mut core::ffi::c_void, + M as u64, + (K / 2) as u64, + (K / 2) as u64, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .expect("A TMA descriptor encoding failed"); + let b_desc: TensorMap = create_tensor_map_tiled( + b_dev.cu_deviceptr() as *mut core::ffi::c_void, + N as u64, + (K / 2) as u64, + (K / 2) as u64, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .expect("B TMA descriptor encoding failed"); + let sfa_atom_rows = (M / 128) * K_STAGES; + let sfb_atom_rows = (N / 128) * K_STAGES; + let sfa_desc: TensorMap = create_tensor_map_tiled( + sfa_dev.cu_deviceptr() as *mut core::ffi::c_void, + sfa_atom_rows as u64, + 256, + 256, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .expect("SFA TMA descriptor encoding failed"); + let sfb_desc: TensorMap = create_tensor_map_tiled( + sfb_dev.cu_deviceptr() as *mut core::ffi::c_void, + sfb_atom_rows as u64, + 256, + 256, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .expect("SFB TMA descriptor encoding failed"); + let poison = vec![f16::NAN; M * N]; + let out = cuda_core::DeviceBuffer::from_host(&stream, &poison).unwrap(); + let out_desc: TensorMap = create_tensor_map_tiled( + out.cu_deviceptr() as *mut core::ffi::c_void, + M as u64, + N as u64, + N as u64, + EncodeOptions { + l2_promotion: L2Promotion::B128, + }, + ) + .expect("C TMA descriptor encoding failed"); + // Ask the driver how many copies of this CTA fit on each SM. The + // persistent scheduler needs that number to choose its physical grid. + let probe = module + .prepare_blockscale_gemm(LaunchConfig1D::new( + 1, + THREADS as u32, + DYNAMIC_SMEM_BYTES as u32, + )) + .expect("launch contract validation failed"); + let active_ctas = probe + .function() + .max_active_blocks_per_multiprocessor(THREADS as u32, DYNAMIC_SMEM_BYTES as u32) + .expect("query active CTAs per SM"); + let sm_count = ctx.multiprocessor_count().expect("query SM count"); + let persistent_ctas = sm_count + .checked_mul(active_ctas) + .expect("resident grid overflow") + .min(TOTAL_TILES as u32) + .max(1); + drop(probe); + let persistent_ctas = ctas.unwrap_or(persistent_ctas); + println!("Persistent grid: {persistent_ctas} CTAs"); + + let config = LaunchConfig1D::new(persistent_ctas, THREADS as u32, DYNAMIC_SMEM_BYTES as u32); + let launch = module + .prepare_blockscale_gemm(config) + .expect("persistent launch contract validation failed"); + // SAFETY: 288-thread blocks contain eight compute warps and one producer + // warp. The grid-stride scheduler partitions all 256 output tiles into + // disjoint CTA streams, and the warp/lane mapping partitions each tile. + // SAFETY: all descriptor allocations and tensor buffers remain live through stream completion. + unsafe { + module.blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + } + .expect("kernel launch failed"); + let got = out.to_host_vec(&stream).unwrap(); + + let mut errors = 0; + let mut nonzero = 0; + for m in 0..M { + for n in 0..N { + if expected[m * N + n] != 0.0 { + nonzero += 1; + } + if got[m * N + n].to_bits() != expected_f16[m * N + n].to_bits() { + if errors < 5 { + eprintln!( + " C({m},{n}): expected {} (0x{:04x}), got {} (0x{:04x})", + expected_f16[m * N + n], + expected_f16[m * N + n].to_bits(), + got[m * N + n], + got[m * N + n].to_bits(), + ); + } + errors += 1; + } + } + } + if errors != 0 { + println!("\n✗ FAILED: {errors} mismatches"); + std::process::exit(1); + } + println!( + "Numeric check: all {} FP16 output bit patterns exact ({nonzero} nonzero FP32 accumulators), {} K-blocks", + M * N, + KB, + ); + + // Single-kernel device timing: every sample places exactly one launch + // between CUDA events. The throughput count is 2*M*N*K FLOPs per launch. + + for _ in 0..warmup { + // SAFETY: all descriptor allocations and tensor buffers remain live through stream completion. + unsafe { + module.blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + } + .expect("warmup launch failed"); + } + stream.synchronize().expect("warmup sync failed"); + + let mut times_ms = Vec::with_capacity(samples); + for _ in 0..samples { + let start = ctx + .new_event_with_flags(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create start event"); + let end = ctx + .new_event_with_flags(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create end event"); + + start.record(&stream).expect("record start event"); + // SAFETY: all descriptor allocations and tensor buffers remain live through stream completion. + unsafe { + module.blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + } + .expect("timed launch failed"); + end.record(&stream).expect("record end event"); + times_ms.push(start.elapsed_ms(&end).expect("event timing")); + } + times_ms.sort_by(f32::total_cmp); + let median_ms = percentile(×_ms, 1, 2); + let p10_ms = percentile(×_ms, 1, 10); + let p90_ms = percentile(×_ms, 9, 10); + let tflops = (2.0 * (M * N * K) as f64) / (f64::from(median_ms) * 1e-3) / 1e12; + println!( + "Single-kernel device time: median={median_ms:.6} ms, p10={p10_ms:.6} ms, \ + p90={p90_ms:.6} ms ({samples} samples, {warmup} warmups); {tflops:.2} TFLOP/s" + ); + + let replay = out + .to_host_vec(&stream) + .expect("copy repeated-launch output"); + assert!( + replay + .iter() + .zip(&expected_f16) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "repeated launches changed numerical results" + ); + if let Some(path) = json_path { + let report = serde_json::json!({"verification":"bit-exact-pass", "mnk":[M,N,K], + "compute_capability":[major,minor], "ctas":persistent_ctas, + "warmup":warmup,"samples":samples,"median_ms":median_ms, + "p10_ms":p10_ms,"p90_ms":p90_ms,"tflops":tflops, + "checked_outputs":M*N,"replay_verified":true}); + std::fs::write(path, serde_json::to_string_pretty(&report).unwrap()).expect("write report"); + } + println!( + "\n✓ SUCCESS: block-scaled MXFP4 GEMM is bit-exact through Rust instruction primitives" + ); +} diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/physical.rs b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/physical.rs new file mode 100644 index 0000000000..bddc6f1c1c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/src/physical.rs @@ -0,0 +1,165 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +//! Explicit physical addresses for this kernel, in bytes unless named otherwise. + +/// Logical N-column offsets of the eight accumulator fragment slots. +pub const N_OFFSETS: [usize; 8] = [0, 8, 32, 40, 64, 72, 96, 104]; +/// Return the packed u32 word index within a scale atom from a row's low seven bits. +#[inline(always)] +pub const fn scale_word(row: usize) -> usize { + // Only the lowest seven bits participate, even for a full-matrix row. + let row = row as u32; + ((row & 31) * 4 + ((row >> 5) & 3)) as usize +} +/// Return a scale byte offset for a full-matrix row and K-group index in the blocked atom layout. +#[inline(always)] +pub const fn scale_offset(row: usize, group: usize, groups: usize) -> usize { + (row / 128) * (groups / 4) * 512 + (group / 4) * 512 + scale_word(row) * 4 + group % 4 +} +/// Byte offset for `atom < 8`, `half < 2`, and `lane < 32`. +#[inline(always)] +pub const fn matrix_row(atom: usize, half: usize, lane: usize) -> usize { + let atom = atom as u32; + let half = half as u32; + let lane = lane as u32; + let sub = lane / 8; + let row = 16 * atom + 8 * (sub % 2) + lane % 8; + let base = row * 64 + 16 * (sub / 2); + // The K64 half selects bit 5 independently of the lane/atom swizzle. + ((base ^ ((base & 0x180) >> 3)) ^ (half * 32)) as usize +} +/// Element offset for `row < 128` and `column < 128`. +#[inline(always)] +pub const fn output_element(row: usize, column: usize) -> usize { + let row = row as u32; + let column = column as u32; + let plain = row * 64 + column % 64; + ((column / 64) * 8192 + (plain ^ ((plain & 0x1c0) >> 3))) as usize +} +/// Element offset for `warp < 8`, `lane < 32`, `band < 2`, and `slot < 8`. +#[cfg(test)] +#[inline(always)] +pub const fn output_row(warp: usize, lane: usize, band: usize, slot: usize) -> usize { + let warp = warp as u32; + let lane = lane as u32; + let band = band as u32; + let row = 16 * (warp % 4) + 64 * band + (lane & 15); + let column = 16 * (warp / 4) + N_OFFSETS[slot] as u32; + output_element(row as usize, column as usize) +} + +#[cfg(test)] +mod tests { + //! Exhaustive CPU checks for bounded scale, matrix-load, and accumulator address calculations. + + use super::*; + + /// Check narrowed arithmetic matches the original usize layout throughout each legal coordinate domain. + #[test] + fn bounded_addresses_match_the_original_usize_mapping() { + /// Independent usize output-element calculation used to validate the narrowed implementation. + fn reference_output(row: usize, column: usize) -> usize { + let plain = row * 64 + column % 64; + (column / 64) * 8192 + (plain ^ ((plain & 0x1c0) >> 3)) + } + + for row in 0..2048 { + assert_eq!(scale_word(row), (row & 31) * 4 + ((row >> 5) & 3)); + } + // Truncation must also preserve scale_word's full usize domain. + for low in 0..128 { + let row = (usize::MAX & !127) | low; + assert_eq!(scale_word(row), (row & 31) * 4 + ((row >> 5) & 3)); + } + for atom in 0..8 { + for half in 0..2 { + for lane in 0..32 { + let sub = lane / 8; + let row = 16 * atom + 8 * (sub % 2) + lane % 8; + let column = 16 * half + 8 * (sub / 2); + let plain = row * 64 + column * 2; + assert_eq!(matrix_row(atom, half, lane), plain ^ ((plain & 0x180) >> 3)); + } + } + } + for row in 0..128 { + for column in 0..128 { + assert_eq!(output_element(row, column), reference_output(row, column)); + } + } + for warp in 0..8 { + for lane in 0..32 { + for band in 0..2 { + for (slot, offset) in N_OFFSETS.into_iter().enumerate() { + let row = 16 * (warp % 4) + 64 * band + (lane & 15); + let column = 16 * (warp / 4) + offset; + assert_eq!( + output_row(warp, lane, band, slot), + reference_output(row, column) + ); + } + } + } + } + } + + /// Check scale atoms are byte-complete bijections with adjacent bytes for four K groups. + #[test] + fn scale_atoms_cover_every_byte_once_and_keep_four_k_groups_adjacent() { + let mut seen = [false; 512]; + for row in 0..128 { + for group in 0..4 { + let offset = scale_word(row) * 4 + group; + assert!(offset < 512 && !seen[offset]); + seen[offset] = true; + } + } + assert!(seen.into_iter().all(|value| value)); + assert_eq!((scale_word(0), scale_word(32), scale_word(1)), (0, 1, 4)); + assert_eq!(scale_offset(128, 4, 8), 1536); + } + /// Check each K64 half provides distinct in-bounds 16-byte-aligned matrix-load rows. + #[test] + fn ldmatrix_addresses_cover_each_k64_half_with_aligned_rows() { + for atom in 0..8 { + for half in 0..2 { + let mut seen = [false; 512]; + for lane in 0..32 { + let offset = matrix_row(atom, half, lane); + assert!(offset < 8192 && offset.is_multiple_of(16)); + assert!(!seen[offset / 16]); + seen[offset / 16] = true; + } + } + } + assert_eq!(matrix_row(0, 0, 1), 64); + assert_eq!(matrix_row(0, 1, 0), 32); + } + /// Check all warp/lane/register fragments cover the full output tile exactly once. + #[test] + fn accumulator_mapping_covers_the_output_tile_exactly_once() { + let mut seen = [false; 128 * 128]; + for warp in 0..8 { + for lane in 0..32 { + for band in 0..2 { + for (slot, offset) in N_OFFSETS.into_iter().enumerate() { + for register in 0..4 { + let row = 16 * (warp % 4) + 64 * band + lane / 4 + 8 * (register / 2); + let col = 16 * (warp / 4) + offset + 2 * (lane % 4) + register % 2; + assert!(row < 128 && col < 128); + let physical = output_element(row, col); + assert!(physical < seen.len() && !seen[physical]); + seen[physical] = true; + } + // STMATRIX address-provider lanes name the start of each + // eight-element row, in the independently computed swizzle. + let row = 16 * (warp % 4) + 64 * band + (lane & 15); + let col = 16 * (warp / 4) + offset; + assert_eq!(output_row(warp, lane, band, slot), output_element(row, col)); + } + } + } + } + assert!(seen.into_iter().all(|value| value)); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/test_compare.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/test_compare.py new file mode 100644 index 0000000000..478a0a9590 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/test_compare.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only acceptance tests for Nsight Systems comparison evidence.""" + +import contextlib +import io +import json +import os +from pathlib import Path +import sqlite3 +import tempfile +import time +import unittest +from unittest import mock + +try: + import compare +except ImportError: + compare = None + + +def gpu_report(**changes): + report = { + "verification": "bit-exact-pass", "checked_outputs": 4194304, + "mnk": [2048, 2048, 1024], "compute_capability": [12, 0], + "ctas": 170, "warmup": 2, "samples": 3, "replay_verified": True, + "median_ms": 999, "p10_ms": 998, "p90_ms": 1000, "tflops": 1, + } + report.update(changes) + return report + + +def trace(path, launches): + """Insert in reverse order so readers must order by device timestamp.""" + with sqlite3.connect(path) as db: + db.execute("CREATE TABLE StringIds (id INTEGER PRIMARY KEY, value TEXT)") + db.execute("CREATE TABLE CUPTI_ACTIVITY_KIND_KERNEL " + "(start INTEGER, end INTEGER, shortName INTEGER)") + names = {name: i for i, name in enumerate(dict.fromkeys(x[0] for x in launches))} + db.executemany("INSERT INTO StringIds VALUES (?, ?)", + [(i, name) for name, i in names.items()]) + db.executemany("INSERT INTO CUPTI_ACTIVITY_KIND_KERNEL VALUES (?, ?, ?)", + [(i * 1000000, i * 1000000 + duration, names[name]) + for i, (name, duration) in reversed(list(enumerate(launches)))]) + + +class ComparisonTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(compare, "comparison runner has not been implemented") + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + + def write_report(self, report): + path = self.directory / "result.json" + path.write_text(json.dumps(report)) + return path + + def test_accepts_full_initial_and_replay_verification(self): + started = time.time_ns() + path = self.write_report(gpu_report()) + result = compare.validate_report(path, 170, 2, 3, started) + self.assertEqual(result["checked_outputs"], 4194304) + + def test_rejects_failed_incomplete_or_mismatched_verification(self): + for key, bad_value in [ + ("verification", "failed"), ("checked_outputs", 1), + ("mnk", [2048, 2048, 2048]), ("compute_capability", [12, 1]), + ("ctas", 169), ("warmup", 3), ("samples", 2), + ("replay_verified", False), ("replay_verified", 1), + ]: + with self.subTest(key=key, bad_value=bad_value): + path = self.write_report(gpu_report(**{key: bad_value})) + with self.assertRaises(ValueError): + compare.validate_report(path, 170, 2, 3, 0) + for key in gpu_report(): + if key in {"median_ms", "p10_ms", "p90_ms", "tflops"}: + continue + with self.subTest(missing=key): + report = gpu_report() + del report[key] + path = self.write_report(report) + with self.assertRaises(ValueError): + compare.validate_report(path, 170, 2, 3, 0) + + def test_rejects_stale_missing_and_malformed_reports(self): + path = self.directory / "result.json" + with self.assertRaises(ValueError): + compare.validate_report(path, 170, 2, 3, 0) + self.write_report(gpu_report()) + os.utime(path, ns=(1, 1)) + with self.assertRaisesRegex(ValueError, "stale"): + compare.validate_report(path, 170, 2, 3, time.time_ns()) + for contents in ["{", "[]", "null"]: + path.write_text(contents) + with self.assertRaises(ValueError): + compare.validate_report(path, 170, 2, 3, 0) + + def test_discards_correctness_and_warmup_launches_in_timestamp_order(self): + path = self.directory / "trace.sqlite" + trace(path, [("blockscale_gemm", n) for n in [900000, 800000, 700000, 31000, 11000, 21000]]) + self.assertEqual(compare.kernel_samples(path, 2, 3), [31.0, 11.0, 21.0]) + + def test_accepts_explicit_mangled_name_for_all_132_launches(self): + name = "kernel_cutlass_blockscale_gemm_1..." + path = self.directory / "mangled.sqlite" + trace(path, [(name, 900000)] * 101 + [(name, n * 1000) for n in range(1, 32)]) + self.assertEqual(compare.kernel_samples(path, 100, 31, expected_name=name), + list(range(1, 32))) + + def test_explicit_kernel_name_rejects_partial_matches_and_wrong_counts(self): + expected = "kernel_cutlass_blockscale_gemm_1..." + for index, wrong in enumerate(["blockscale_gemm", expected + "_other", + "prefix_" + expected, expected[:-1]]): + with self.subTest(wrong=wrong): + path = self.directory / f"wrong-mangled-{index}.sqlite" + trace(path, [(wrong, 1000)] + [(expected, 1000)] * 131) + with self.assertRaisesRegex(ValueError, "unexpected kernel shortName"): + compare.kernel_samples(path, 100, 31, expected_name=expected) + for count in (131, 133): + with self.subTest(count=count): + path = self.directory / f"wrong-count-{count}.sqlite" + trace(path, [(expected, 1000)] * count) + with self.assertRaisesRegex(ValueError, "expected 132 kernel launches"): + compare.kernel_samples(path, 100, 31, expected_name=expected) + + def test_rejects_wrong_kernel_or_wrong_launch_count(self): + for name, launches in [ + ("wrong", [("other_kernel", 1000)] * 6), + ("missing", [("blockscale_gemm", 1000)] * 5), + ("extra", [("blockscale_gemm", 1000)] * 7), + ("mixed", [("blockscale_gemm", 1000)] * 5 + [("other_kernel", 1000)]), + ("invalid-duration", [("blockscale_gemm", 1000)] * 5 + [("blockscale_gemm", -1)]), + ]: + with self.subTest(name=name): + path = self.directory / (name + ".sqlite") + trace(path, launches) + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3) + + def test_rejects_unresolved_names_and_missing_trace_without_creating_it(self): + path = self.directory / "missing.sqlite" + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3) + self.assertFalse(path.exists()) + trace(path, [("blockscale_gemm", 1000)] * 6) + with sqlite3.connect(path) as db: + db.execute("DELETE FROM StringIds") + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3) + + def test_percentiles_use_rust_rounded_rank_including_even_median(self): + result = compare.summarize_samples([10, 1, 9, 2, 8, 3, 7, 4, 6, 5]) + self.assertEqual(result, {"median_us": 6, "p10_us": 2, "p90_us": 9}) + self.assertEqual(compare.summarize_samples([7]), + {"median_us": 7, "p10_us": 7, "p90_us": 7}) + + def test_aggregate_uses_round_medians_and_keeps_every_sample(self): + entries = [ + {"label": "rust", "samples_us": [1, 100, 1], "median_us": 1}, + {"label": "cpp", "samples_us": [4, 5, 6], "median_us": 5}, + {"label": "rust", "samples_us": [9, 10, 11], "median_us": 10}, + ] + result = compare.aggregate_rounds(entries) + self.assertEqual(result["rust"]["round_medians_us"], [1, 10]) + self.assertEqual(result["rust"]["median_of_round_medians_us"], 10) + self.assertEqual(result["rust"]["samples_us"], [1, 100, 1, 9, 10, 11]) + self.assertEqual(result["cpp"]["median_of_round_medians_us"], 5) + + def test_rotation_balances_starting_order(self): + self.assertEqual(compare.execution_order(["a", "b", "c"], 4), + [["a", "b", "c"], ["b", "c", "a"], + ["c", "a", "b"], ["a", "b", "c"]]) + + def arguments(self): + for name in ("rust", "cpp"): + path = self.directory / name + path.write_text(name) + path.chmod(0o755) + return ["--binary", f"rust={self.directory / 'rust'}", + "--binary", f"cpp={self.directory / 'cpp'}", + "--output-dir", str(self.directory / "evidence"), + "--rounds", "2", "--warmup", "2", "--samples", "3"] + + def fake_command(self, command, timeout, environment, log_path): + self.assertEqual(environment["CUDA_VISIBLE_DEVICES"], "1") + log_path.write_text("test tool output\n") + if command[1:2] == ["profile"]: + result = Path(command[command.index("--json") + 1]) + result.write_text(json.dumps(gpu_report())) + prefix = Path(command[command.index("--output") + 1]) + prefix.with_suffix(".nsys-rep").write_bytes(b"synthetic trace") + elif command[1:2] == ["export"]: + path = Path(command[command.index("--output") + 1]) + trace(path, [("blockscale_gemm", n) for n in [900000, 800000, 700000, 31000, 11000, 21000]]) + return 0 + + def test_runner_persists_serial_order_active_timings_and_fresh_evidence(self): + arguments = self.arguments() + with mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare.main(arguments), 0) + self.assertEqual(compare.main(arguments), 0) + runs = sorted((self.directory / "evidence").iterdir()) + self.assertEqual(len(runs), 2) + for directory in runs: + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "passed") + self.assertEqual(result["kernel_names"], + {"rust": "blockscale_gemm", "cpp": "blockscale_gemm"}) + self.assertEqual(result["execution_order"], [["rust", "cpp"], ["cpp", "rust"]]) + self.assertEqual([r["label"] for r in result["runs"]], ["rust", "cpp", "cpp", "rust"]) + self.assertEqual(result["summary"]["rust"]["median_of_round_medians_us"], 21) + self.assertEqual(result["summary"]["rust"]["samples_us"], [31, 11, 21, 31, 11, 21]) + self.assertEqual(len(result["binaries"]["rust"]["sha256"]), 64) + self.assertIn("nsys_version", result) + self.assertIn("gpu_snapshot", result) + self.assertTrue(all(r["status"] == "passed" for r in result["runs"])) + + def test_runner_records_and_applies_exact_kernel_name_per_label(self): + name = "kernel_cutlass_blockscale_gemm_1..." + + def mangled_command(command, timeout, environment, log_path): + status = self.fake_command(command, timeout, environment, log_path) + if command[1:2] == ["export"]: + path = Path(command[command.index("--output") + 1]) + if path.stem.endswith("-cpp"): + with sqlite3.connect(path) as db: + db.execute("UPDATE StringIds SET value = ?", (name,)) + return status + + with mock.patch.object(compare, "run_command", side_effect=mangled_command), \ + contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare.main(self.arguments() + [ + "--kernel-name", "rust=blockscale_gemm", "--kernel-name", f"cpp={name}", + ]), 0) + directory, = (self.directory / "evidence").iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["kernel_names"], {"rust": "blockscale_gemm", "cpp": name}) + self.assertEqual(result["kernel_name_field"], "shortName") + self.assertEqual(result["status"], "passed") + self.assertEqual(result["summary"]["cpp"]["samples_us"], [31, 11, 21, 31, 11, 21]) + + def test_kernel_name_options_reject_unknown_empty_and_duplicate_labels(self): + arguments = self.arguments() + for specifications in [["unknown=kernel"], ["rust="], ["rust= "], ["rust"], + ["=kernel"], ["rust=first", "rust=second"]]: + extra = [value for specification in specifications + for value in ("--kernel-name", specification)] + with self.subTest(specifications=specifications), \ + contextlib.redirect_stderr(io.StringIO()), \ + mock.patch.object(compare, "run_command") as run: + with self.assertRaises(SystemExit) as caught: + compare.main(arguments + extra) + self.assertEqual(caught.exception.code, 2) + run.assert_not_called() + + def test_missing_correctness_report_stops_before_accepting_timings(self): + def missing_report(command, timeout, environment, log_path): + status = self.fake_command(command, timeout, environment, log_path) + if command[1:2] == ["profile"]: + Path(command[command.index("--json") + 1]).unlink() + return status + + with mock.patch.object(compare, "run_command", side_effect=missing_report), \ + contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare.main(self.arguments()), 1) + directory, = (self.directory / "evidence").iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "failed") + self.assertEqual(len(result["runs"]), 1) + self.assertEqual(result["runs"][0]["status"], "failed") + self.assertNotIn("summary", result) + self.assertNotIn("samples_us", result["runs"][0]) + + def test_invalid_cli_is_rejected_before_running_tools(self): + arguments = self.arguments() + for extra in [["--rounds", "0"], ["--warmup", "-1"], + ["--samples", "0"], ["--ctas", "257"], + ["--timeout", "nan"], ["--timeout", "0"], + ["--binary", f"rust={self.directory / 'rust'}"]]: + with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit) as caught: + compare.main(arguments + extra) + self.assertEqual(caught.exception.code, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/validate.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/validate.py new file mode 100644 index 0000000000..416942506a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/validate.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run the full bit-exact MXFP4 oracle, persistent grids and optional sanitizers.""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import signal +import subprocess +import tempfile + + +def run(command, timeout): + with subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, start_new_session=True) as process: + try: + output, _ = process.communicate(timeout=timeout) + return process.returncode, output + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + output, _ = process.communicate() + return None, output + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--timeout", type=float, default=180) + parser.add_argument("--sanitizer", action="store_true") + args = parser.parse_args() + if not 0 < args.timeout < float("inf"): + parser.error("timeout must be finite and positive") + binary = args.binary.resolve(strict=True) + args.output_dir.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="run-", dir=args.output_dir.resolve())) + report = {"binary": str(binary), "binary_sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + "status": "running", "checks": []} + cases = [("resident", [], []), ("single-cta", [], ["--ctas", "1"]), + ("phase-wraps", [], ["--ctas", "17"])] + if args.sanitizer: + cases += [("memcheck", ["compute-sanitizer", "--tool", "memcheck", "--check-tensor-ops", "yes", "--error-exitcode", "1"], ["--ctas", "17"]), + ("synccheck", ["compute-sanitizer", "--tool", "synccheck", "--error-exitcode", "1"], ["--ctas", "17"])] + for name, prefix, options in cases: + result_file = directory / f"{name}.json" + command = prefix + [str(binary), *options, "--warmup", "0", "--samples", "2", "--json", str(result_file)] + try: + status, output = run(command, args.timeout) + except OSError as error: + status, output = -1, str(error) + (directory / f"{name}.log").write_text(output) + entry = {"name": name, "command": command, "exit_code": status} + report["checks"].append(entry) + try: + if status != 0: + raise ValueError("timeout" if status is None else f"exit {status}") + result = json.loads(result_file.read_text()) + if not (result["verification"] == "bit-exact-pass" + and result["checked_outputs"] == 2048*2048 + and result["mnk"] == [2048,2048,1024] + and result["compute_capability"] == [12,0] + and result["replay_verified"] is True): + raise ValueError("incomplete or invalid GPU verification") + if options and result["ctas"] != int(options[1]): + raise ValueError("persistent grid mismatch") + entry.update(status="passed", result=result) + except (ValueError, OSError, KeyError, TypeError) as error: + entry.update(status="failed", error=str(error)) + report["status"] = "failed" + (directory / "report.json").write_text(json.dumps(report, indent=2)+"\n") + print(f"FAIL {name}: {error}\n{output}\nEvidence: {directory}", flush=True) + return 1 + (directory / "report.json").write_text(json.dumps(report, indent=2)+"\n") + print(f"PASS {name}: all 4,194,304 FP16 outputs exact", flush=True) + report["status"] = "passed" + (directory / "report.json").write_text(json.dumps(report, indent=2)+"\n") + print(f"Evidence: {directory}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.lock new file mode 100644 index 0000000000..e620fa17a3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "blockscale_gemm_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.toml new file mode 100644 index 0000000000..f3a103422d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "blockscale_gemm_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "Block-scaled MXFP4 GEMM composing typed TMA pipelines, shared-memory tiled MMA, and TMA epilogues" + +# Standalone example, matching the other cuda-oxide example crates. +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/README.md b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/README.md new file mode 100644 index 0000000000..a87c78f22f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/README.md @@ -0,0 +1,86 @@ +# Block-scaled MXFP4 GEMM + +This cuda-oxide example uses `cute-rs` to multiply matrices with packed FP4 +values and a shared scale for each group of 32 values: + +```text +C[M,N] = A[M,K] × B[K,N] +M = 2048, N = 2048, K = 1024 +inputs: E2M1 (FP4) + UE8M0 scales +accumulation: f32; output: f16 +``` + +The scale moves each group of small FP4 values into a useful numerical range. +It is applied as part of the matrix multiply. The output is split into +`128 × 128` tiles, each owned by one thread block (CTA). + +## CuTe data flow + +```text +global A/B + scale Tensors + │ TMA tile copies + ▼ +shared-memory Tensors → per-warp tiled views + │ load input and scale fragments + ▼ +block-scaled MMA → f32 accumulator tile + │ tiled f16 epilogue + ▼ +shared output Tensor → TMA stores → global C Tensor +``` + +TMA (Tensor Memory Accelerator) copies whole tiles. MMA is the tensor-core +matrix multiply-and-accumulate operation. One producer warp fills three +shared-memory slots while eight compute warps consume completed slots. +A persistent scheduler lets each CTA process further output tiles. + +Typed layouts connect the host's TMA descriptors to device copies. A and B +also carry distinct operand-role types through the tiled MMA operation. +The compiler lowers this CuTe computation directly. + +## Build and run + +This kernel needs an `sm_120a` GPU such as the RTX 5090 and builds for `sm_120a` +only. Other GPUs skip execution and cannot validate this example. +From `cuda-oxide/`, install the pinned CUTLASS 4.7 compiler once, then run: + +```bash +cargo oxide toolchain install cutlass + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ + cargo oxide build blockscale_gemm_cute --arch sm_120a + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ + cargo oxide run blockscale_gemm_cute --arch sm_120a +``` + +The Rust program builds deterministic inputs and computes a substantial CPU +reference. Success means all **4,194,304 FP16 output bit patterns** match. +It then reports CUDA-event median, p10, p90 and TFLOP/s from **100 warmups and +31 single-launch samples**. This command runs Rust only; the +[CuTeDSL harness](reference/README.md) runs the separate comparison and +requires its Python dependencies. + +## CuTeDSL comparison + +RTX 5090 at `a0d5776b2` (2026-10-18), CUTLASS 4.7.0 at optimization level 2 +on both sides. Both programs pass their respective correctness checks. +Each value is the median of three round medians, with 100 warmups and 31 +single-launch samples per round, with reused buffers. Times are microseconds: + +| M / N / K | CUDA events: Rust | CUDA events: CuTeDSL | Nsight kernels: Rust | Nsight kernels: CuTeDSL | +| --- | ---: | ---: | ---: | ---: | +| 2048 / 2048 / 1024 | 13.536 | 16.832 | 11.360 | 10.176 | + +CUDA events include submission gaps; Nsight Systems reports only GEMM kernel +duration, including stalls within the kernel. Profiling runs separately, so +subtracting columns does not isolate host overhead. Rust has the shorter event +interval here; CuTeDSL has the shorter kernel duration. + +Clocks were unlocked and Rust ran first. The programs use different input +fixtures; Rust starts a new process per round, while CuTeDSL keeps one process. +These results do not establish a stable speed ranking. + +- [Reference sources, dependencies and provenance](reference/README.md) +- [Separate Rust CTM example](../blockscale_gemm_ctm) +- [CUDA C++ primitives counterpart](reference/cpp) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/.gitignore b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/.gitignore new file mode 100644 index 0000000000..2f79b547e0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/.gitignore @@ -0,0 +1,4 @@ +# Locally generated timing reports and Python bytecode. +/cutedsl-events.json +/cutedsl-profiled-events.json +/__pycache__/ diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/README.md b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/README.md new file mode 100644 index 0000000000..79c23c85d7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/README.md @@ -0,0 +1,96 @@ +# Block-scaled GEMM references + +These are the comparison sources for [CuTe-RS](../) and the separate +[Rust CTM example](../../blockscale_gemm_ctm). All use +M=N=2048, K=1024, L=1, E2M1 inputs, UE8M0 scales per 32 K values, +FP32 accumulation, and FP16 output on SM120a. + +| Path | Implementation | +| --- | --- | +| [`dense_blockscaled_gemm_persistent_cooperative.py`](dense_blockscaled_gemm_persistent_cooperative.py) | Upstream Python CuTeDSL kernel. | +| [`blockscaled_gemm_dispatch.py`](blockscaled_gemm_dispatch.py) | Its required upstream instruction-selection helper. | +| [`benchmark.py`](benchmark.py) | The Python harness used for the recorded comparison. | +| [CTM reference](../../blockscale_gemm_ctm/reference/) | Matching GEMM port using the actual CUTLASS experimental primitives API, with the Rust fixture and exact oracle. | +| [`cpp/`](cpp) | The CUDA C++ instruction-primitives counterpart, including build and layout checks. | + +The C++ counterpart was written for cuda-oxide using custom inline PTX. Its +historical results compare Rust with custom CUDA C++. Use [CTM reference](../../blockscale_gemm_ctm/reference/) for the +comparison against CUTLASS's actual primitives API. That port uses the upstream +SM120 operations; the available complete upstream blockscale primitives example +targets SM100. + +## Python setup and execution + +From this `reference/` directory, use Python 3.12 and a CUDA-enabled PyTorch +installation supporting the GPU. The pinned packages reproduce the versions +used by the recorded CuTeDSL 4.6.2 run: + +```sh +python3.12 -m venv /tmp/blockscale-reference-venv +/tmp/blockscale-reference-venv/bin/python -m pip install -r requirements.txt +CUDA_VISIBLE_DEVICES=1 CUDA_DEVICE_ORDER=PCI_BUS_ID \ + /tmp/blockscale-reference-venv/bin/python benchmark.py +``` + +The harness performs the upstream numerical reference check (`atol=0.1`, +`rtol=0.01`) before timing. It uses seed 2026, hot buffer reuse, no CUDA graphs, +three rounds, 100 warmups and 31 single-kernel samples per round. Its generated +`cutedsl-events.json` contains **CUDA-event intervals**, which include submission +gaps. Those values are not the Nsight kernel-active times in the comparison +tables. Generated reports are ignored by Git. + +This command runs CuTeDSL only; `cargo oxide run blockscale_gemm_cute` prints +the Rust executable's own CUDA-event timing. + +To measure kernel-active time, profile the same harness: + +```sh +CUDA_VISIBLE_DEVICES=1 CUDA_DEVICE_ORDER=PCI_BUS_ID BENCH_PROFILED=1 \ + nsys profile --trace=cuda,nvtx --sample=none --cpuctxsw=none \ + --output=/tmp/blockscale-cutedsl \ + /tmp/blockscale-reference-venv/bin/python benchmark.py +nsys export --type=sqlite --output=/tmp/blockscale-cutedsl.sqlite \ + /tmp/blockscale-cutedsl.nsys-rep +``` + +Use a fresh output basename for each run. The harness marks each measured +window with `cutedsl_samples_round_*`. Extract only the 31 main GEMM launches +within each window; setup, conversion, and correctness kernels are excluded: + +```sh +python3 - <<'PY' +import sqlite3 +import statistics +with sqlite3.connect('/tmp/blockscale-cutedsl.sqlite') as db: + windows = db.execute("SELECT start,end,text FROM NVTX_EVENTS " + "WHERE text LIKE 'cutedsl_samples_round_%' ORDER BY start").fetchall() + assert len(windows) == 3 + medians = [] + for start, end, name in windows: + samples = [row[0] for row in db.execute( + "SELECT (k.end-k.start)/1000.0 FROM CUPTI_ACTIVITY_KIND_KERNEL k " + "JOIN StringIds s ON s.id=k.shortName WHERE k.start>=? AND k.end<=? " + "AND s.value LIKE 'kernel_cutlass_kernel_dense%' ORDER BY k.start", (start, end))] + assert len(samples) == 31 + medians.append(statistics.median(samples)) + print(name, medians[-1], 'us kernel-active') + print('Median of round medians:', statistics.median(medians), 'us') +PY +``` + +The exact-oracle [comparison runner](../../blockscale_gemm_ctm/compare.py) accepts +the Rust and C++ binaries and the new [CTM executable harness](../../blockscale_gemm_ctm/reference/benchmark.py). +The older CuTeDSL `benchmark.py` above has a different interface. + +## Source provenance + +The two upstream Python files are unmodified copies from NVIDIA CUTLASS +`v4.6.2`, commit `6c65a175668952f09bcbf66cb97a8de1b734b4a0`, under +`examples/python/CuTeDSL/cute/blackwell_geforce/kernel/blockscaled_gemm/`. +They retain their full BSD-3-Clause notices. The harness and C++ counterpart +use cuda-oxide's Apache-2.0 license. + +| File | SHA-256 | +| --- | --- | +| `dense_blockscaled_gemm_persistent_cooperative.py` | `e3e8503239716b4d68db7dfec4971dcff7b386470c8515c3d2cc8698c761540e` | +| `blockscaled_gemm_dispatch.py` | `2d4ab5d6c150d938c0dbae3f63f5ddbbc9abf33975d66810cd6fd47a45ff3996` | diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/benchmark.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/benchmark.py new file mode 100644 index 0000000000..61627a5bfe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/benchmark.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Measure the unmodified CUTLASS v4.6.2 example with one launch per event pair.""" +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path +import statistics + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute.testing as testing +import torch +import dense_blockscaled_gemm_persistent_cooperative as example + +OUT = Path(__file__).parent + + +def checked(result): + if result[0] != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(str(result)) + return result[1] if len(result) == 2 else result[1:] + + +def single_launch_benchmark(fn, *, workspace_generator, workspace_count, + stream, warmup_iterations, iterations, **kwargs): + assert workspace_count == 1 and warmup_iterations == 100 and iterations == 31 + args = workspace_generator() + checked(cuda.cuStreamSynchronize(stream)) + runs = [] + for round_index in range(3): + for _ in range(warmup_iterations): + fn(*args.args, **args.kwargs) + checked(cuda.cuStreamSynchronize(stream)) + samples = [] + torch.cuda.nvtx.range_push(f"cutedsl_samples_round_{round_index}") + for _ in range(iterations): + start = checked(cuda.cuEventCreate(0)) + end = checked(cuda.cuEventCreate(0)) + checked(cuda.cuEventRecord(start, stream)) + fn(*args.args, **args.kwargs) + checked(cuda.cuEventRecord(end, stream)) + checked(cuda.cuEventSynchronize(end)) + samples.append(checked(cuda.cuEventElapsedTime(start, end)) * 1000) + checked(cuda.cuEventDestroy(start)) + checked(cuda.cuEventDestroy(end)) + torch.cuda.nvtx.range_pop() + ordered = sorted(samples) + run = dict(round=round_index, median_us=statistics.median(samples), + p10_us=ordered[3], p90_us=ordered[27], samples_us=samples) + runs.append(run) + print(json.dumps(run), flush=True) + report = dict(implementation="Python CuTeDSL", version=importlib.metadata.version("nvidia-cutlass-dsl"), + source_tag="v4.6.2", shape=[2048,2048,1024,1], tile=[128,128,128], + input="Float4E2M1FN", scale="Float8E8M0FNU", scale_vector=32, + output="Float16", gpu=torch.cuda.get_device_name(), + gpu_visible_device=os.getenv("CUDA_VISIBLE_DEVICES"), + warmups=100, samples=31, graphs=False, cold_l2=False, + timing="one direct launch per CUDA driver event pair", + profiler_run=os.getenv("BENCH_PROFILED") == "1", + correctness="upstream reference check passed before timing", + source_sha256=hashlib.sha256(Path(example.__file__).read_bytes()).hexdigest(), + runs=runs, median_us=statistics.median(r["median_us"] for r in runs)) + name = "cutedsl-profiled-events.json" if report["profiler_run"] else "cutedsl-events.json" + (OUT / name).write_text(json.dumps(report, indent=2) + "\n") + return report["median_us"] + + +testing.benchmark = single_launch_benchmark +torch.manual_seed(2026) +example.run_bs( + (2048,2048,1024,1), cutlass.Float4E2M1FN, cutlass.Float4E2M1FN, + cutlass.Float8E8M0FNU, 32, cutlass.Float16, cutlass.Float32, + "k", "k", "n", (128,128,128), (128,128), 0.1, 100, 31, False, False, +) +print("PASS", flush=True) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/blockscaled_gemm_dispatch.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/blockscaled_gemm_dispatch.py new file mode 100644 index 0000000000..4aaf0adfd1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/blockscaled_gemm_dispatch.py @@ -0,0 +1,284 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""SM120 block-scaled GEMM dispatch helpers shared by the cooperative and +pingpong examples in this directory. Both examples are otherwise standalone; +they only depend on this file (and the rest of the cutlass DSL library) and +not on each other. + +Dispatch table: + Same-dtype paths: + (FP4, FP4, *, Float8E4M3FN, 16) -> MmaMXF4NVF4Op, use_mxf8f6f4=False, mma_K=64 + (FP4, FP4, *, Float8E8M0FNU, 32) -> MmaMXF4Op, use_mxf8f6f4=False, mma_K=64 + (FP8, FP8, *, Float8E8M0FNU, 32) -> MmaMXF8Op, use_mxf8f6f4=True, mma_K=32 + (FP8 same-dtype is restricted to a_dtype == b_dtype.) + Mixed-precision: + (FP4, FP8, *, Float8E8M0FNU, 32) -> MmaMXF8F6F4Op, use_mxf8f6f4=True, mma_K=32 + (FP8, FP4, *, Float8E8M0FNU, 32) -> MmaMXF8F6F4Op, use_mxf8f6f4=True, mma_K=32 + (Same-width mixed-FP8 (E4M3 + E5M2) and FP6 mixed pairs are not supported.) + +`use_mxf8f6f4` is the third argument to +`cutlass.utils.blackwell_helpers.get_permutation_mnk` and selects perm_k = 32 +(FP8 / mixed warp-MMA shape (16,8,32)) instead of perm_k = 64 (FP4 shape (16,8,64)). +""" + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu.warp.mma import MXF8F6F4_SUPPORTED_PAIRS + + +# Number of bits to shift the FP4 byte left by before mma.sync.kind::mxf8f6f4. +# ldsm.b4x16_p64 places the FP4 nibble in the LOW half of the 8-bit register +# byte; mma.sync.kind::mxf8f6f4 reads the FP4 from the MIDDLE of the byte. The +# fixed `<< 2` shift moves it into position. See the C++ reference +# implementation `cute::fp4_shift_A/B` in +# cutlass/include/cute/atom/mma_traits_sm120.hpp. +FP4_SHIFT_BITS = 2 + +_FP8_DTYPES = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) + + +def make_ldmatrix_atom(operand_dtype, transpose, num_matrices=4, mixed_mode=False): + """Build a warp-level ldmatrix copy atom for the SM120 block-scaled GEMM + A/B SMEM->RMEM path. + """ + if mixed_mode and operand_dtype.width == 4: + assert not transpose, "LdMatrix8x16x8bOp does not support transpose" + return cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x16x8bOp( + transpose=False, + num_matrices=num_matrices, + unpack_bits=4, + ), + cutlass.Int8, + ) + atom_element_type = operand_dtype + if mixed_mode and operand_dtype.width < 8: + atom_element_type = cutlass.Int8 + return cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + transpose=transpose, + num_matrices=num_matrices, + ), + atom_element_type, + ) + + +def make_sm120_blockscaled_mma_op(a_dtype, b_dtype, acc_dtype, sf_dtype, sf_vec_size): + """Dispatch the SM120 block-scaled MMA op for the given operand / scale + factor combination. + + Returns + ------- + Tuple[MmaOp, bool] + (mma_op, use_mxf8f6f4) where use_mxf8f6f4 is True for FP8 / mixed + paths (atom shape (16,8,32)) and False for FP4 same-dtype paths + (atom shape (16,8,64)). `use_mxf8f6f4` is consumed by + `cutlass.utils.blackwell_helpers.get_permutation_mnk`. + + Raises + ------ + ValueError + On any unsupported combination, with a diagnostic that names the + offending combination and the supported alternatives. + """ + # Same-dtype paths use the FP4 / FP8 same-dtype atoms. + if a_dtype == b_dtype: + if a_dtype == cutlass.Float4E2M1FN: + if sf_vec_size == 16: + if sf_dtype != cutlass.Float8E4M3FN: + raise ValueError( + f"Float4E2M1FN + sf_vec_size=16 requires sf_dtype=Float8E4M3FN, " + f"got sf_dtype={sf_dtype}" + ) + return ( + cute.nvgpu.warp.MmaMXF4NVF4Op(a_dtype, acc_dtype, sf_dtype), + False, + ) + if sf_vec_size == 32: + if sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"Float4E2M1FN + sf_vec_size=32 requires sf_dtype=Float8E8M0FNU, " + f"got sf_dtype={sf_dtype}" + ) + return ( + cute.nvgpu.warp.MmaMXF4Op(a_dtype, acc_dtype, sf_dtype), + False, + ) + raise ValueError( + f"Float4E2M1FN requires sf_vec_size in (16, 32), got {sf_vec_size}" + ) + if a_dtype in _FP8_DTYPES: + if sf_vec_size != 32: + raise ValueError( + f"FP8 ab_dtype ({a_dtype}) requires sf_vec_size=32, " + f"got {sf_vec_size}" + ) + if sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"FP8 ab_dtype + sf_vec_size=32 requires sf_dtype=Float8E8M0FNU, " + f"got sf_dtype={sf_dtype}" + ) + return ( + cute.nvgpu.warp.MmaMXF8Op(a_dtype, acc_dtype, sf_dtype), + True, + ) + raise ValueError( + f"Unsupported same-dtype ab_dtype={a_dtype} for SM120 block-scaled GEMM. " + f"Supported same-dtype: Float4E2M1FN, Float8E4M3FN, Float8E5M2." + ) + # Mixed-precision FP4 x FP8 paths. + is_a_fp4_b_fp8 = a_dtype == cutlass.Float4E2M1FN and b_dtype in _FP8_DTYPES + is_a_fp8_b_fp4 = a_dtype in _FP8_DTYPES and b_dtype == cutlass.Float4E2M1FN + if is_a_fp4_b_fp8 or is_a_fp8_b_fp4: + if sf_vec_size != 32: + raise ValueError( + f"FP4 x FP8 mixed-precision requires sf_vec_size=32, got {sf_vec_size}" + ) + if sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"FP4 x FP8 mixed-precision requires sf_dtype=Float8E8M0FNU, " + f"got sf_dtype={sf_dtype}" + ) + return ( + cute.nvgpu.warp.MmaMXF8F6F4Op(a_dtype, b_dtype, acc_dtype, sf_dtype), + True, + ) + # Reject same-width mixed-FP8 (E4M3 + E5M2) and FP6 with named diagnostics. + if a_dtype in _FP8_DTYPES and b_dtype in _FP8_DTYPES: + raise ValueError( + f"same-width mixed-FP8 (a_dtype={a_dtype}, b_dtype={b_dtype}) is not supported. " + f"Supported FP4 x FP8 mixed pairs: (Float4E2M1FN, Float8E4M3FN), " + f"(Float4E2M1FN, Float8E5M2), (Float8E4M3FN, Float4E2M1FN), " + f"(Float8E5M2, Float4E2M1FN). Same-dtype FP8 uses MmaMXF8Op." + ) + raise ValueError( + f"Unsupported (a_dtype={a_dtype}, b_dtype={b_dtype}, sf_dtype={sf_dtype}, " + f"sf_vec_size={sf_vec_size}) for SM120 block-scaled GEMM. FP6 mixed pairs " + f"are not supported; supported mixed pairs are FP4 x FP8 only." + ) + + +def validate_blockscaled_args(args, fp4_allowed_tiles, fp8_allowed_tiles): + """Post-argparse validation of (a_dtype, b_dtype, sf_dtype, sf_vec_size, tile_shape_mnk). + + Same-dtype paths use the FP4 / FP8 same-dtype atoms. Mixed-precision FP4 x FP8 + is permitted only for the four supported pairs. Same-width mixed-FP8 and FP6 + are explicitly rejected with named diagnostics. + + Tile-K constraints come from the BlockScaled SF SMEM layout + (`sm120_make_smem_layout_sfa`), which requires + ``tile_K >= sf_vec_size * blk_sf == sf_vec_size * 4``: + * sf_vec_size=16 (NVFP4): tile_K must be a multiple of 64 + * sf_vec_size=32 (MXFP4 / MXFP8 / mixed): tile_K must be a multiple of 128 + A K=64 SF block cannot be filled at sf_vec_size=32 (only 2 SFs along K + fit in the K=128-required basic chunk), so tile_K=64 is rejected for + sf_vec_size=32 even though the FP4 same-dtype path otherwise allows it. + """ + tile = tuple(args.tile_shape_mnk) + a_dtype = args.a_dtype + b_dtype = args.b_dtype + # Generic sf_vec_size sanity check applies to every dtype branch below. + if args.sf_vec_size not in (16, 32): + raise ValueError( + f"--sf_vec_size must be 16 or 32, got {args.sf_vec_size}" + ) + # Mixed-precision A/B: only the four FP4 x FP8 pairs are allowed. + if a_dtype != b_dtype: + if (a_dtype, b_dtype) not in MXF8F6F4_SUPPORTED_PAIRS: + if a_dtype in _FP8_DTYPES and b_dtype in _FP8_DTYPES: + raise ValueError( + f"same-width mixed-FP8 (--a_dtype {a_dtype} --b_dtype {b_dtype}) " + f"is not supported. Supported mixed pairs: FP4 x FP8 only." + ) + raise ValueError( + f"unsupported mixed (--a_dtype {a_dtype} --b_dtype {b_dtype}). " + f"Supported mixed pairs are FP4 x FP8 only: " + f"{sorted(repr(p) for p in MXF8F6F4_SUPPORTED_PAIRS)}. " + f"FP6 mixed pairs are not supported." + ) + if args.sf_vec_size != 32: + raise ValueError( + f"FP4 x FP8 mixed-precision requires --sf_vec_size 32, " + f"got {args.sf_vec_size}" + ) + if args.sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"FP4 x FP8 mixed-precision requires --sf_dtype Float8E8M0FNU, " + f"got --sf_dtype {args.sf_dtype}" + ) + if tile not in fp8_allowed_tiles: + raise ValueError( + f"tile_shape {tile} is not supported for FP4 x FP8 mixed-precision. " + f"Allowed mixed tile shapes: {sorted(fp8_allowed_tiles)}." + ) + return + # Same-dtype paths. + if a_dtype in _FP8_DTYPES: + if args.sf_vec_size != 32: + raise ValueError( + f"FP8 a_dtype ({a_dtype}) requires --sf_vec_size 32, " + f"got {args.sf_vec_size}" + ) + if args.sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"FP8 a_dtype + sf_vec_size=32 requires --sf_dtype Float8E8M0FNU, " + f"got --sf_dtype {args.sf_dtype}" + ) + if tile not in fp8_allowed_tiles: + raise ValueError( + f"tile_shape {tile} is not supported for MXFP8. " + f"Allowed FP8 tile shapes: {sorted(fp8_allowed_tiles)}." + ) + elif a_dtype == cutlass.Float4E2M1FN: + if args.sf_vec_size == 16 and args.sf_dtype != cutlass.Float8E4M3FN: + raise ValueError( + f"FP4 + --sf_vec_size 16 requires --sf_dtype Float8E4M3FN, " + f"got {args.sf_dtype}" + ) + if args.sf_vec_size == 32 and args.sf_dtype != cutlass.Float8E8M0FNU: + raise ValueError( + f"FP4 + --sf_vec_size 32 requires --sf_dtype Float8E8M0FNU, " + f"got {args.sf_dtype}" + ) + if tile not in fp4_allowed_tiles: + raise ValueError( + f"tile_shape {tile} is not supported for FP4 path. " + f"Allowed FP4 tile shapes: {sorted(fp4_allowed_tiles)}." + ) + if args.sf_vec_size == 32 and tile[2] % 128 != 0: + raise ValueError( + f"FP4 + sf_vec_size=32 (MXFP4) requires tile_K to be a " + f"multiple of 128, " + f"got tile_K={tile[2]}." + ) + else: + raise ValueError( + f"--a_dtype must be Float4E2M1FN, Float8E4M3FN, or Float8E5M2; " + f"got {a_dtype}" + ) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/README.md b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/README.md new file mode 100644 index 0000000000..1c4a488baf --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/README.md @@ -0,0 +1,65 @@ +# MXFP4 GEMM with C++ instruction primitives + +This CUDA C++ example implements the same fixed SM120a GEMM as +[`blockscale_gemm_ctm`](../../../blockscale_gemm_ctm): M=N=2048, K=1024, +packed E2M1 inputs, UE8M0 scales per 32 K values, FP32 accumulators, and FP16 output. +It keeps the 128×128×128 tile, three-stage TMA ring, eight compute warps and one +producer warp, persistent tile schedule, and STMATRIX/TMA epilogue. The five +tensor maps are passed by value with `__grid_constant__`. + +The kernel uses narrow inline-PTX wrappers and explicit physical addresses; +it requires no CuTe algebra or CUTLASS headers. A/B fragments and scale carriers +have distinct C++ types. This is a new C++ counterpart to the Rust example, +not an upstream CUTLASS benchmark. + +Build with a CUDA toolkit supporting SM120a and a C++17 host compiler: + +```sh +python3 build.py --output-dir /tmp/blockscale-cpp +``` + +The script runs CPU layout checks and compiles the executable in a fresh +`run-*` directory. Its manifest records exact commands, source/binary hashes, +compiler versions, retained PTX/cubin intermediates, and resource usage. The +explicit `-gencode=arch=compute_120a,code=sm_120a` target avoids a generic SM120 +fallback, which cannot express the architecture-specific blockscale instruction. +`--nvcc` and `--cxx` accept executable paths or names on `PATH`; `--cxx` selects +the compiler for both CPU tests and nvcc's host compilation. + +Use the existing shared validation harness with the reported binary path: + +```sh +CUDA_VISIBLE_DEVICES=1 python3 ../../../blockscale_gemm_ctm/validate.py \ + --binary /tmp/blockscale-cpp/run-*/blockscale_gemm_cpp \ + --output-dir /tmp/blockscale-cpp-validation --sanitizer +``` + +Replace the wildcard with one exact run directory if multiple builds exist. +Each process uses the same deterministic input fixture as Rust, computes an +independent dense CPU oracle, checks all 4,194,304 FP16 bit patterns from poisoned +output, and rechecks after repeated launches. Validation covers the resident +grid, forced one-CTA and 17-CTA grids, memcheck, and synccheck. CPU tests check +physical-layout coverage, alignment, and nonaliasing; they do not require a GPU. + +The executable accepts `--ctas`, `--warmup`, `--samples`, and `--json`. Defaults +are the occupancy-based resident grid, 100 warmups, and 31 single-kernel CUDA +event samples. Use the same explicit CTA count for comparisons (170 on the +tested RTX 5090). Compare Nsight **kernel-active** durations with Rust under the +same protocol; CUDA event intervals include submission gaps and are a separate +metric. Compiler versions and generated resource usage remain relevant to the +comparison. No SOL claim is made. + +## Compare with the Rust CTM + +Build both executables freshly, then run the exact-oracle comparison runner: + +```sh +python3 ../../../blockscale_gemm_ctm/compare.py \ + --binary rust=/path/to/blockscale_gemm_ctm \ + --binary cpp=/path/to/blockscale_gemm_cpp \ + --output-dir /tmp/blockscale-comparison --device 1 \ + --ctas 170 --rounds 3 --warmup 100 --samples 31 +``` + +The runner requires 132 kernel launches and full correctness in every timing +process and reports Nsight kernel-active time, excluding host submission gaps. diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/blockscale_gemm.cu b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/blockscale_gemm.cu new file mode 100644 index 0000000000..22a0103b89 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/blockscale_gemm.cu @@ -0,0 +1,467 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// Instruction-level CUDA C++ counterpart to blockscale_gemm_ctm. +// Physical addresses, issue order, barrier phases, and ownership stay explicit. +#include +#include +#include + +#include "layout.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace primitives { +struct Mxf4A { std::uint32_t x[4]; }; +struct Mxf4B { std::uint32_t x[2]; }; +struct ScaleA { std::uint16_t bits; }; +struct ScaleB { std::uint16_t bits; }; +struct Accumulator { float x[4]; }; + +__device__ __forceinline__ unsigned shared_address(const void* pointer) { + return static_cast(__cvta_generic_to_shared(pointer)); +} +__device__ __forceinline__ void init(std::uint64_t* bar, unsigned count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" :: + "r"(shared_address(bar)), "r"(count) : "memory"); +} +__device__ __forceinline__ void wait(std::uint64_t* bar, unsigned phase) { + unsigned done; + do { + asm volatile("{ .reg .pred p; " + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], %2; " + "selp.u32 %0, 1, 0, p; }" + : "=r"(done) : "r"(shared_address(bar)), "r"(phase) : "memory"); + } while (!done); +} +__device__ __forceinline__ void arrive(std::uint64_t* bar) { + [[maybe_unused]] std::uint64_t token; + asm volatile("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(token) : + "r"(shared_address(bar)) : "memory"); + (void)token; +} +__device__ __forceinline__ void expect(std::uint64_t* bar, unsigned bytes) { + [[maybe_unused]] std::uint64_t token; + asm volatile("mbarrier.arrive.expect_tx.shared.b64 %0, [%1], %2;" : "=l"(token) : + "r"(shared_address(bar)), "r"(bytes) : "memory"); + (void)token; +} +__device__ __forceinline__ void load(void* destination, const CUtensorMap* map, + int x, int y, std::uint64_t* bar) { + asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.tile." + "mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" :: + "r"(shared_address(destination)), "l"(map), "r"(x), "r"(y), + "r"(shared_address(bar)) : "memory"); +} +__device__ __forceinline__ void store(const void* source, const CUtensorMap* map, + int x, int y) { + asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group " + "[%0, {%1, %2}], [%3];" :: "l"(map), "r"(x), "r"(y), + "r"(shared_address(source)) : "memory"); +} +__device__ __forceinline__ Mxf4A ldmatrix(const void* source) { + Mxf4A result; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(result.x[0]), "=r"(result.x[1]), "=r"(result.x[2]), "=r"(result.x[3]) + : "r"(shared_address(source)) : "memory"); + return result; +} +__device__ __forceinline__ void mma(Accumulator& c, Mxf4A a, Mxf4B b, + ScaleA sa, ScaleB sb) { + // Two UE8M0 bytes per lane represent consecutive K32 scale groups. + const std::uint16_t selector = 0; + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4.block_scale.scale_vec::2X." + "f32.e2m1.e2m1.f32.ue8m0 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " + "%10, {%12, %12}, %11, {%12, %12};" + : "+f"(c.x[0]), "+f"(c.x[1]), "+f"(c.x[2]), "+f"(c.x[3]) + : "r"(a.x[0]), "r"(a.x[1]), "r"(a.x[2]), "r"(a.x[3]), + "r"(b.x[0]), "r"(b.x[1]), "r"(static_cast(sa.bits)), + "r"(static_cast(sb.bits)), "h"(selector)); +} +__device__ __forceinline__ void stmatrix(void* destination, const Accumulator& c) { + unsigned top, bottom; + // PTX takes the high half first; the accumulator is low-column first. + asm("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(top) : "f"(c.x[1]), "f"(c.x[0])); + asm("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(bottom) : "f"(c.x[3]), "f"(c.x[2])); + asm volatile("stmatrix.sync.aligned.m8n8.x2.shared.b16 [%0], {%1, %2};" :: + "r"(shared_address(destination)), "r"(top), "r"(bottom) : "memory"); +} +__device__ __forceinline__ void wait_output() { + asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory"); +} +__device__ __forceinline__ void output_barrier() { + asm volatile("bar.sync 2, 256;" ::: "memory"); +} +} // namespace primitives + +__global__ __launch_bounds__(288) void blockscale_gemm( + const __grid_constant__ CUtensorMap a_map, + const __grid_constant__ CUtensorMap b_map, + const __grid_constant__ CUtensorMap sa_map, + const __grid_constant__ CUtensorMap sb_map, + const __grid_constant__ CUtensorMap c_map) { + using namespace blockscale; + using namespace primitives; + extern __shared__ __align__(1024) unsigned char shared[]; + const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32; + auto* full = reinterpret_cast(shared + BARRIERS); + auto* empty = full + STAGES; + if (tid == 0) { +#pragma unroll + for (int s = 0; s < STAGES; ++s) { + init(full + s, 1); + init(empty + s, 8); + } + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + __syncthreads(); + int tile = blockIdx.x; + const int stride = gridDim.x; + int slot = 0; + if (warp < 8) { + const int wm = warp % 4, wn = warp / 4; + unsigned phase = 0; + for (; tile < TOTAL_TILES; tile += stride) { + const int tile_m = tile % (M / 128), tile_n = tile / (M / 128); + Accumulator c[2][8]{}; + for (int stage = 0; stage < K_STAGES; ++stage) { + wait(full + slot, phase); + const auto* a = shared + slot * DATA_BYTES; + const auto* b = shared + B_OFFSET + slot * DATA_BYTES; + const auto* sa = reinterpret_cast(shared + SFA_OFFSET + slot * SCALE_BYTES); + const auto* sb = reinterpret_cast(shared + SFB_OFFSET + slot * SCALE_BYTES); + const int provider = lane / 4 + 8 * (lane & 1); + const unsigned sa0 = sa[scale_word(16 * wm + provider)]; + const unsigned sa1 = sa[scale_word(16 * wm + provider + 64)]; + unsigned scales_b[8]; +#pragma unroll + for (int n = 0; n < 8; ++n) + scales_b[n] = sb[scale_word(16 * wn + lane / 4 + n_offset(n))]; +#pragma unroll + for (int half = 0; half < 2; ++half) { + const auto a0 = ldmatrix(a + matrix_row(wm, half, lane)); + const auto a1 = ldmatrix(a + matrix_row(wm + 4, half, lane)); + const ScaleA as0{static_cast(sa0 >> (16 * half))}; + const ScaleA as1{static_cast(sa1 >> (16 * half))}; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const auto loaded_b = ldmatrix(b + matrix_row(wn + 2 * pair, half, lane)); + const Mxf4B b0{{loaded_b.x[0], loaded_b.x[2]}}; + const Mxf4B b1{{loaded_b.x[1], loaded_b.x[3]}}; + const ScaleB bs0{static_cast(scales_b[2 * pair] >> (16 * half))}; + const ScaleB bs1{static_cast(scales_b[2 * pair + 1] >> (16 * half))}; + mma(c[0][2 * pair], a0, b0, as0, bs0); + mma(c[0][2 * pair + 1], a0, b1, as0, bs1); + mma(c[1][2 * pair], a1, b0, as1, bs0); + mma(c[1][2 * pair + 1], a1, b1, as1, bs1); + } + } + if (lane == 0) arrive(empty + slot); + if (++slot == STAGES) { slot = 0; phase ^= 1; } + } + // The producer does not participate in these 256-thread barriers. + if (tid == 0) wait_output(); + output_barrier(); + auto* output = reinterpret_cast<__half*>(shared + OUTPUT); +#pragma unroll + for (int band = 0; band < 2; ++band) { +#pragma unroll + for (int n = 0; n < 8; ++n) + stmatrix(output + output_row(warp, lane, band, n), c[band][n]); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + output_barrier(); + if (tid == 0) { + store(shared + OUTPUT, &c_map, tile_n * 128, tile_m * 128); + store(shared + OUTPUT + 16384, &c_map, tile_n * 128 + 64, tile_m * 128); + asm volatile("cp.async.bulk.commit_group;" ::: "memory"); + } + } + if (tid == 0) wait_output(); + } else if (lane == 0) { + unsigned phase = 1; + for (; tile < TOTAL_TILES; tile += stride) { + const int tile_m = tile % (M / 128), tile_n = tile / (M / 128); + for (int stage = 0; stage < K_STAGES; ++stage) { + wait(empty + slot, phase); + expect(full + slot, TX_BYTES); + load(shared + slot * DATA_BYTES, &a_map, stage * 64, tile_m * 128, full + slot); + load(shared + B_OFFSET + slot * DATA_BYTES, &b_map, stage * 64, tile_n * 128, full + slot); + load(shared + SFA_OFFSET + slot * SCALE_BYTES, &sa_map, 0, + tile_m * K_STAGES + stage, full + slot); + load(shared + SFB_OFFSET + slot * SCALE_BYTES, &sb_map, 0, + tile_n * K_STAGES + stage, full + slot); + if (++slot == STAGES) { slot = 0; phase ^= 1; } + } + } + for (int tail = 0; tail < STAGES; ++tail) { + wait(empty + slot, phase); + if (++slot == STAGES) { slot = 0; phase ^= 1; } + } + } +} + +namespace { +void check_cuda(cudaError_t status, const char* expression) { + if (status != cudaSuccess) + throw std::runtime_error(std::string(expression) + ": " + cudaGetErrorString(status)); +} +void check_driver(CUresult status, const char* expression) { + if (status != CUDA_SUCCESS) { + const char* text = nullptr; + cuGetErrorString(status, &text); + throw std::runtime_error(std::string(expression) + ": " + (text ? text : "CUDA driver error")); + } +} +#define CUDA_CHECK(expression) check_cuda((expression), #expression) +#define DRIVER_CHECK(expression) check_driver((expression), #expression) + +template struct DeviceBuffer { + T* pointer = nullptr; + explicit DeviceBuffer(std::size_t count) { + CUDA_CHECK(cudaMalloc(reinterpret_cast(&pointer), count * sizeof(T))); + } + ~DeviceBuffer() { if (pointer) cudaFree(pointer); } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; +}; +struct Event { + cudaEvent_t value{}; + Event() { CUDA_CHECK(cudaEventCreate(&value)); } + ~Event() { cudaEventDestroy(value); } + Event(const Event&) = delete; + Event& operator=(const Event&) = delete; +}; + +CUtensorMap encode(void* pointer, CUtensorMapDataType type, std::uint64_t columns, + std::uint64_t rows, std::uint64_t row_bytes, + std::uint32_t tile_columns, std::uint32_t tile_rows, + CUtensorMapSwizzle swizzle) { + alignas(128) CUtensorMap map{}; + const std::uint64_t dimensions[] = {columns, rows}; + const std::uint64_t strides[] = {row_bytes}; + const std::uint32_t box[] = {tile_columns, tile_rows}, element_strides[] = {1, 1}; + DRIVER_CHECK(cuTensorMapEncodeTiled(&map, type, 2, pointer, dimensions, strides, + box, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return map; +} +std::uint32_t mix32(std::uint32_t x) { + x ^= x >> 16; x *= 0x7feb352dU; x ^= x >> 15; x *= 0x846ca68bU; + return x ^ (x >> 16); +} +float decode(std::uint8_t nibble) { + constexpr float values[] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + return (nibble & 8) ? -values[nibble & 7] : values[nibble & 7]; +} +std::vector pack(const std::vector& nibbles) { + std::vector result(nibbles.size() / 4); + for (std::size_t i = 0; i < result.size(); ++i) + result[i] = nibbles[4 * i] | (nibbles[4 * i + 1] << 4) | + (nibbles[4 * i + 2] << 8) | (nibbles[4 * i + 3] << 12); + return result; +} +std::vector pack_scales(const std::vector& dense, int rows) { + std::vector result(dense.size()); + for (int row = 0; row < rows; ++row) + for (int group = 0; group < blockscale::K_GROUPS; ++group) + result[blockscale::scale_offset(row, group)] = dense[row * blockscale::K_GROUPS + group]; + return result; +} +unsigned parse_count(const std::string& argument, const std::string& value) { + unsigned result = 0; + auto parsed = std::from_chars(value.data(), value.data() + value.size(), result); + if (parsed.ec != std::errc{} || parsed.ptr != value.data() + value.size()) + throw std::runtime_error(argument + " requires an unsigned integer"); + return result; +} +struct Options { + unsigned ctas = 0, warmup = 100, samples = 31; + std::filesystem::path json; +}; +Options options(int argc, char** argv) { + Options result; + for (int i = 1; i < argc; ++i) { + std::string argument = argv[i]; + if (argument != "--ctas" && argument != "--warmup" && argument != "--samples" && argument != "--json") + throw std::runtime_error("unknown option " + argument + "; use --ctas, --warmup, --samples, --json"); + if (++i == argc) throw std::runtime_error("missing value for " + argument); + if (argument == "--json") { + result.json = argv[i]; + if (result.json.empty()) throw std::runtime_error("--json path cannot be empty"); + } else { + unsigned count = parse_count(argument, argv[i]); + if (argument == "--ctas") { + if (!count || count > blockscale::TOTAL_TILES) + throw std::runtime_error("--ctas must be between 1 and 256"); + result.ctas = count; + } else if (argument == "--samples") { + if (!count) throw std::runtime_error("--samples must be positive"); + result.samples = count; + } else result.warmup = count; + } + } + return result; +} +float percentile(const std::vector& sorted, std::size_t numerator, std::size_t denominator) { + return sorted[((sorted.size() - 1) * numerator + denominator / 2) / denominator]; +} + +int run(int argc, char** argv) { + using namespace blockscale; + const auto config = options(argc, argv); + if (!config.json.empty()) std::filesystem::remove(config.json); + CUDA_CHECK(cudaSetDevice(0)); // CUDA_VISIBLE_DEVICES controls the physical GPU. + cudaDeviceProp properties{}; + CUDA_CHECK(cudaGetDeviceProperties(&properties, 0)); + if (properties.major != 12 || properties.minor != 0) + throw std::runtime_error("mxf4 primitives require an SM120 GPU"); + std::cout << "=== Block-scaled MXFP4 GEMM with C++ instruction primitives ===\n" + << "C(2048x2048) = A(2048x1024 e2m1) x B(1024x2048 e2m1), " + "ue8m0 scales per 32-element block\n"; + + std::vector a(M * K), b(N * K), sa(M * K_GROUPS), sb(N * K_GROUPS); + for (int m = 0; m < M; ++m) { + for (int k = 0; k < K; ++k) + a[m * K + k] = mix32(std::uint32_t(m) * 0x9e3779b9U ^ std::uint32_t(k) ^ 0xa511e9b3U) & 15; + for (int group = 0; group < K_GROUPS; ++group) + sa[m * K_GROUPS + group] = 126 + mix32(std::uint32_t(m) * 0xc2b2ae35U ^ std::uint32_t(group)) % 3; + } + for (int n = 0; n < N; ++n) { + for (int k = 0; k < K; ++k) + b[n * K + k] = mix32(std::uint32_t(n) * 0x85ebca6bU ^ std::uint32_t(k) ^ 0x63d83595U) & 15; + for (int group = 0; group < K_GROUPS; ++group) + sb[n * K_GROUPS + group] = 126 + mix32(std::uint32_t(n) * 0x27d4eb2fU ^ std::uint32_t(group) ^ 0x165667b1U) % 3; + } + // Dense CPU oracle independent of shared/register layouts. Products and sums + // are exact multiples of 1/16 within FP32's exact range for this fixture. + std::cout << "computing exact host reference (2048x2048x1024)...\n" << std::flush; + std::vector decoded_a(a.size()), decoded_b(b.size()); + std::transform(a.begin(), a.end(), decoded_a.begin(), decode); + std::transform(b.begin(), b.end(), decoded_b.begin(), decode); + std::vector expected(M * N); + std::size_t nonzero = 0; + for (int m = 0; m < M; ++m) { + for (int n = 0; n < N; ++n) { + float total = 0; + for (int group = 0; group < K_GROUPS; ++group) { + float dot = 0; + for (int j = 0; j < 32; ++j) { + int k = group * 32 + j; + dot += decoded_a[m * K + k] * decoded_b[n * K + k]; + } + total += std::ldexp(dot, int(sa[m * K_GROUPS + group]) + int(sb[n * K_GROUPS + group]) - 254); + } + nonzero += total != 0; + expected[m * N + n] = __half_as_ushort(__float2half_rn(total)); + } + } + auto packed_a = pack(a), packed_b = pack(b); + auto packed_sa = pack_scales(sa, M), packed_sb = pack_scales(sb, N); + DeviceBuffer da(packed_a.size()), db(packed_b.size()), dc(M * N); + DeviceBuffer dsa(packed_sa.size()), dsb(packed_sb.size()); + CUDA_CHECK(cudaMemcpy(da.pointer, packed_a.data(), packed_a.size() * 2, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(db.pointer, packed_b.data(), packed_b.size() * 2, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(dsa.pointer, packed_sa.data(), packed_sa.size(), cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(dsb.pointer, packed_sb.data(), packed_sb.size(), cudaMemcpyHostToDevice)); + std::vector got(M * N, 0x7e00); // FP16 quiet NaN poison. + CUDA_CHECK(cudaMemcpy(dc.pointer, got.data(), got.size() * 2, cudaMemcpyHostToDevice)); + const auto a_map = encode(da.pointer, CU_TENSOR_MAP_DATA_TYPE_UINT8, K / 2, M, K / 2, + 64, 128, CU_TENSOR_MAP_SWIZZLE_64B); + const auto b_map = encode(db.pointer, CU_TENSOR_MAP_DATA_TYPE_UINT8, K / 2, N, K / 2, + 64, 128, CU_TENSOR_MAP_SWIZZLE_64B); + const auto sa_map = encode(dsa.pointer, CU_TENSOR_MAP_DATA_TYPE_UINT16, 256, (M / 128) * K_STAGES, + 512, 256, 1, CU_TENSOR_MAP_SWIZZLE_NONE); + const auto sb_map = encode(dsb.pointer, CU_TENSOR_MAP_DATA_TYPE_UINT16, 256, (N / 128) * K_STAGES, + 512, 256, 1, CU_TENSOR_MAP_SWIZZLE_NONE); + const auto c_map = encode(dc.pointer, CU_TENSOR_MAP_DATA_TYPE_FLOAT16, N, M, N * 2, + 64, 128, CU_TENSOR_MAP_SWIZZLE_128B); + CUDA_CHECK(cudaFuncSetAttribute(blockscale_gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, + DYNAMIC_SMEM_BYTES)); + int active = 0; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, blockscale_gemm, THREADS, + DYNAMIC_SMEM_BYTES)); + if (!active) throw std::runtime_error("kernel has no resident CTAs on this GPU"); + const unsigned ctas = config.ctas ? config.ctas : + static_cast(std::max(1, std::min(TOTAL_TILES, active * properties.multiProcessorCount))); + std::cout << "Persistent grid: " << ctas << " CTAs\n"; + auto launch = [&] { + blockscale_gemm<<>>(a_map, b_map, sa_map, sb_map, c_map); + CUDA_CHECK(cudaGetLastError()); + }; + auto verify = [&] { + CUDA_CHECK(cudaMemcpy(got.data(), dc.pointer, got.size() * 2, cudaMemcpyDeviceToHost)); + std::size_t errors = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + if (got[i] != expected[i]) { + if (errors < 5) + std::cerr << "C(" << i / N << ',' << i % N << "): expected 0x" << std::hex + << expected[i] << ", got 0x" << got[i] << std::dec << '\n'; + ++errors; + } + } + if (errors) throw std::runtime_error(std::to_string(errors) + " FP16 bit mismatches"); + }; + launch(); + verify(); + std::cout << "Numeric check: all " << M * N << " FP16 output bit patterns exact (" + << nonzero << " nonzero FP32 accumulators), " << K_GROUPS << " K-blocks\n"; + for (unsigned i = 0; i < config.warmup; ++i) launch(); + CUDA_CHECK(cudaDeviceSynchronize()); + std::vector times; + times.reserve(config.samples); + for (unsigned i = 0; i < config.samples; ++i) { + Event start, end; + CUDA_CHECK(cudaEventRecord(start.value)); + launch(); + CUDA_CHECK(cudaEventRecord(end.value)); + CUDA_CHECK(cudaEventSynchronize(end.value)); + float ms = 0; + CUDA_CHECK(cudaEventElapsedTime(&ms, start.value, end.value)); + times.push_back(ms); + } + verify(); + std::sort(times.begin(), times.end()); + const float median = percentile(times, 1, 2), p10 = percentile(times, 1, 10), p90 = percentile(times, 9, 10); + const double tflops = 2.0 * M * N * K / (median * 1.e9); + std::cout << std::fixed << std::setprecision(6) << "Single-kernel device time: median=" << median + << " ms, p10=" << p10 << " ms, p90=" << p90 << " ms (" << config.samples + << " samples, " << config.warmup << " warmups); " << tflops << " TFLOP/s\n"; + if (!config.json.empty()) { + std::ofstream report(config.json); + report.exceptions(std::ios::failbit | std::ios::badbit); + report << std::setprecision(9) + << "{\n \"verification\": \"bit-exact-pass\",\n \"mnk\": [2048, 2048, 1024],\n" + " \"compute_capability\": [12, 0],\n \"ctas\": " << ctas + << ",\n \"warmup\": " << config.warmup << ",\n \"samples\": " << config.samples + << ",\n \"median_ms\": " << median << ",\n \"p10_ms\": " << p10 + << ",\n \"p90_ms\": " << p90 << ",\n \"tflops\": " << tflops + << ",\n \"checked_outputs\": 4194304,\n \"replay_verified\": true\n}\n"; + report.close(); + } + std::cout << "SUCCESS: block-scaled MXFP4 GEMM is bit-exact through C++ instruction primitives\n"; + return 0; +} +} // namespace + +int main(int argc, char** argv) { + try { return run(argc, argv); } + catch (const std::exception& error) { + std::cerr << "FAILED: " << error.what() << '\n'; + return 1; + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/build.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/build.py new file mode 100644 index 0000000000..4eb3864341 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/build.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Build the SM120a C++ primitives example and record reproducible evidence; no GPU work.""" +import argparse +import hashlib +import json +from pathlib import Path +import shutil +import subprocess +import tempfile + + +def sha256(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--nvcc", default=shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc", + help="CUDA compiler executable (path or command on PATH)") + parser.add_argument("--cxx", default="c++", + help="host compiler for both CPU tests and nvcc -ccbin") + args = parser.parse_args() + for option in ("nvcc", "cxx"): + executable = shutil.which(getattr(args, option)) + if executable is None: + parser.error(f"--{option}: executable not found: {getattr(args, option)}") + setattr(args, option, str(Path(executable).resolve(strict=True))) + source = Path(__file__).resolve().parent + args.output_dir.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="run-", dir=args.output_dir.resolve())) + keep = directory / "keep" + keep.mkdir() + manifest = {"status": "building", "directory": str(directory), "commands": [], + "compilers": {"cuda": args.nvcc, "host": args.cxx}, + "sources": {name: sha256(source / name) for name in + ("blockscale_gemm.cu", "layout.cuh", "test_layout.cpp", "build.py")}} + manifest_path = directory / "manifest.json" + + def record(): + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + + def run(command, log_name): + manifest["commands"].append(command) + record() + result = subprocess.run(command, cwd=directory, text=True, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + (directory / log_name).write_text(result.stdout) + if result.returncode: + raise RuntimeError(f"exit {result.returncode}: {command}\n{result.stdout}") + return result.stdout + + record() + try: + manifest["nvcc_version"] = run([args.nvcc, "--version"], "nvcc-version.log") + manifest["cxx_version"] = run([args.cxx, "--version"], "cxx-version.log") + host_test = directory / "test_layout" + run([args.cxx, "-std=c++17", "-O2", "-Wall", "-Wextra", "-Werror", + str(source / "test_layout.cpp"), "-o", str(host_test)], "test-build.log") + manifest["host_tests"] = run([str(host_test)], "test-layout.log").strip() + binary = directory / "blockscale_gemm_cpp" + # An explicit gencode avoids nvcc's additional generic compute_120 + # fallback, which cannot represent architecture-specific MXFP4 MMA. + run([args.nvcc, "-ccbin", args.cxx, "-std=c++17", "-O3", + "-gencode=arch=compute_120a,code=sm_120a", "-lineinfo", "-Xptxas=-v", + "--keep", "--keep-dir", str(keep), str(source / "blockscale_gemm.cu"), + "-lcuda", "-o", str(binary)], "build.log") + cuobjdump = str(Path(args.nvcc).resolve().parent / "cuobjdump") + manifest["resources"] = run([cuobjdump, "--dump-resource-usage", str(binary)], "resources.log") + run([cuobjdump, "--dump-sass", str(binary)], "kernel.sass") + # These are the exact intermediate files consumed by this executable's build. + manifest["artifacts"] = {str(path.relative_to(directory)): sha256(path) + for path in sorted(keep.iterdir()) + if path.suffix in (".ptx", ".cubin")} + manifest.update(status="compiled", binary=str(binary), binary_sha256=sha256(binary)) + record() + except (OSError, RuntimeError) as error: + manifest.update(status="failed", error=str(error)) + record() + print(f"FAIL: {error}\nEvidence: {manifest_path}") + return 1 + print(f"PASS CPU physical-layout checks and SM120a compilation\nBinary: {binary}\nEvidence: {manifest_path}") + print(manifest["resources"]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/layout.cuh b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/layout.cuh new file mode 100644 index 0000000000..636bb050d9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/layout.cuh @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#if defined(__CUDACC__) +#define BLOCKSCALE_HD __host__ __device__ +#else +#define BLOCKSCALE_HD +#endif + +namespace blockscale { +constexpr int M = 2048, N = 2048, K = 1024; +constexpr int K_GROUPS = K / 32, K_STAGES = K / 128; +constexpr int THREADS = 288, TOTAL_TILES = (M / 128) * (N / 128); +constexpr int STAGES = 3, DATA_BYTES = 8192, SCALE_BYTES = 512; +constexpr int B_OFFSET = 24576, SFA_OFFSET = 49152, SFB_OFFSET = 50688; +constexpr int BARRIERS = 52224, OUTPUT = 53248, DYNAMIC_SMEM_BYTES = 86016; +constexpr unsigned TX_BYTES = 2 * DATA_BYTES + 2 * SCALE_BYTES; +static_assert(B_OFFSET == STAGES * DATA_BYTES); +static_assert(SFA_OFFSET == B_OFFSET + STAGES * DATA_BYTES); +static_assert(SFB_OFFSET == SFA_OFFSET + STAGES * SCALE_BYTES); +static_assert(BARRIERS == SFB_OFFSET + STAGES * SCALE_BYTES); +static_assert(OUTPUT >= BARRIERS + 2 * STAGES * 8); +static_assert(OUTPUT % 1024 == 0 && OUTPUT + 2 * 16384 == DYNAMIC_SMEM_BYTES); +static_assert(DATA_BYTES % 512 == 0 && SCALE_BYTES % 128 == 0); + +BLOCKSCALE_HD constexpr int scale_word(int row) { + return (row & 31) * 4 + ((row >> 5) & 3); +} +BLOCKSCALE_HD constexpr int scale_offset(int row, int group) { + return (row / 128) * (K_GROUPS / 4) * 512 + (group / 4) * 512 + + scale_word(row) * 4 + group % 4; +} +BLOCKSCALE_HD constexpr int matrix_row(int atom, int half, int lane) { + int sub = lane / 8; + int row = 16 * atom + 8 * (sub % 2) + lane % 8; + int column = 16 * half + 8 * (sub / 2); + int plain = row * 64 + column * 2; + return plain ^ ((plain & 0x180) >> 3); +} +BLOCKSCALE_HD constexpr int n_offset(int slot) { + return 32 * (slot / 2) + 8 * (slot % 2); +} +// Return FP16 element offsets; matrix_row returns byte offsets. +BLOCKSCALE_HD constexpr int output_element(int row, int column) { + int plain = row * 64 + column % 64; + return (column / 64) * 8192 + (plain ^ ((plain & 0x1c0) >> 3)); +} +BLOCKSCALE_HD constexpr int output_row(int warp, int lane, int band, int slot) { + return output_element(16 * (warp % 4) + 64 * band + (lane & 15), + 16 * (warp / 4) + n_offset(slot)); +} +} // namespace blockscale +#undef BLOCKSCALE_HD diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/test_layout.cpp b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/test_layout.cpp new file mode 100644 index 0000000000..908f795844 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/cpp/test_layout.cpp @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "layout.cuh" + +#include +#include +#include + +int main() { + using namespace blockscale; + // The complete scale tensor must be a permutation of the dense input bytes, + // including boundaries between 128-row tiles and four-K-group atoms. + std::array scales{}; + for (int row = 0; row < M; ++row) { + for (int group = 0; group < K / 32; ++group) { + auto offset = scale_offset(row, group); + assert(offset >= 0 && static_cast(offset) < scales.size() && !scales[offset]); + scales[offset] = true; + } + } + for (bool visited : scales) assert(visited); + assert(scale_offset(0, 0) == 0); + assert(scale_offset(32, 0) == 4); + assert(scale_offset(128, 0) == 4096); + + // Each ldmatrix address provider names one distinct aligned 16-byte row. + for (int half = 0; half < 2; ++half) { + std::array matrix{}; + for (int atom = 0; atom < 8; ++atom) { + for (int lane = 0; lane < 32; ++lane) { + auto offset = matrix_row(atom, half, lane); + assert(offset < DATA_BYTES && offset % 16 == 0); + assert(!matrix[offset / 16]); + matrix[offset / 16] = true; + } + } + } + + // Independently enumerate the PTX accumulator's logical (row,column) + // ownership and ensure the output swizzle neither aliases nor loses values. + std::array output{}; + constexpr int columns[] = {0, 8, 32, 40, 64, 72, 96, 104}; + for (int warp = 0; warp < 8; ++warp) { + for (int lane = 0; lane < 32; ++lane) { + for (int band = 0; band < 2; ++band) { + for (int slot = 0; slot < 8; ++slot) { + for (int reg = 0; reg < 4; ++reg) { + int row = 16 * (warp % 4) + 64 * band + lane / 4 + 8 * (reg / 2); + int col = 16 * (warp / 4) + columns[slot] + 2 * (lane % 4) + reg % 2; + assert(row < 128 && col < 128); + auto offset = output_element(row, col); + assert(offset >= 0 && static_cast(offset) < output.size() && !output[offset]); + output[offset] = true; + } + int provider_row = 16 * (warp % 4) + 64 * band + (lane & 15); + int provider_col = 16 * (warp / 4) + columns[slot]; + assert(output_row(warp, lane, band, slot) == + output_element(provider_row, provider_col)); + } + } + } + } + for (bool visited : output) assert(visited); + std::cout << "PASS physical scale, matrix, and accumulator mappings\n"; +} diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/dense_blockscaled_gemm_persistent_cooperative.py b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/dense_blockscaled_gemm_persistent_cooperative.py new file mode 100644 index 0000000000..08bea8cebe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/dense_blockscaled_gemm_persistent_cooperative.py @@ -0,0 +1,2037 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +from typing import Tuple, Type, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync +import cutlass.cute.testing as testing +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.cute.runtime import from_dlpack +import cutlass.utils.hopper_helpers as sm90_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.blackwell_helpers as sm120_utils + +try: + from blackwell_geforce.kernel.blockscaled_gemm.blockscaled_gemm_dispatch import ( + FP4_SHIFT_BITS, + make_ldmatrix_atom, + make_sm120_blockscaled_mma_op, + validate_blockscaled_args, + ) +except ImportError: + from blockscaled_gemm_dispatch import ( + FP4_SHIFT_BITS, + make_ldmatrix_atom, + make_sm120_blockscaled_mma_op, + validate_blockscaled_args, + ) + +""" +A high-performance batched dense blockscaled GEMM (C = A*SF_A * B*SF_B) example for the NVIDIA Blackwell Geforce architecture using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can only be row-major("K") +- Matrix B is NxKxL, L is batch dimension, B can only be column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can only be row-major("N") +- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively +- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively +- Source formats for matrices A and B: The only supported source format in this example is E2M1. +- Source formats for matrices SF_A and SF_B are controlled separately. With sf_vec_size=32, the only supported source format is E8. With sf_vec_size=16, the supported source formats are E8 and E4M3. + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes warp-level block-scaled MMA for matrix multiply-accumulate (MMA) operations + - Supports persistent tile scheduling to better overlap memory load/store with MMA between tiles + - Supports warp specialization to avoid explicit pipelining between mainloop load and MMA + - Uses a cooperative schedule: both warp groups work together to process a single output tile. + The two warp groups collaboratively execute the MMA mainloop, each computing a portion of + the tile, and then jointly perform the epilogue. This increases computational throughput + per tile by leveraging both warp groups simultaneously. + +This GEMM works as follows: +1. DMA warp group: + - Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. + - Load scale factor A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp groups (both warp groups cooperate on the same tile): + - Load A/B from shared memory (SMEM) to registers (RMEM) using ldmatrix instruction. + - Load scale factor A/B from shared memory (SMEM) to registers (RMEM) using universal copy. + - Perform matrix multiply-accumulate (MMA) operations using warp-level block-scaled MMA instruction. + - Store C matrix from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM) with TMA operations. + Note: Both MMA warp groups jointly handle MMA and epilogue for the same tile. + +Warp-level block-scaled MMA instructions operate as follows: +- Set matrix scale factor A/B from registers +- Read matrix A/B from registers +- Perform MMA operation and store the result in Accumulator(register) + +To run this example: + +.. code-block:: bash + + python examples/cute/blackwell_geforce/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent_cooperative.py \ + --mnkl 1024,1024,1024,1 --tile_shape_mnk 128,128,128 \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN \ + --c_dtype Float16 --acc_dtype Float32 \ + --sf_dtype Float8E4M3FN --sf_vec_size 16 + +The above example command compute batched gemm with M=1024, N=1024, K=1024, +batch_count=1. The tile shape is 128x128x128 and the cluster shape is (1,1). +The input, mma accumulator and output data type are set as fp4, fp32 +and fp16, respectively. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/cute/blackwell_geforce/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent_cooperative.py \ + --mnkl 1024,1024,1024,1 --tile_shape_mnk 128,128,128 \ + --a_dtype Float4E2M1FN --b_dtype Float4E2M1FN \ + --c_dtype Float16 --acc_dtype Float32 \ + --sf_dtype Float8E4M3FN --sf_vec_size 16 + +Constraints: +* Supported input data types: Float4E2M1FN +* Only Float32 accumulation is supported in FP4 mma +* CTA tile shape M/N/K: + - tile_shape_m should be divisible by 128 + - tile_shape_n should be divisible by 128 + - tile_shape_k should be divisible by 64 (sf_vec_size=16) or 128 (sf_vec_size=32) +* Cluster shape M/N must be [1, 1] for Blackwell Gefore +""" + + +# ///////////////////////////////////////////////////////////////////////////// +# Helpers to parse args +# ///////////////////////////////////////////////////////////////////////////// +def parse_comma_separated_ints(s: str): + try: + return tuple([int(x.strip()) for x in s.split(",")]) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + +# ///////////////////////////////////////////////////////////////////////////// +# Host setup and device kernel launch +# ///////////////////////////////////////////////////////////////////////////// + + +class Sm120BlockScaledGemmKernel: + def __init__( + self, + acc_dtype, + sf_vec_size, + tile_shape_mnk, + epi_tile, + ): + self.acc_dtype = acc_dtype + self.sf_vec_size = sf_vec_size + self.cluster_shape_mnk = (1, 1, 1) + self.tile_shape_mnk = tuple(tile_shape_mnk) + self.epi_tile = tuple(epi_tile) + self.tiled_mma = None + + self.occupancy = 1 + self.num_mma_warps = 8 + self.tma_load_warp_id = self.num_mma_warps + self.num_threads_per_warp = 32 + self.threads_per_cta = ( + self.num_mma_warps + 1 # 1 warp for DMA + ) * self.num_threads_per_warp + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_120") + + self.ab_stage = None + self.epi_stage = None + + self.a_smem_layout_staged = None + self.b_smem_layout_staged = None + self.epi_smem_layout_staged = None + + self.buffer_align_bytes = 1024 + + self.mma_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.num_mma_warps * self.num_threads_per_warp, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.num_mma_warps * self.num_threads_per_warp, + ) + self.load_register_requirement = 40 + self.mma_register_requirement = 232 + + def _setup_attributes(self): + mma_op, use_mxf8f6f4 = make_sm120_blockscaled_mma_op( + self.a_dtype, + self.b_dtype, + self.acc_dtype, + self.sf_dtype, + self.sf_vec_size, + ) + self.mixed_mode = self.a_dtype != self.b_dtype + # a_fp4_in_mixed / b_fp4_in_mixed: this side carries an FP4 operand in the + # mixed FP4 x FP8 mode, so SMEM/TMA see Int8 storage and the mma.sync + # consumer needs the LDSM b4x16_p64 unpack + register `<< FP4_SHIFT_BITS`. + a_fp4_in_mixed = self.mixed_mode and self.a_dtype.width < 8 + b_fp4_in_mixed = self.mixed_mode and self.b_dtype.width < 8 + self.smem_alloc_a_dtype = cutlass.Int8 if a_fp4_in_mixed else self.a_dtype + self.smem_alloc_b_dtype = cutlass.Int8 if b_fp4_in_mixed else self.b_dtype + # `internal_type` for `_make_tma_atoms_and_tensors` — None when the dtype + # already matches (TMA sees the native dtype), Int8 when we recast for FP4. + self.tma_internal_a_dtype = cutlass.Int8 if a_fp4_in_mixed else None + self.tma_internal_b_dtype = cutlass.Int8 if b_fp4_in_mixed else None + atom_shape = (4, 2, 1) + atom_layout = cute.make_layout(atom_shape) + permutation_mnk = sm120_utils.get_permutation_mnk( + self.tile_shape_mnk, self.sf_vec_size, use_mxf8f6f4 + ) + self.tiled_mma = cute.make_tiled_mma( + mma_op, + atom_layout, + permutation_mnk=permutation_mnk, + ) + + self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk) + + # Compute the smem size of SFA/SFB + sfa_smem_layout_per_stage = blockscaled_utils.sm120_make_smem_layout_sfa( + self.tiled_mma, + self.tile_shape_mnk, + self.sf_vec_size, + 1, + ) + + sfb_smem_layout_per_stage = blockscaled_utils.sm120_make_smem_layout_sfb( + self.tiled_mma, + self.tile_shape_mnk, + self.sf_vec_size, + 1, + ) + + # Compute stage before compute smem layout + self.ab_stage, self.epi_stage = self._compute_stages( + self.tile_shape_mnk, + self.smem_alloc_a_dtype, + self.smem_alloc_b_dtype, + self.sf_dtype, + sfa_smem_layout_per_stage, + sfb_smem_layout_per_stage, + self.epi_tile, + self.c_dtype, + self.smem_capacity, + self.occupancy, + ) + + assert self.epi_stage > 0, ( + "epi_stage <= 0, no enough shared memory. This case will be skipped." + ) + + ( + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.epi_smem_layout_staged, + ) = self._make_smem_layouts( + self.tile_shape_mnk, + self.epi_tile, + self.smem_alloc_a_dtype, + self.a_layout, + self.smem_alloc_b_dtype, + self.b_layout, + self.ab_stage, + self.c_dtype, + self.c_layout, + self.epi_stage, + self.sf_vec_size, + self.tiled_mma, + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + sfa: cute.Tensor, + sfb: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Execute the GEMM operation in steps: + - Setup static attributes + - Setup TMA load/store atoms and tensors + - Compute grid size + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + """ + + # setup static attributes before smem/grid/tma computation + self.a_dtype = a.element_type + self.b_dtype = b.element_type + self.c_dtype = c.element_type + self.sf_dtype = sfa.element_type + + self.a_layout = utils.LayoutEnum.from_tensor(a) + self.b_layout = utils.LayoutEnum.from_tensor(b) + self.c_layout = utils.LayoutEnum.from_tensor(c) + + self._setup_attributes() + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + self.sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a.shape, self.sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa.iterator, self.sfa_layout) + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + self.sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b.shape, self.sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb.iterator, self.sfb_layout) + + tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors( + a, + self.a_smem_layout_staged, + (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + 1, + internal_type=self.tma_internal_a_dtype, + ) + + tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors( + b, + self.b_smem_layout_staged, + (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), + 1, + internal_type=self.tma_internal_b_dtype, + ) + + tma_atom_sfa, tma_tensor_sfa = self._make_tma_atoms_and_tensors( + sfa_tensor, + self.sfa_smem_layout_staged, + (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + 1, + internal_type=cutlass.Int16, + ) + + tma_atom_sfb, tma_tensor_sfb = self._make_tma_atoms_and_tensors( + sfb_tensor, + self.sfb_smem_layout_staged, + (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), + 1, + internal_type=cutlass.Int16, + ) + + tma_atom_c, tma_tensor_c = self._make_tma_store_atoms_and_tensors( + c, + self.epi_smem_layout_staged, + self.epi_tile, + ) + + tile_sched_params, grid = self._compute_grid( + c, + self.tile_shape_mnk, + max_active_clusters, + ) + + @cute.struct + class SharedStorage: + mainloop_pipeline_array_ptr: cute.struct.MemRange[ + cutlass.Int64, self.ab_stage * 2 + ] + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_a_dtype, cute.cosize(self.a_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.smem_alloc_b_dtype, cute.cosize(self.b_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, cute.cosize(self.epi_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.kernel( + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + self.tiled_mma, + self.cta_layout_mnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.epi_smem_layout_staged, + tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=[1, 1, 1], + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tiled_mma: cute.TiledMma, + cta_layout_mnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + epi_smem_layout_staged: cute.ComposedLayout, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the batched GEMM computation. + + :param tma_atom_a: TMA copy atom for A tensor + :type tma_atom_a: cute.CopyAtom + :param mA_mkl: Input tensor A + :type mA_mkl: cute.Tensor + :param tma_atom_b: TMA copy atom for B tensor + :type tma_atom_b: cute.CopyAtom + :param mB_nkl: Input tensor B + :type mB_nkl: cute.Tensor + :param tma_atom_c: TMA copy atom for C tensor + :type tma_atom_c: cute.CopyAtom + :param mC_mnl: Output tensor C + :type mC_mnl: cute.Tensor + :param tiled_mma: Tiled MMA object + :type tiled_mma: cute.TiledMma + :param cta_layout_mnk: CTA layout + :type cta_layout_mnk: cute.Layout + :param a_smem_layout_staged: Shared memory layout for A + :type a_smem_layout_staged: cute.ComposedLayout + :param b_smem_layout_staged: Shared memory layout for B + :type b_smem_layout_staged: cute.ComposedLayout + :param epi_smem_layout_staged: Shared memory layout for epilogue + :type epi_smem_layout_staged: cute.ComposedLayout + """ + + # /////////////////////////////////////////////////////////////////////////////// + # Get cta/warp/thread idx + # /////////////////////////////////////////////////////////////////////////////// + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + # bidx, bidy, bidz = cute.arch.block_idx() + # bdimx, bdimy, bdimz = cute.arch.grid_dim() + + # ///////////////////////////////////////////////////////////////////////////// + # Prefetch Tma desc + # ///////////////////////////////////////////////////////////////////////////// + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + cluster_coord_mnk = cta_layout_mnk.get_flat_coord(cta_rank_in_cluster) + + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0)) + sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, 0)) + sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, 0)) + tma_copy_bytes = ( + cute.size_in_bytes(self.a_dtype, a_smem_layout) + + cute.size_in_bytes(self.b_dtype, b_smem_layout) + + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + ) + + # ///////////////////////////////////////////////////////////////////////////// + # Alloc and init AB full/empty + ACC full mbar (pipeline) + # ///////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # mbar arrays + mainloop_pipeline_array_ptr = storage.mainloop_pipeline_array_ptr.data_ptr() + + # Threads/warps participating in this pipeline + mainloop_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread + ) + mainloop_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mma_warps + ) + + cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape)) + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.ab_stage, + producer_group=mainloop_pipeline_producer_group, + consumer_group=mainloop_pipeline_consumer_group, + tx_count=tma_copy_bytes, + barrier_storage=mainloop_pipeline_array_ptr, + cta_layout_vmnk=cta_layout_vmnk, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_arrive_relaxed() + + # /////////////////////////////////////////////////////////////////////////////// + # Generate smem tensor A/B + # /////////////////////////////////////////////////////////////////////////////// + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + sC = storage.sC.get_tensor( + epi_smem_layout_staged.outer, swizzle=epi_smem_layout_staged.inner + ) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # /////////////////////////////////////////////////////////////////////////////// + # Local_tile partition global tensors + # /////////////////////////////////////////////////////////////////////////////// + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + # (tM, tK, loopM, loopK, loopL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (tN, tK, loopN, loopK, loopL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.tile_shape_mnk, (0, None, None)), + (None, None, None), + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, + cute.slice_(self.tile_shape_mnk, (None, None, 0)), + (None, None, None), + ) + + # ////////////////////////////////////////////////////////////////////////////// + # Partition global tensor for TiledMMA_A/B/C + # ////////////////////////////////////////////////////////////////////////////// + thr_mma = tiled_mma.get_slice(tidx) + + # ////////////////////////////////////////////////////////////////////////////// + # Partition shared tensor for TMA load A/B + # ////////////////////////////////////////////////////////////////////////////// + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape) + a_cta_crd = cluster_coord_mnk[1] + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + a_cta_crd, + a_cta_layout, + cute.group_modes(sA, 0, 2), + cute.group_modes(gA_mkl, 0, 2), + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape) + b_cta_crd = cluster_coord_mnk[0] + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + b_cta_crd, + b_cta_layout, + cute.group_modes(sB, 0, 2), + cute.group_modes(gB_nkl, 0, 2), + ) + + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + a_cta_crd, + a_cta_layout, + cute.group_modes(sSFA, 0, 2), + cute.group_modes(gSFA_mkl, 0, 2), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + b_cta_crd, + b_cta_layout, + cute.group_modes(sSFB, 0, 2), + cute.group_modes(gSFB_nkl, 0, 2), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # Make frangments + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrSFA = sm120_utils.partition_fragment_SFA(sSFA[None, None, 0], thr_mma, tidx) + tCrSFB = sm120_utils.partition_fragment_SFB(sSFB[None, None, 0], thr_mma, tidx) + # Keep residual K modes nested to match the C++ SM120 block-scaled mainloop. + tCrSFA = cute.group_modes(tCrSFA, 2, cute.rank(tCrSFA)) + tCrSFB = cute.group_modes(tCrSFB, 2, cute.rank(tCrSFB)) + + tCgC = thr_mma.partition_C(gC_mnl) + acc_shape = tCgC.shape[:3] + accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + + # cluster wait for barrier init + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_wait() + else: + cute.arch.sync_threads() + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Create the tile scheduler + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # Create the pipeline states for producer and consumer + mainloop_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.ab_stage + ) + mainloop_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.ab_stage + ) + + # MMA warp group + if warp_idx < self.num_mma_warps: + cute.arch.setmaxregister_increase(self.mma_register_requirement) + + num_k_blocks = cute.size(tCrA, mode=[2]) + + # /////////////////////////////////////////////////////////////////////////////// + # Copy Atom A/B retiling for TMA load A/B + # /////////////////////////////////////////////////////////////////////////////// + atom_copy_ldmatrix_A = make_ldmatrix_atom( + self.a_dtype, + transpose=self.a_layout.is_m_major_a(), + num_matrices=4, + mixed_mode=self.mixed_mode, + ) + atom_copy_ldmatrix_B = make_ldmatrix_atom( + self.b_dtype, + transpose=self.b_layout.is_n_major_b(), + num_matrices=4, + mixed_mode=self.mixed_mode, + ) + smem_tiled_copy_A = cute.make_tiled_copy_A(atom_copy_ldmatrix_A, tiled_mma) + smem_tiled_copy_B = cute.make_tiled_copy_B(atom_copy_ldmatrix_B, tiled_mma) + + atom_copy_ldmatrix_SF = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.sf_dtype, + ) + smem_tiled_copy_SFA = cute.make_tiled_copy( + atom_copy_ldmatrix_SF, + sm120_utils.get_layoutSFA_TV(tiled_mma), + ( + cute.size(tiled_mma.permutation_mnk[0]), + cute.size(tiled_mma.permutation_mnk[2]), + ), + ) + smem_tiled_copy_SFB = cute.make_tiled_copy( + atom_copy_ldmatrix_SF, + sm120_utils.get_layoutSFB_TV(tiled_mma), + ( + cute.size(tiled_mma.permutation_mnk[1]), + cute.size(tiled_mma.permutation_mnk[2]), + ), + ) + + thr_copy_ldmatrix_A = smem_tiled_copy_A.get_slice(tidx) + thr_copy_ldmatrix_B = smem_tiled_copy_B.get_slice(tidx) + tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) + tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) + tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB) + tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB) + + thr_copy_ldmatrix_SFA = smem_tiled_copy_SFA.get_slice(tidx) + thr_copy_ldmatrix_SFB = smem_tiled_copy_SFB.get_slice(tidx) + tCsSFA_copy_view = thr_copy_ldmatrix_SFA.partition_S(sSFA) + tCrSFA_copy_view = thr_copy_ldmatrix_SFA.retile(tCrSFA) + tCsSFB_copy_view = thr_copy_ldmatrix_SFB.partition_S(sSFB) + tCrSFB_copy_view = thr_copy_ldmatrix_SFB.retile(tCrSFB) + + epi_buffer = cutlass.Int32(0) + + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] + # Clear the accumulator + accumulators.fill(0.0) + + # ///////////////////////////////////////////////////////////////////////////// + # Pipelined MAINLOOP + # ///////////////////////////////////////////////////////////////////////////// + + mainloop_consumer_state.reset_count() + + peek_ab_full_status = cutlass.Boolean(1) + if mainloop_consumer_state.count < k_tile_cnt: + peek_ab_full_status = mainloop_pipeline.consumer_try_wait( + mainloop_consumer_state + ) + + # Wait for TMA copies to complete + mainloop_pipeline.consumer_wait( + mainloop_consumer_state, peek_ab_full_status + ) + # tCsA_p: (MMA, (4, MMA_M / 4), MMA_K), tCsA_p: (MMA, (4, MMA_N / 4), MMA_K) + tCsA_p = tCsA_copy_view[None, None, None, mainloop_consumer_state.index] + tCsB_p = tCsB_copy_view[None, None, None, mainloop_consumer_state.index] + tCsSFA_p = tCsSFA_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + tCsSFB_p = tCsSFB_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, 0], + tCrA_copy_view[None, None, 0], + ) + cute.copy( + smem_tiled_copy_B, + tCsB_p[None, None, 0], + tCrB_copy_view[None, None, 0], + ) + + tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p) + tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p) + tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view) + tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view) + + cute.copy( + smem_tiled_copy_SFA, + tCsSFA_p_filtered[None, None, 0], + tCrSFA_copy_view_filtered[None, None, 0], + ) + cute.copy( + smem_tiled_copy_SFB, + tCsSFB_p_filtered[None, None, 0], + tCrSFB_copy_view_filtered[None, None, 0], + ) + + for k_tile in range(0, k_tile_cnt - 1, 1, unroll=1): + # unroll the loop + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_next = ( + 0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1 + ) + + if k_block_idx == num_k_blocks - 1: + mainloop_pipeline.consumer_release(mainloop_consumer_state) + mainloop_consumer_state.advance() + + peek_ab_full_status = cutlass.Boolean(1) + peek_ab_full_status = mainloop_pipeline.consumer_try_wait( + mainloop_consumer_state + ) + + tCsA_p = tCsA_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + tCsB_p = tCsB_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + tCsSFA_p = tCsSFA_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + tCsSFB_p = tCsSFB_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + mainloop_pipeline.consumer_wait( + mainloop_consumer_state, peek_ab_full_status + ) + + # Mixed FP4 x FP8 register-side bit shift before mma.sync + # to move the FP4 nibble (in low half of each byte after + # b4x16_p64 ldmatrix) into the middle as mxf8f6f4 expects. + if cutlass.const_expr( + self.mixed_mode and self.a_dtype.width < 8 + ): + a_view = cute.recast_tensor( + tCrA[None, None, k_block_idx], cutlass.Int8 + ) + for _i in cutlass.range_constexpr(cute.size(a_view)): + a_view[_i] = cutlass.Int8( + a_view[_i] << FP4_SHIFT_BITS + ) + if cutlass.const_expr( + self.mixed_mode and self.b_dtype.width < 8 + ): + b_view = cute.recast_tensor( + tCrB[None, None, k_block_idx], cutlass.Int8 + ) + for _i in cutlass.range_constexpr(cute.size(b_view)): + b_view[_i] = cutlass.Int8( + b_view[_i] << FP4_SHIFT_BITS + ) + cute.gemm( + tiled_mma, + accumulators, + [ + tCrA[None, None, k_block_idx], + tCrSFA[None, None, k_block_idx], + ], + [ + tCrB[None, None, k_block_idx], + tCrSFB[None, None, k_block_idx], + ], + accumulators, + ) + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + cute.copy( + smem_tiled_copy_B, + tCsB_p[None, None, k_block_next], + tCrB_copy_view[None, None, k_block_next], + ) + + tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p) + tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p) + tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view) + tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view) + cute.copy( + smem_tiled_copy_SFA, + tCsSFA_p_filtered[None, None, k_block_next], + tCrSFA_copy_view_filtered[None, None, k_block_next], + ) + cute.copy( + smem_tiled_copy_SFB, + tCsSFB_p_filtered[None, None, k_block_next], + tCrSFB_copy_view_filtered[None, None, k_block_next], + ) + + # end of for loop + # end of for loop + # Hoist out last k_tile + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_next = ( + 0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1 + ) + + if k_block_idx == num_k_blocks - 1: + cute.arch.fence_proxy("async.shared", space="cta") + mainloop_pipeline.consumer_release(mainloop_consumer_state) + mainloop_consumer_state.advance() + + if k_block_next > 0: + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + cute.copy( + smem_tiled_copy_B, + tCsB_p[None, None, k_block_next], + tCrB_copy_view[None, None, k_block_next], + ) + tCsSFA_p_filtered = cute.filter_zeros(tCsSFA_p) + tCsSFB_p_filtered = cute.filter_zeros(tCsSFB_p) + tCrSFA_copy_view_filtered = cute.filter_zeros(tCrSFA_copy_view) + tCrSFB_copy_view_filtered = cute.filter_zeros(tCrSFB_copy_view) + cute.copy( + smem_tiled_copy_SFA, + tCsSFA_p_filtered[None, None, k_block_next], + tCrSFA_copy_view_filtered[None, None, k_block_next], + ) + cute.copy( + smem_tiled_copy_SFB, + tCsSFB_p_filtered[None, None, k_block_next], + tCrSFB_copy_view_filtered[None, None, k_block_next], + ) + # Mixed FP4 x FP8 register-side bit shift before mma.sync (hoisted tail). + if cutlass.const_expr( + self.mixed_mode and self.a_dtype.width < 8 + ): + a_view_h = cute.recast_tensor( + tCrA[None, None, k_block_idx], cutlass.Int8 + ) + for _i in cutlass.range_constexpr(cute.size(a_view_h)): + a_view_h[_i] = cutlass.Int8( + a_view_h[_i] << FP4_SHIFT_BITS + ) + if cutlass.const_expr( + self.mixed_mode and self.b_dtype.width < 8 + ): + b_view_h = cute.recast_tensor( + tCrB[None, None, k_block_idx], cutlass.Int8 + ) + for _i in cutlass.range_constexpr(cute.size(b_view_h)): + b_view_h[_i] = cutlass.Int8( + b_view_h[_i] << FP4_SHIFT_BITS + ) + cute.gemm( + tiled_mma, + accumulators, + [ + tCrA[None, None, k_block_idx], + tCrSFA[None, None, k_block_idx], + ], + [ + tCrB[None, None, k_block_idx], + tCrSFB[None, None, k_block_idx], + ], + accumulators, + ) + + # ///////////////////////////////////////////////////////////////////////////// + # EPILOG + # ///////////////////////////////////////////////////////////////////////////// + + copy_atom_r2s = sm120_utils.sm120_get_smem_store_op( + self.c_layout, + elem_ty_d=self.c_dtype, + elem_ty_acc=self.acc_dtype, + ) + + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + self.c_layout.is_m_major_c(), + 2, + ), + self.c_dtype, + ) + + tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) + + tiled_copy_r2s = cute.make_tiled_copy_S( + copy_atom_r2s, + tiled_copy_C_Atom, + ) + + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + # (R2S, R2S_M, R2S_N, PIPE_D) + tRS_sD = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rAcc = tiled_copy_r2s.retile(accumulators) + + # Allocate D registers. + rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) + tRS_rD_layout = cute.make_layout(rD_shape[:3]) + tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype) + _size_tRS_rD = cute.size(tRS_rD) + + sepi_for_tma_partition = cute.group_modes(sC, 0, 2) + tcgc_for_tma_partition = cute.zipped_divide(gC_mnl_slice, self.epi_tile) + + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sepi_for_tma_partition, + tcgc_for_tma_partition, + ) + + # Initialize tma store pipeline + tma_store_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mma_warps * self.num_threads_per_warp, + ) + tma_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.epi_stage, + producer_group=tma_store_producer_group, + ) + + epi_rest_m = bSG_gD.shape[1][0] + epi_rest_n = bSG_gD.shape[1][1] + epi_tile_m = self.epi_tile[0] + epi_tile_n = self.epi_tile[1] + mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rAcc, mode=[1]) + mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rAcc, mode=[2]) + + for epi_m in cutlass.range_constexpr(epi_rest_m): + for epi_n in cutlass.range_constexpr(epi_rest_n): + MmaMPerEpiM = epi_tile_m // mma_tile_m + MmaNPerEpiN = epi_tile_n // mma_tile_n + for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN): + for mma_m_in_epi in cutlass.range_constexpr(MmaMPerEpiM): + mma_n = (epi_n * MmaNPerEpiN) + mma_n_in_epi + mma_m = (epi_m * MmaMPerEpiM) + mma_m_in_epi + tRS_rD_slice = tRS_rD[ + (None, mma_m_in_epi, mma_n_in_epi) + ] + tRS_rAcc_slice = tRS_rAcc[(None, mma_m, mma_n)] + for elem_idx in cutlass.range_constexpr( + cute.size(tRS_rD_slice) + ): + tRS_rD_slice[elem_idx] = tRS_rAcc_slice[elem_idx] + + # Type conversion + tRS_rD_out = cute.make_rmem_tensor( + tRS_rD_layout.shape, self.c_dtype + ) + acc_vec = tRS_rD.load() + tRS_rD_out.store(acc_vec.to(self.c_dtype)) + + # Register to shared memory + epi_buffer = epi_buffer + 1 + epi_buffer = epi_buffer % cute.size( + tRS_sD, mode=[3] + ) + self.epilog_sync_barrier.arrive_and_wait() + cute.copy( + tiled_copy_r2s, + tRS_rD_out, + tRS_sD[(None, None, None, epi_buffer)], + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # barrier for sync + self.epilog_sync_barrier.arrive_and_wait() + + # Get the global memory coordinate for the current epi tile. + gmem_coord = (epi_m, epi_n) + # Copy from shared memory to global memory + if warp_idx == 0: + cute.copy( + tma_atom_c, + bSG_sD[(None, epi_buffer)], + bSG_gD[(None, gmem_coord)], + ) + tma_store_pipeline.producer_commit() + tma_store_pipeline.producer_acquire() + + # Advance to the next work tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # End of for k_tile loop + # End of while loop + # End of MMA warp group + # Start of DMA warp group + elif warp_idx == self.tma_load_warp_id: + cute.arch.setmaxregister_decrease(self.load_register_requirement) + + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] + tBgB_nkl = tBgB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])] + tAgSFA_mkl = tAgSFA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] + tBgSFB_nkl = tBgSFB[(None, tile_coord_mnl[1], None, tile_coord_mnl[2])] + + mainloop_producer_state.reset_count() + + for k_tile in range(0, k_tile_cnt, 1, unroll=1): + # ///////////////////////////////////////////////////////////////////////////// + # Wait for A/B buffers to be empty before loading into them + # Also sets the transaction barrier for the A/B buffers + # ///////////////////////////////////////////////////////////////////////////// + mainloop_pipeline.producer_acquire(mainloop_producer_state) + + # ///////////////////////////////////////////////////////////////////////////// + # Slice to global/shared memref to current k_tile + # ///////////////////////////////////////////////////////////////////////////// + tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)] + tAsA_pipe = tAsA[(None, mainloop_producer_state.index)] + + tBgB_k = tBgB_nkl[(None, mainloop_producer_state.count)] + tBsB_pipe = tBsB[(None, mainloop_producer_state.index)] + + tAgSFA_k = tAgSFA_mkl[(None, mainloop_producer_state.count)] + tAsSFA_pipe = tAsSFA[(None, mainloop_producer_state.index)] + + tBgSFB_k = tBgSFB_nkl[(None, mainloop_producer_state.count)] + tBsSFB_pipe = tBsSFB[(None, mainloop_producer_state.index)] + + # ///////////////////////////////////////////////////////////////////////////// + # TMA load A/B + # ///////////////////////////////////////////////////////////////////////////// + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + cute.copy( + tma_atom_sfa, + tAgSFA_k, + tAsSFA_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + cute.copy( + tma_atom_sfb, + tBgSFB_k, + tBsSFB_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + # Mainloop pipeline's producer commit is a NOP + mainloop_pipeline.producer_commit(mainloop_producer_state) + mainloop_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # end of while loop + + # Wait A/B buffer empty + mainloop_pipeline.producer_tail(mainloop_producer_state) + return + + @staticmethod + def _compute_stages( + tile_shape_mnk: tuple[int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + sf_dtype: type[cutlass.Numeric], + sfa_smem_layout: cute.Layout, + sfb_smem_layout: cute.Layout, + epi_tile: tuple[int, int], + c_dtype: type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + ) -> tuple[int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type tile_shape_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (A/B operand stages, epilogue stages) + :rtype: tuple[int, int] + """ + + epi_stage_max = (tile_shape_mnk[1] // epi_tile[1]) * ( + tile_shape_mnk[0] // epi_tile[0] + ) + epi_stage = min(epi_stage_max, 4) + c_bytes_per_stage = cute.size(epi_tile) * c_dtype.width // 8 + epi_bytes = c_bytes_per_stage * epi_stage + + # A/B shared memory size + a_shape = cute.slice_(tile_shape_mnk, (None, 0, None)) + b_shape = cute.slice_(tile_shape_mnk, (0, None, None)) + ab_bytes_per_stage = ( + cute.size(a_shape) * a_dtype.width // 8 + + cute.size(b_shape) * b_dtype.width // 8 + ) + # SFA/SFB shared memory size + sf_bytes_per_stage = ( + cute.size(cute.filter_zeros(sfa_smem_layout).shape) * sf_dtype.width // 8 + + cute.size(cute.filter_zeros(sfb_smem_layout).shape) * sf_dtype.width // 8 + ) + # Mbar helpers size + mbar_helpers_bytes = 1024 + + ab_stage = ( + (smem_capacity - occupancy * 1024) // occupancy + - mbar_helpers_bytes + - epi_bytes + ) // (ab_bytes_per_stage + sf_bytes_per_stage) + return ab_stage, epi_stage + + @staticmethod + def _make_smem_layouts( + tile_shape_mnk: tuple[int, int, int], + epi_tile: tuple[int, int], + a_dtype: type[cutlass.Numeric], + a_layout: cute.Layout, + b_dtype: type[cutlass.Numeric], + b_layout: cute.Layout, + ab_stage: int, + c_dtype: type[cutlass.Numeric], + c_layout: cute.Layout, + epi_stage: int, + sf_vec_size: int, + tiled_mma: cute.TiledMma, + ) -> tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]: + """Create shared memory layouts for A, B, and C tensors. + + :param tile_shape_mnk: CTA tile shape (M,N,K) + :type tile_shape_mnk: Tuple[int, int, int] + :param epi_tile: Epilogue tile shape + :type epi_tile: Tuple[int, int] + :param a_dtype: Data type for matrix A + :type a_dtype: type[cutlass.Numeric] + :param a_layout: Layout for matrix A + :type a_layout: Layout + :param b_dtype: Data type for matrix B + :type b_dtype: type[cutlass.Numeric] + :param b_layout: Layout for matrix B + :type b_layout: Layout + :param ab_stage: Number of stages for A/B tensors + :type ab_stage: int + :param c_dtype: Data type for output matrix C + :type c_dtype: type[cutlass.Numeric] + :param c_layout: leading dimension of the output matrix C + :type c_layout: Layout + :param epi_stage: Number of epilogue stages + :type epi_stage: int + + :return: Tuple of shared memory layouts for A, B, and C + :rtype: Tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout] + """ + a_smem_shape = cute.slice_(tile_shape_mnk, (None, 0, None)) + + a_is_k_major = a_layout.is_k_major_a() + b_is_k_major = b_layout.is_k_major_b() + a_major_mode_size = tile_shape_mnk[2 if a_is_k_major else 0] + + a_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + a_layout, + a_dtype, + a_major_mode_size, + ), + a_dtype, + ) + a_smem_layout_staged = cute.tile_to_shape( + a_smem_layout_atom, + cute.append(a_smem_shape, ab_stage), + order=(0, 1, 2) if a_is_k_major else (1, 0, 2), + ) + + b_smem_shape = cute.slice_(tile_shape_mnk, (0, None, None)) + + b_major_mode_size = tile_shape_mnk[2 if b_is_k_major else 1] + b_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + b_layout, + b_dtype, + b_major_mode_size, + ), + b_dtype, + ) + b_smem_layout_staged = cute.tile_to_shape( + b_smem_layout_atom, + cute.append(b_smem_shape, ab_stage), + order=(0, 1, 2) if b_is_k_major else (1, 0, 2), + ) + + sfa_smem_layout_staged = blockscaled_utils.sm120_make_smem_layout_sfa( + tiled_mma, + tile_shape_mnk, + sf_vec_size, + ab_stage, + ) + + sfb_smem_layout_staged = blockscaled_utils.sm120_make_smem_layout_sfb( + tiled_mma, + tile_shape_mnk, + sf_vec_size, + ab_stage, + ) + + c_smem_shape = epi_tile + c_major_mode_size = epi_tile[1] if c_layout.is_n_major_c() else epi_tile[0] + c_smem_layout_atom = cute.nvgpu.warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom( + c_layout, + c_dtype, + c_major_mode_size, + ), + c_dtype, + ) + epi_smem_layout_staged = cute.tile_to_shape( + c_smem_layout_atom, + cute.append(c_smem_shape, epi_stage), + order=(1, 0, 2) if c_layout.is_m_major_c() else (0, 1, 2), + ) + + return ( + a_smem_layout_staged, + b_smem_layout_staged, + sfa_smem_layout_staged, + sfb_smem_layout_staged, + epi_smem_layout_staged, + ) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + tile_shape_mnk: tuple[int, int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[int, int, int]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type tile_shape_mnk: tuple[int, int, int] + + :return: Grid shape for kernel launch. + :rtype: tuple[int, int, int] + """ + + c_shape = cute.slice_(tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (1, 1, 1) + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + return tile_sched_params, grid + + @staticmethod + def _make_tma_store_atoms_and_tensors( + tensor_c: cute.Tensor, + epi_smem_layout_staged: cute.ComposedLayout, + epi_tile: tuple[int, int], + ) -> tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for C tensor storage. + + :param tensor_c: Output tensor C + :type tensor_c: cute.Tensor + :param epi_smem_layout_staged: Shared memory layout for epilogue + :type epi_smem_layout_staged: cute.ComposedLayout + :param epi_tile: Epilogue tile shape + :type epi_tile: Tuple[int, int] + + :return: TMA atom and tensor for C + :rtype: Tuple[cute.CopyAtom, cute.Tensor] + """ + epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + tensor_c, + epi_smem_layout, + epi_tile, + ) + + return tma_atom_c, tma_tensor_c + + @staticmethod + def _make_tma_atoms_and_tensors( + tensor: cute.Tensor, + smem_layout_staged: cute.ComposedLayout, + smem_tile: tuple[int, int], + mcast_dim: int, + internal_type: Optional[Type[cutlass.Numeric]] = None, + ) -> tuple[cute.CopyAtom, cute.Tensor]: + """Create TMA atoms and tensors for input tensors. + + :param tensor: Input tensor (A or B) + :type tensor: cute.Tensor + :param smem_layout_staged: Shared memory layout for the tensor + :type smem_layout_staged: cute.ComposedLayout + :param smem_tile: Shared memory tile shape + :type smem_tile: Tuple[int, int] + :param mcast_dim: Multicast dimension + :type mcast_dim: int + + :return: TMA atom and tensor + :rtype: Tuple[cute.CopyAtom, cute.Tensor] + """ + op = ( + cpasync.CopyBulkTensorTileG2SOp() + if mcast_dim == 1 + else cpasync.CopyBulkTensorTileG2SMulticastOp() + ) + + smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + op, + tensor, + smem_layout, + smem_tile, + num_multicast=mcast_dim, + internal_type=internal_type, + ) + + return tma_atom, tma_tensor + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +def run( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str = "k", + b_major: str = "k", + c_major: str = "n", + tile_shape_mnk: Tuple[int, int, int] = (128, 128, 128), + epi_tile: Tuple[int, int] = (128, 128), + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Perf-framework-compatible entry point. + + FP4 MMA only supports Float32 accumulation, so acc_dtype is always + Float32 regardless of what the caller passes. + """ + return run_bs( + mnkl=mnkl, + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + acc_dtype=cutlass.Float32, + a_major=a_major, + b_major=b_major, + c_major=c_major, + tile_shape_mnk=tile_shape_mnk, + epi_tile=epi_tile, + tolerance=tolerance, + warmup_iterations=warmup_iterations, + iterations=iterations, + skip_ref_check=skip_ref_check, + use_cold_l2=use_cold_l2, + ) + + +def run_bs( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + tile_shape_mnk: Tuple[int, int, int], + epi_tile: Tuple[int, int], + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + import torch + import cutlass.torch as cutlass_torch + + print("Running Blackwell Geforce Blockscaled Dense GEMM with:") + print(f"mnkl: {mnkl}") + print( + f"A dtype: {a_dtype}, B dtype: {b_dtype}, A/B scale factor dtype: {sf_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}" + ) + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Tile Shape: {tile_shape_mnk}") + print(f"Epilogue tile: {epi_tile}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {use_cold_l2}") + + m, n, k, l = mnkl + + if not Sm120BlockScaledGemmKernel.is_valid_tensor_alignment( + m, n, k, l, a_dtype, c_dtype, a_major, b_major, c_major + ): + raise ValueError("Invalid tensor alignment") + + a_dtype = getattr(cutlass, a_dtype) if isinstance(a_dtype, str) else a_dtype + b_dtype = getattr(cutlass, b_dtype) if isinstance(b_dtype, str) else b_dtype + c_dtype = getattr(cutlass, c_dtype) if isinstance(c_dtype, str) else c_dtype + acc_dtype = getattr(cutlass, acc_dtype) if isinstance(acc_dtype, str) else acc_dtype + + # Unpack parameters + m, n, k, l = mnkl + cluster_shape_mnk = (1, 1, 1) + + # Prepare pytorch tensors: A, B (random from 0 to 2) and C (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor to be byte aligned + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=2 if a_dtype == cutlass.Float4E2M1FN else 1, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=2 if b_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + # Create scale factor tensor SFA/SFB + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + # The A, C, and D matrices are row-major whereas the B matrix is column-major. + # So only k-major (A/B) is supported. + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=0, + max_val=1, + ), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + + return ref_f32_torch_tensor_cpu, cute_tensor, torch_tensor + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype + ) + + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype + ) + + # configure gemm kernel + gemm = Sm120BlockScaledGemmKernel( + acc_dtype, + sf_vec_size, + tile_shape_mnk, + epi_tile, + ) + + # Compute max active clusters on current device + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1] + ) + + # Initialize stream + stream = cutlass_torch.default_stream() + + # compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + max_active_clusters, + stream, + ) + + if not skip_ref_check: + print("Reference checking ...") + # execution + compiled_gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, stream) + torch.cuda.synchronize() + + # Ref check + res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Convert c back to f32 for comparison. + c_ref_device = c_ref.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype is cutlass.Float4E2M1FN: + # Convert ref : f32 -> f4 -> f32 + ref_f4_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f4 = from_dlpack(ref_f4_, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + ref_f4.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f4) + cute.testing.convert(ref_f4, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + + def generate_tensors(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, a_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, b_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + # Mark tensor to be byte aligned + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=2 if a_dtype == cutlass.Float4E2M1FN else 1, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=2 if b_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) + return cute.testing.JitArguments( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, stream + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_ref.numel() * a_ref.element_size() + + b_ref.numel() * b_ref.element_size() + + sfa_ref.numel() * sfa_ref.element_size() + + sfb_ref.numel() * sfb_ref.element_size() + + c_ref.numel() * c_ref.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + gflop = 2 * m * n * k / 1e9 + gflops = gflop / exec_time * 1e6 + + print(f"Execution time: {exec_time} microseconds per iteration") + print(f"GFLOPS: {gflops}") + + return exec_time # Return execution time in microseconds + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Example of MxNxKxL GEMM on Blackwell Geforce." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=( + 1024, + 1024, + 1024, + 1, + ), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--tile_shape_mnk", + type=parse_comma_separated_ints, + choices=[ + (128, 128, 128), + (128, 128, 256), + ], + default=(128, 128, 128), + help="CTA tile shape (comma-separated)", + ) + parser.add_argument( + "--epi_tile", + type=parse_comma_separated_ints, + choices=[ + (128, 128), + (64, 32), + ], + default=(128, 128), + help="Epilogue tile shape (comma-separated)", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Float4E2M1FN, + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.Float4E2M1FN, + ) + parser.add_argument( + "--sf_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + ) + parser.add_argument( + "--sf_vec_size", + type=int, + choices=[16, 32], # 16 for NVFP4, 32 for MXFP4 / MXFP8. + default=16, + ) + parser.add_argument( + "--c_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + ) + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + ) + parser.add_argument("--a_major", choices=["k"], type=str, default="k") + parser.add_argument("--b_major", choices=["k"], type=str, default="k") + parser.add_argument("--c_major", choices=["n"], type=str, default="n") + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", + action="store_true", + default=False, + help="Skip reference checking", + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + fp4_allowed_tiles = {(128, 128, 128), (128, 128, 256)} + fp8_allowed_tiles = {(128, 128, 128)} + try: + validate_blockscaled_args(args, fp4_allowed_tiles, fp8_allowed_tiles) + except ValueError as e: + parser.error(str(e)) + + return args + + +if __name__ == "__main__": + args = parse_arguments() + run_bs( + args.mnkl, + args.a_dtype, + args.b_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.tile_shape_mnk, + args.epi_tile, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + print("PASS") diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/requirements.txt b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/requirements.txt new file mode 100644 index 0000000000..4c1fac17af --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/reference/requirements.txt @@ -0,0 +1,6 @@ +# Versions used for the recorded CuTeDSL reference run. +nvidia-cutlass-dsl==4.6.2 +cuda-python==13.3.1 +cuda-bindings==13.3.1 +torch==2.13.0 +numpy==2.5.2 diff --git a/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/src/main.rs new file mode 100644 index 0000000000..b1de97f842 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/src/main.rs @@ -0,0 +1,702 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#![feature(f16)] + +//! Block-scaled MXFP4 GEMM through cute layouts. +//! +//! C(MxN, f16) = A(MxK, e2m1) x B(KxN, e2m1), with one ue8m0 scale per +//! 32-element K-block per row of A and per column of B. Four FP4 values share +//! one 16-bit carrier, and one typed 16x8x64 block-scaled MMA consumes two +//! scale groups per operation. +//! +//! A typed three-stage TMA pipeline overlaps lookahead loads with the existing +//! MMA mapping and tiled FP16 TMA epilogue: +//! +//! ```text +//! host: packed A/B + canonical SFA/SFB +//! | +//! v +//! four typed TMA-load descriptors + one TMA-store descriptor +//! | +//! v +//! resident CTAs: blockIdx.x, then += gridDim.x +//! | +//! v +//! for each scheduled 128x128 output tile { +//! warp 8, lane 0: fill the three-slot ring until EMPTY blocks +//! each stage is 17,408 bytes from four TMA loads +//! warps 0..7: wait FULL, then consume two packed K64 steps +//! accumulators.gemm(...) +//! typed epilogue: FragmentC -> f16x2 -> stmatrix -> shared +//! warp 0, lane 0: two 128x64 TMA stores -> global C +//! } +//! producer/consumer slot + phase survive between output tiles +//! ``` +//! +//! `TensorMap` makes the host encoder and device copy name the same +//! placement. The byte TMA view and f16 consumer overlay encode the same B64 +//! swizzle; the scale view keeps the hardware's canonical atom unchanged. +//! +//! B stages N-major (columns of B, k contiguous), so both operands share +//! one tile type and both copies land through the same 64-byte swizzle. +//! +//! Why B needs no transposed source view: storing B as N x K gives both +//! operands K-contiguous packed rows. Each f16 carrier holds four consecutive +//! nibbles, low to high, and the backend selects the matching fragment loads +//! and same-format block-scaled MMA. +//! +//! Data and scales are chosen so every f32 intermediate is exact and the +//! comparison is bitwise. The accelerated operation only exists on +//! sm_120/121; other GPUs report the requirement and skip execution. After +//! correctness, CUDA events measure 31 single launches following 100 warmups. +//! +//! Build and run with: +//! cargo oxide run blockscale_gemm_cute --arch sm_120a + +use core::{marker::PhantomData, mem::size_of}; +use ctm_rs::mbarrier::Barrier; +use ctm_rs::memory::DynamicSharedArray; +use ctm_rs::{cuda_module, kernel, launch_bounds, launch_contract, thread, warp}; +use cuda_core::{CudaContext, LaunchConfig1D}; +use cute_rs::copy::tma::{TmaEncodeOptions, create_tensor_map_tiled_with_options}; +use cute_rs::{ + C, ComposedLayout, Consumer, Layout, Mkl, Mxf4AccumulatorTile2x8, Mxf4E2M1, Mxf4TiledMma, Nkl, + PipelineState, PipelineTmaAsync, PipelineTmaStore, Producer, RowMajor, SharedScaleAtom, + SharedTensor, Sm1xxBlockScaleKMajor, Sm120Epilogue128x128, Sm120EpilogueHalfLayout, + Sm120ScaleAtom, SmemTile, StaticPersistentTileScheduler, Swizzle, TensorMap, copy_tma_2d, + copy_tma_s2g_2d, +}; + +/// Output row count for the fixed benchmark problem. +const M: usize = 2048; +/// Output column count for the fixed benchmark problem. +const N: usize = 2048; +/// Reduction extent in logical FP4 elements before packing. +const K: usize = 1024; // logical e2m1 elements; four packed into each f16 carrier +/// Number of 32-element K groups, each with one UE8M0 scale. +const KB: usize = K / 32; // K-blocks: one ue8m0 scale per 32 elements +/// Eight compute warps arranged as four along M and two along N. +const COMPUTE_WARPS: usize = 8; // 4 along M, 2 along N +/// Warp index reserved for issuing TMA loads into the shared rings. +const PRODUCER_WARP: usize = COMPUTE_WARPS; +/// Threads required in each CTA by the kernel launch contract. +const THREADS: usize = (COMPUTE_WARPS + 1) * 32; +/// Output rows computed by each CTA tile. +const BLOCK_M: usize = 128; +/// Output columns computed by each CTA tile. +const BLOCK_N: usize = 128; +/// Number of independent matrices in the static tile scheduler. +const BATCHES: usize = 1; +/// Persistent assignment of the 16x16 output-tile grid. +type TileScheduler = StaticPersistentTileScheduler<16, 16, BATCHES>; + +// TMA moves one 128-row x 64-byte packed-FP4 stage per operand. The consumer +// overlays those exact bytes as 128x32 f16 carriers for ldmatrix. +/// Unswizzled 128x32 FP16-carrier layout overlaying packed FP4 input bytes. +type TileP = Layout<(C<128>, C<32>), (C<32>, C<1>)>; +/// Element-domain swizzle for the FP16 carriers consumed by matrix loads. +type SmemP = ComposedLayout, 0, TileP>; +/// The same packed input stage expressed as 128x64 bytes for TMA. +type ByteSmem = ComposedLayout, 0, RowMajor<128, 64>>; +/// Contiguous 512-byte scale stage carried as 256 u16 values. +type ScaleTmaSmem = RowMajor<1, 256>; +/// Typed shared A view carrying packed MXF4 values in FP16 transport lanes. +type SharedA = SharedTensor; +/// Typed shared B view carrying packed MXF4 values in FP16 transport lanes. +type SharedB = SharedTensor; + +/// Number of 128-element K tiles consumed for each output tile. +const K_STAGES: usize = KB / 4; +/// Byte footprint of one packed input-operand stage. +const DATA_STAGE_BYTES: usize = BLOCK_M * 64; +/// Byte footprint of one SM120 scale atom. +const SCALE_STAGE_BYTES: usize = Sm120ScaleAtom::BYTES; +/// Total A/B value and scale bytes announced to each full-stage barrier. +const TMA_TX_BYTES: u32 = (2 * DATA_STAGE_BYTES + 2 * SCALE_STAGE_BYTES) as u32; +/// Number of shared slots overlapped by the load pipeline. +const PIPELINE_STAGES: usize = 3; +/// Three-stage load pipeline with eight consumer warps and an exact transaction-byte count. +type MainloopPipeline = PipelineTmaAsync; +/// TMA store queue allowing one outstanding committed output group. +type EpiloguePipeline = PipelineTmaStore<1>; + +// One 1024-byte-aligned dynamic allocation, partitioned without padding +// surprises. Each operand owns a three-slot ring, followed by the six 8-byte +// full/empty barriers, alignment padding, and the 32 KiB FP16 epilogue tile. +/// Byte offset of the three-slot A-value ring in dynamic shared memory. +const SMEM_A_OFFSET: usize = 0; +/// Byte offset of the three-slot B-value ring after the A ring. +const SMEM_B_OFFSET: usize = SMEM_A_OFFSET + PIPELINE_STAGES * DATA_STAGE_BYTES; +/// Byte offset of the three-slot A-scale ring after the value rings. +const SMEM_SFA_OFFSET: usize = SMEM_B_OFFSET + PIPELINE_STAGES * DATA_STAGE_BYTES; +/// Byte offset of the three-slot B-scale ring after the A scales. +const SMEM_SFB_OFFSET: usize = SMEM_SFA_OFFSET + PIPELINE_STAGES * SCALE_STAGE_BYTES; +/// Byte offset of the full/empty load barriers after all staged input storage. +const PIPELINE_OFFSET: usize = SMEM_SFB_OFFSET + PIPELINE_STAGES * SCALE_STAGE_BYTES; +/// Byte extent occupied by input rings and their load-pipeline barriers. +const MAINLOOP_SMEM_BYTES: usize = PIPELINE_OFFSET + MainloopPipeline::STORAGE_BYTES; +/// Aligned byte offset of the shared FP16 output tile. +const EPILOGUE_OFFSET: usize = + MAINLOOP_SMEM_BYTES.next_multiple_of(Sm120Epilogue128x128::ALIGNMENT); +/// Total dynamic shared-memory allocation in bytes required per CTA. +const DYNAMIC_SMEM_BYTES: usize = EPILOGUE_OFFSET + Sm120Epilogue128x128::BYTES; +/// Compile-time checks tying tile counts, byte offsets, alignments, and launch storage to the schedule. +const _: () = { + assert!(M / BLOCK_M == 16); + assert!(N / BLOCK_N == 16); + assert!(TileScheduler::TOTAL_TILES == 256); + assert!(TMA_TX_BYTES == 17_408); + assert!(SMEM_A_OFFSET == 0); + assert!(SMEM_B_OFFSET == 24_576); + assert!(SMEM_SFA_OFFSET == 49_152); + assert!(SMEM_SFB_OFFSET == 50_688); + assert!(PIPELINE_OFFSET == 52_224); + assert!(SMEM_A_OFFSET.is_multiple_of(1024)); + assert!(SMEM_B_OFFSET.is_multiple_of(1024)); + assert!(SMEM_SFA_OFFSET.is_multiple_of(512)); + assert!(SMEM_SFB_OFFSET.is_multiple_of(512)); + assert!(PIPELINE_OFFSET.is_multiple_of(MainloopPipeline::STORAGE_ALIGN)); + assert!(MAINLOOP_SMEM_BYTES == 52_272); + assert!(EPILOGUE_OFFSET == 53_248); + assert!(EPILOGUE_OFFSET.is_multiple_of(Sm120Epilogue128x128::ALIGNMENT)); + assert!(DYNAMIC_SMEM_BYTES == 86_016); +}; + +/// SM120 block-scaled GEMM entry point with typed CuTe views and pipelines. +#[cuda_module] +mod kernels { + + use super::*; + + /// Four typed input tensor maps fill each 17,408-byte mainloop stage. The + /// output tensor map carries the complete FP16 epilogue tile to global C. + #[launch_bounds(288)] + #[launch_contract( + domain = 1, + block = (288, 1, 1), + dynamic_shared = 86016, + dynamic_shared_alignment = 1024, + min_compute_capability = (12, 0), + )] + #[kernel] + pub fn blockscale_gemm( + #[grid_constant] a_tma: &TensorMap, + #[grid_constant] b_tma: &TensorMap, + #[grid_constant] sfa_tma: &TensorMap, + #[grid_constant] sfb_tma: &TensorMap, + #[grid_constant] out_tma: &TensorMap, + ) { + let tid = thread::threadIdx_x(); + + // All aliases are fixed offsets into one externally sized allocation. + // The launch contract and host LaunchConfig both reserve exactly the + // same byte count and request the same 1024-byte base alignment. + let dynamic_smem = DynamicSharedArray::::get_raw(); + let smem_a_ptr = unsafe { dynamic_smem.add(SMEM_A_OFFSET) }; + let smem_b_ptr = unsafe { dynamic_smem.add(SMEM_B_OFFSET) }; + let smem_sfa_ptr = unsafe { dynamic_smem.add(SMEM_SFA_OFFSET) }; + let smem_sfb_ptr = unsafe { dynamic_smem.add(SMEM_SFB_OFFSET) }; + let pipeline_base = unsafe { dynamic_smem.add(PIPELINE_OFFSET).cast::() }; + let pipeline = unsafe { MainloopPipeline::from_raw_base(pipeline_base) }; + let epilogue = unsafe { + Sm120Epilogue128x128::from_raw(dynamic_smem.add(EPILOGUE_OFFSET).cast::()) + }; + unsafe { pipeline.init(0) }; + + let warp_id = warp::warp_id() as usize; + let lane = tid % 32; + if warp_id < COMPUTE_WARPS { + // CuTe's (4,2) MMA-atom layout. The 128-M/N permutations give + // every compute warp two M atoms and four separated N16 pairs. + let wm = warp_id % 4; + let wn = warp_id / 4; + let tiled_mma = Mxf4TiledMma::::get_slice(lane); + let warp_epilogue = unsafe { epilogue.get_slice(warp_id, lane as u32) }; + let epilogue_pipeline = EpiloguePipeline::new(); + let mut scheduler = TileScheduler::new_1d(); + // Slot and phase continue across output tiles. With eight K stages + // and a three-slot ring, resetting this state per tile would + // restart at the wrong barrier and can deadlock. + let mut consumer_state = PipelineState::::new(); + while scheduler.has_work() { + let work_tile = scheduler.get_current_work(); + let (m_tile, n_tile, batch) = work_tile.tile_idx(); + let mut accumulators = Mxf4AccumulatorTile2x8::zero(); + + // Only the logical K coordinate resets per output tile. + let mut stage = 0usize; + while stage < K_STAGES { + unsafe { pipeline.consumer_wait(&consumer_state) }; + + let consumer_slot = consumer_state.index(); + let a_stage_ptr = unsafe { smem_a_ptr.add(consumer_slot * DATA_STAGE_BYTES) }; + let b_stage_ptr = unsafe { smem_b_ptr.add(consumer_slot * DATA_STAGE_BYTES) }; + let sfa_stage_ptr = + unsafe { smem_sfa_ptr.add(consumer_slot * SCALE_STAGE_BYTES) }; + let sfb_stage_ptr = + unsafe { smem_sfb_ptr.add(consumer_slot * SCALE_STAGE_BYTES) }; + + // SAFETY: whole warps in uniform control flow; windows lie + // inside the completed TMA stage. The byte and f16 layouts + // encode the same B64 physical swizzle; scales preserve the + // SM120 atom. + unsafe { + let s_a = SharedA::from_raw_parts(a_stage_ptr.cast::(), BLOCK_M * 32); + let s_b = SharedB::from_raw_parts(b_stage_ptr.cast::(), BLOCK_N * 32); + let s_sfa = SharedScaleAtom::::from_raw_parts( + sfa_stage_ptr.cast::(), + Sm120ScaleAtom::WORDS, + ); + let s_sfb = SharedScaleAtom::::from_raw_parts( + sfb_stage_ptr.cast::(), + Sm120ScaleAtom::WORDS, + ); + let scale_tile = tiled_mma.load_scale_atom_128(&s_sfa, &s_sfb, wm, wn); + // Two packed K=64 steps per stage: window column kk + // selects one. + let mut kk = 0usize; + while kk < 2 { + // SAFETY: the fixed loop visits exactly the two K64 + // halves of this K128 stage. + let scales = scale_tile.pairs_at_unchecked(kk); + let a_fragments = tiled_mma.load_a_128(&s_a, wm, kk); + let b_fragments = tiled_mma.get_b_tile_k64(&s_b, wn, kk); + accumulators.gemm(&tiled_mma, a_fragments, b_fragments, scales); + kk += 1; + } + } + // Every compute warp calls uniformly; the pipeline elects + // one lane-zero arrival, giving EMPTY eight arrivals. + unsafe { pipeline.consumer_release(&consumer_state) }; + consumer_state.advance(); + stage += 1; + } + + // A single shared epilogue buffer is reused by every logical + // tile assigned to this persistent CTA. Lane zero waits for + // the previous TMA store's shared reads; the counted barrier + // releases only the eight compute warps after that wait. + if tid == 0 { + epilogue_pipeline.producer_acquire(); + } + unsafe { + epilogue.sync_reusable(); + + warp_epilogue.store_tile(accumulators); + + // This semantic boundary publishes every writer's + // generic-proxy stores before the counted hand-off to + // the one TMA issuer. + epilogue.sync_ready_for_tma(); + } + + if tid == 0 { + let c_outer_m = batch * (M / BLOCK_M) + m_tile; + let c_outer_n = n_tile * 2; + unsafe { + let left = epilogue.tma_half::<0>(); + copy_tma_s2g_2d(out_tma, (c_outer_m, c_outer_n), left); + let right = epilogue.tma_half::<1>(); + copy_tma_s2g_2d(out_tma, (c_outer_m, c_outer_n + 1), right); + } + epilogue_pipeline.producer_commit(); + } + scheduler.advance_to_next_work(); + } + if tid == 0 { + epilogue_pipeline.producer_tail(); + } + } else { + // Lane zero is the single TMA issuer; the other producer-warp + // lanes remain inactive after the uniform pipeline setup. + if warp_id == PRODUCER_WARP && lane == 0 { + let mut scheduler = TileScheduler::new_1d(); + // Keep the circular cursor live across all scheduled output + // tiles; only the descriptor's logical K coordinate resets. + let mut producer_state = PipelineState::::new(); + while scheduler.has_work() { + let work_tile = scheduler.get_current_work(); + let (m_tile, n_tile, batch) = work_tile.tile_idx(); + let a_outer = batch * (M / BLOCK_M) + m_tile; + let b_outer = batch * (N / BLOCK_N) + n_tile; + let mut stage = 0usize; + while stage < K_STAGES { + let producer_slot = producer_state.index(); + let a_stage_ptr = + unsafe { smem_a_ptr.add(producer_slot * DATA_STAGE_BYTES) }; + let b_stage_ptr = + unsafe { smem_b_ptr.add(producer_slot * DATA_STAGE_BYTES) }; + let sfa_stage_ptr = + unsafe { smem_sfa_ptr.add(producer_slot * SCALE_STAGE_BYTES) }; + let sfb_stage_ptr = + unsafe { smem_sfb_ptr.add(producer_slot * SCALE_STAGE_BYTES) }; + + unsafe { + pipeline.producer_acquire(&producer_state); + let full_barrier = pipeline.producer_expect_tx(&producer_state); + let mut a_tile = SmemTile:: { + base: a_stage_ptr, + capacity: DATA_STAGE_BYTES, + layout: PhantomData, + }; + let mut b_tile = SmemTile:: { + base: b_stage_ptr, + capacity: DATA_STAGE_BYTES, + layout: PhantomData, + }; + let mut sfa_tile = SmemTile:: { + base: sfa_stage_ptr.cast::(), + capacity: SCALE_STAGE_BYTES / size_of::(), + layout: PhantomData, + }; + let mut sfb_tile = SmemTile:: { + base: sfb_stage_ptr.cast::(), + capacity: SCALE_STAGE_BYTES / size_of::(), + layout: PhantomData, + }; + copy_tma_2d(a_tma, (a_outer, stage), &mut a_tile, full_barrier); + copy_tma_2d(b_tma, (b_outer, stage), &mut b_tile, full_barrier); + copy_tma_2d( + sfa_tma, + (a_outer * K_STAGES + stage, 0), + &mut sfa_tile, + full_barrier, + ); + copy_tma_2d( + sfb_tma, + (b_outer * K_STAGES + stage, 0), + &mut sfb_tile, + full_barrier, + ); + } + producer_state.advance(); + stage += 1; + } + scheduler.advance_to_next_work(); + } + // Drain only once, after the final scheduled tile. Per-tile + // tails would destroy overlap between consecutive work tiles. + unsafe { pipeline.producer_tail(producer_state) }; + } + } + } +} + +/// The eight non-negative e2m1 values, indexed by nibble 0x0..=0x7; +/// nibbles 0x8..=0xF are their negations (0x8 is -0.0). +const E2M1_VALUES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + +/// Decode the sign and three-bit magnitude of one FP4 nibble for the CPU oracle. +fn e2m1_decode(nibble: u8) -> f32 { + let magnitude = E2M1_VALUES[(nibble & 0x7) as usize]; + if nibble & 0x8 != 0 { + -magnitude + } else { + magnitude + } +} + +/// Small deterministic hash used only to build the host fixture. +/// +/// Mixing both coordinates prevents tile-sized row/column permutations from +/// hiding behind a short periodic input pattern. +fn mix32(mut value: u32) -> u32 { + value ^= value >> 16; + value = value.wrapping_mul(0x7feb_352d); + value ^= value >> 15; + value = value.wrapping_mul(0x846c_a68b); + value ^ (value >> 16) +} + +/// Deterministic data covering all 16 nibbles (zeros, negatives, both -0). +fn a_nibble(m: usize, k: usize) -> u8 { + (mix32((m as u32).wrapping_mul(0x9e37_79b9) ^ k as u32 ^ 0xa511_e9b3) & 0xf) as u8 +} + +/// Generate a deterministic B value from its K/N coordinates, including signed zeros. +fn b_nibble(k: usize, n: usize) -> u8 { + (mix32((n as u32).wrapping_mul(0x85eb_ca6b) ^ k as u32 ^ 0x63d8_3595) & 0xf) as u8 +} + +/// ue8m0 scale bytes: 2^(byte-127). 126..=128 keeps every product a small +/// binary fraction that `f32` represents exactly. +fn sfa_byte(m: usize, kb: usize) -> u8 { + 126 + (mix32((m as u32).wrapping_mul(0xc2b2_ae35) ^ kb as u32) % 3) as u8 +} + +/// Generate a deterministic B scale exponent in 126..=128 from its row and K-group indices. +fn sfb_byte(n: usize, kb: usize) -> u8 { + 126 + (mix32((n as u32).wrapping_mul(0x27d4_eb2f) ^ kb as u32 ^ 0x1656_67b1) % 3) as u8 +} + +/// Decode a finite UE8M0 exponent byte as 2 raised to byte minus 127. +fn ue8m0_decode(byte: u8) -> f32 { + 2f32.powi(byte as i32 - 127) +} + +/// Pack four increasing-K FP4 nibbles into each 16-bit transport carrier. +fn pack_mxf4_view(nibbles: &[u8]) -> Vec { + nibbles + .as_chunks::<4>() + .0 + .iter() + .map(|q| { + f16::from_bits( + (q[0] as u16 & 0xf) + | ((q[1] as u16 & 0xf) << 4) + | ((q[2] as u16 & 0xf) << 8) + | ((q[3] as u16 & 0xf) << 12), + ) + }) + .collect() +} + +/// Reorder dense `[row][K/32]` scales into SM120's 128-row x four-group +/// atoms, then expose the exact bytes as the u16 carrier used by TMA. +fn pack_canonical_scales(dense: &[u8], rows: usize) -> Vec { + let layout = Sm1xxBlockScaleKMajor::<32>::new(rows, K); + let mut packed = vec![0u8; layout.storage_len(1)]; + for row in 0..rows { + for group in 0..KB { + packed[layout.crd2idx(0, row, group)] = dense[row * KB + group]; + } + } + packed + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u16::from_le_bytes([pair[0], pair[1]])) + .collect() +} + +/// Select a rounded percentile index from nonempty sorted timings using an integer fraction. +fn percentile(sorted: &[f32], numerator: usize, denominator: usize) -> f32 { + assert!(!sorted.is_empty()); + assert!(denominator > 0 && numerator <= denominator); + let scaled = (sorted.len() - 1) + .checked_mul(numerator) + .expect("percentile index overflow"); + sorted[(scaled + denominator / 2) / denominator] +} + +/// Build packed inputs, verify the GPU result against the CPU oracle, and report CUDA-event timings. +fn main() { + println!("=== Block-scaled MXFP4 GEMM with cute layouts ==="); + println!("C({M}x{N}) = A({M}x{K} e2m1) x B({K}x{N} e2m1), ue8m0 scales per 32-element block\n"); + + let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); + + // Accelerated MXFP4 block-scaled MMA exists on SM120/121 targets. + let (major, minor) = ctx.compute_capability().expect("query compute capability"); + println!("GPU Compute Capability: sm_{major}{minor}"); + if major != 12 { + println!("skipping execution: mxf4 block-scale MMA requires sm_120/sm_121"); + return; + } + + // Host-side packed FP4 payloads; B is stored transposed as N x K. + let a_nibbles: Vec = (0..M * K).map(|i| a_nibble(i / K, i % K)).collect(); + let b_nibbles: Vec = (0..N * K).map(|i| b_nibble(i % K, i / K)).collect(); + let sfa: Vec = (0..M * KB).map(|i| sfa_byte(i / KB, i % KB)).collect(); + let sfb: Vec = (0..N * KB).map(|i| sfb_byte(i / KB, i % KB)).collect(); + + // Exact f32 accumulator reference: per K-block, an exact 32-term dot times the two + // power-of-two scales. Every partial sum is an exact binary fraction well + // inside f32's exact range, so accumulation order cannot matter. + println!("computing exact host reference ({M}x{N}x{K})..."); + let mut expected = vec![0.0f32; M * N]; + for m in 0..M { + let a_row = &a_nibbles[m * K..(m + 1) * K]; + let sfa_row = &sfa[m * KB..(m + 1) * KB]; + for n in 0..N { + let b_row = &b_nibbles[n * K..(n + 1) * K]; + let sfb_row = &sfb[n * KB..(n + 1) * KB]; + let mut total = 0.0f32; + for kb in 0..KB { + let mut block = 0.0f32; + for j in 0..32 { + let k = kb * 32 + j; + block += e2m1_decode(a_row[k]) * e2m1_decode(b_row[k]); + } + total += ue8m0_decode(sfa_row[kb]) * ue8m0_decode(sfb_row[kb]) * block; + } + expected[m * N + n] = total; + } + } + let expected_f16: Vec = expected.iter().copied().map(|value| value as f16).collect(); + + let stream = ctx.default_stream(); + // SAFETY: this executable is the sole owner of the embedded device module. + let module = unsafe { kernels::load(&ctx) }.expect("Failed to load embedded CUDA module"); + + let a_packed = pack_mxf4_view(&a_nibbles); + let b_packed = pack_mxf4_view(&b_nibbles); + let sfa_packed = pack_canonical_scales(&sfa, M); + let sfb_packed = pack_canonical_scales(&sfb, N); + let a_dev = cuda_core::DeviceBuffer::from_host(&stream, &a_packed).unwrap(); + let b_dev = cuda_core::DeviceBuffer::from_host(&stream, &b_packed).unwrap(); + let sfa_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfa_packed).unwrap(); + let sfb_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfb_packed).unwrap(); + + // The descriptor types are the host/kernel agreement: A/B land in the + // B64 byte layout, while each flattened scale row is one canonical atom. + let a_desc: TensorMap = create_tensor_map_tiled_with_options( + a_dev.cu_deviceptr() as *mut core::ffi::c_void, + M as u64, + (K / 2) as u64, + (K / 2) as u64, + TmaEncodeOptions::L2_128B, + ) + .expect("A TMA descriptor encoding failed"); + let b_desc: TensorMap = create_tensor_map_tiled_with_options( + b_dev.cu_deviceptr() as *mut core::ffi::c_void, + N as u64, + (K / 2) as u64, + (K / 2) as u64, + TmaEncodeOptions::L2_128B, + ) + .expect("B TMA descriptor encoding failed"); + let sfa_atom_rows = (M / BLOCK_M) * K_STAGES; + let sfb_atom_rows = (N / BLOCK_N) * K_STAGES; + let sfa_desc: TensorMap = create_tensor_map_tiled_with_options( + sfa_dev.cu_deviceptr() as *mut core::ffi::c_void, + sfa_atom_rows as u64, + 256, + 256, + TmaEncodeOptions::L2_128B, + ) + .expect("SFA TMA descriptor encoding failed"); + let sfb_desc: TensorMap = create_tensor_map_tiled_with_options( + sfb_dev.cu_deviceptr() as *mut core::ffi::c_void, + sfb_atom_rows as u64, + 256, + 256, + TmaEncodeOptions::L2_128B, + ) + .expect("SFB TMA descriptor encoding failed"); + let out = cuda_core::DeviceBuffer::::zeroed(&stream, M * N).unwrap(); + let out_desc: TensorMap = create_tensor_map_tiled_with_options( + out.cu_deviceptr() as *mut core::ffi::c_void, + M as u64, + N as u64, + N as u64, + TmaEncodeOptions::L2_128B, + ) + .expect("C TMA descriptor encoding failed"); + // Ask the driver how many copies of this CTA fit on each SM. The + // persistent scheduler needs that number to choose its physical grid. + let probe = module + .prepare_blockscale_gemm(LaunchConfig1D::new( + 1, + THREADS as u32, + DYNAMIC_SMEM_BYTES as u32, + )) + .expect("launch contract validation failed"); + let active_ctas = probe + .function() + .max_active_blocks_per_multiprocessor(THREADS as u32, DYNAMIC_SMEM_BYTES as u32) + .expect("query active CTAs per SM"); + let sm_count = ctx.multiprocessor_count().expect("query SM count"); + let persistent_ctas = TileScheduler::get_grid_shape(sm_count, active_ctas) + .expect("persistent grid must have nonzero u32 capacity"); + drop(probe); + + let config = LaunchConfig1D::new(persistent_ctas, THREADS as u32, DYNAMIC_SMEM_BYTES as u32); + let launch = module + .prepare_blockscale_gemm(config) + .expect("persistent launch contract validation failed"); + // SAFETY: 288-thread blocks contain eight compute warps and one producer + // warp. The grid-stride scheduler partitions all 256 output tiles into + // disjoint CTA streams, and the warp/lane mapping partitions each tile. + unsafe { + module + .blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + .expect("kernel launch failed"); + } + let got = out.to_host_vec(&stream).unwrap(); + + let mut errors = 0; + let mut nonzero = 0; + for m in 0..M { + for n in 0..N { + if expected[m * N + n] != 0.0 { + nonzero += 1; + } + if got[m * N + n].to_bits() != expected_f16[m * N + n].to_bits() { + if errors < 5 { + eprintln!( + " C({m},{n}): expected {} (0x{:04x}), got {} (0x{:04x})", + expected_f16[m * N + n], + expected_f16[m * N + n].to_bits(), + got[m * N + n], + got[m * N + n].to_bits(), + ); + } + errors += 1; + } + } + } + if errors != 0 { + println!("\n✗ FAILED: {errors} mismatches"); + std::process::exit(1); + } + println!( + "Numeric check: all {} FP16 output bit patterns exact ({nonzero} nonzero FP32 accumulators), {} K-blocks", + M * N, + KB, + ); + + // Single-kernel device timing: every sample places exactly one launch + // between CUDA events. The throughput count is 2*M*N*K FLOPs per launch. + /// Untimed kernel launches used to stabilize the event timing samples. + const WARMUP: usize = 100; + /// Number of measured kernel latencies used for percentile reporting. + const SAMPLES: usize = 31; + for _ in 0..WARMUP { + unsafe { + module + .blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + .expect("warmup launch failed"); + } + } + stream.synchronize().expect("warmup sync failed"); + + let mut times_ms = Vec::with_capacity(SAMPLES); + for _ in 0..SAMPLES { + let start = ctx + .new_event_with_flags(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create start event"); + let end = ctx + .new_event_with_flags(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create end event"); + + start.record(&stream).expect("record start event"); + unsafe { + module + .blockscale_gemm( + &stream, &launch, a_desc, b_desc, sfa_desc, sfb_desc, out_desc, + ) + .expect("timed launch failed"); + } + end.record(&stream).expect("record end event"); + times_ms.push(start.elapsed_ms(&end).expect("event timing")); + } + times_ms.sort_by(f32::total_cmp); + let median_ms = percentile(×_ms, 1, 2); + let p10_ms = percentile(×_ms, 1, 10); + let p90_ms = percentile(×_ms, 9, 10); + let tflops = (2.0 * (M * N * K) as f64) / (f64::from(median_ms) * 1e-3) / 1e12; + println!( + "Single-kernel device time: median={median_ms:.6} ms, p10={p10_ms:.6} ms, \ + p90={p90_ms:.6} ms ({SAMPLES} samples, {WARMUP} warmups); {tflops:.2} TFLOP/s" + ); + + println!("\n✓ SUCCESS: block-scaled MXFP4 GEMM is bit-exact through CuTe semantics"); +} diff --git a/cuda-oxide/cute-rs-preview/examples/common/rmsnorm.rs b/cuda-oxide/cute-rs-preview/examples/common/rmsnorm.rs new file mode 100644 index 0000000000..c1d5cae5b3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/common/rmsnorm.rs @@ -0,0 +1,266 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared FP32/BF16 RMSNorm kernels with explicit BF16 rounding boundaries. + +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; +use cute_rs::BFloat16; +use cute_rs::tensor::partition::{StridedTensor, StridedTensorMut}; + +/// Four complete warps own four independent rows in every block. +pub(crate) const ROWS_PER_BLOCK: usize = 4; + +/// Scalar row partitions and fixed-width register tiles for FP32/BF16 RMSNorm. +/// +/// Normalization entries require 128 threads per block: four complete warps, +/// one row per warp. Host code proves nonoverlapping output ownership, ready +/// input buffers, and lifetimes through completion; the types do not prove +/// collective convergence. All dimensions and strides here count elements. +#[cuda_module] +pub mod kernels { + use super::*; + + /// Normalize dense `[rows, cols]` FP32 input using one complete warp per row. + /// + /// `weight` contains `cols` shared column scales and output holds `rows * cols` + /// distinct writable elements. `cols` and finite `epsilon` must be positive. + /// All lanes, including column tails, reach the FP32 square-sum collective; + /// the output order is `(value * inverse_rms) * weight`. Launch exactly four + /// warps per block with nonoverflowing checked extents and ready sources. + #[kernel] + pub fn rmsnorm_f32( + input: &[f32], + weight: &[f32], + mut output: DisjointSlice, + rows: usize, + cols: usize, + epsilon: f32, + ) { + let tid = thread::threadIdx_x() as usize; + let row = thread::blockIdx_x() as usize * ROWS_PER_BLOCK + tid / 32; + // This branch is uniform within each warp, including the final block. + if row < rows { + let lane = tid % 32; + let base = row * cols; + let values = StridedTensor::from_slice(input) + .subview(base, cols) + .local_partition::<32>(lane); + let scales = StridedTensor::from_slice(weight) + .subview(0, cols) + .local_partition::<32>(lane); + let squares = values.fold(0.0_f32, |sum, value| sum + value * value); + // SAFETY: all 32 lanes of this row reach the collective. Lanes + // with no valid column contribute zero, instead of returning. + let total = unsafe { cute_rs::arch::warp_reduction_sum_f32(squares) }; + let inverse_rms = cute_rs::arch::rsqrt_f32(total / cols as f32 + epsilon); + let mut destination = StridedTensorMut::from_disjoint_slice(&mut output) + .subview(base, cols) + .local_partition::<32>(lane); + // SAFETY: the checked launch partitions rows by warp and each + // row by lane; sources and destination own distinct allocations. + unsafe { + values.zip_map_to(scales, &mut destination, |value, scale| { + (value * inverse_rms) * scale + }); + } + } + } + + /// Normalize dense BF16 rows, rounding before and after multiplying weights. + /// + /// Square sums and inverse RMS use FP32. Each normalized value rounds to + /// BF16, is multiplied by its BF16 column weight in FP32, then rounds again. + /// Shape, positive epsilon, complete-warps, readiness, and nonaliasing output + /// requirements match `rmsnorm_f32`; scalar partitions permit column tails. + #[kernel] + pub fn rmsnorm_bf16( + input: &[BFloat16], + weight: &[BFloat16], + mut output: DisjointSlice, + rows: usize, + cols: usize, + epsilon: f32, + ) { + let tid = thread::threadIdx_x() as usize; + let row = thread::blockIdx_x() as usize * ROWS_PER_BLOCK + tid / 32; + if row < rows { + let lane = tid % 32; + let base = row * cols; + let values = StridedTensor::from_slice(input) + .subview(base, cols) + .local_partition::<32>(lane); + let scales = StridedTensor::from_slice(weight) + .subview(0, cols) + .local_partition::<32>(lane); + let squares = values.fold(0.0_f32, |sum, value| { + let value = value.to_f32(); + sum + value * value + }); + // SAFETY: each launched block contains four complete warps; the + // row predicate is warp-uniform and column tails do not skip lanes. + let total = unsafe { cute_rs::arch::warp_reduction_sum_f32(squares) }; + let inverse_rms = cute_rs::arch::rsqrt_f32(total / cols as f32 + epsilon); + let mut destination = StridedTensorMut::from_disjoint_slice(&mut output) + .subview(base, cols) + .local_partition::<32>(lane); + // SAFETY: complete warps own disjoint rows and cyclic lane + // partitions own disjoint columns, with separate input storage. + unsafe { + values.zip_map_to(scales, &mut destination, |value, scale| { + let normalized = BFloat16::from_f32(value.to_f32() * inverse_rms); + BFloat16::from_f32(normalized.to_f32() * scale.to_f32()) + }); + } + } + } + + /// Keep each lane's values live across the reduction. Constant trip counts + /// remove repeated address/control work without changing the serial FFMA + /// order, warp XOR tree or either BF16 rounding boundary. + /// + /// Storage is dense `[rows, ITEMS * 32]`; `ITEMS` must be positive and the + /// runtime `divisor` must equal that width for RMSNorm. Extent checks are + /// warp-uniform, and the launch must supply exactly four complete warps. + /// Epsilon must be finite and positive. Ready input/weights and distinct + /// output storage must remain live through completion. + #[inline(always)] + fn rmsnorm_bf16_fixed( + input: &[BFloat16], + weight: &[BFloat16], + mut output: DisjointSlice, + rows: usize, + divisor: usize, + epsilon: f32, + ) { + let tid = thread::threadIdx_x() as usize; + let row = thread::blockIdx_x() as usize * ROWS_PER_BLOCK + tid / 32; + if row < rows + && rows <= input.len() / (ITEMS * 32) + && rows <= output.len() / (ITEMS * 32) + && weight.len() >= ITEMS * 32 + { + let lane = tid % 32; + let base = row * (ITEMS * 32); + // SAFETY: the uniform extent proof bounds both rows. The cyclic + // partition expresses the original lane + item * 32 mapping. + let (values, scales) = unsafe { + let values = StridedTensor::from_slice(input) + .subview_unchecked(base, ITEMS * 32) + .local_partition::<32>(lane) + .load_unchecked::() + .map(BFloat16::to_f32); + let scales = StridedTensor::from_slice(weight) + .subview_unchecked(0, ITEMS * 32) + .local_partition::<32>(lane) + .load_unchecked::() + .map(BFloat16::to_f32); + (values, scales) + }; + let squares = values.fold(0.0_f32, |sum, value| sum + value * value); + // SAFETY: full warps own rows, and the row predicate is uniform. + // warp_reduction_sum_f32 retains XOR offsets 16, 8, 4, 2, 1. + let total = unsafe { cute_rs::arch::warp_reduction_sum_f32(squares) }; + // Keep the divisor a runtime argument, as in the original kernel. + // Constant power-of-two division may otherwise contract with the + // epsilon addition and change subnormal rounding boundaries. + let inverse_rms = cute_rs::arch::rsqrt_f32(total / divisor as f32 + epsilon); + let normalized = values.zip_map(scales, |value, scale| { + let normalized = BFloat16::from_f32(value * inverse_rms); + BFloat16::from_f32(normalized.to_f32() * scale) + }); + // SAFETY: the checked row extent and cyclic partition assign each + // output one writer. The view consumes its parent writable borrow. + unsafe { + StridedTensorMut::from_disjoint_slice(&mut output) + .subview_unchecked(base, ITEMS * 32) + .local_partition::<32>(lane) + .store(normalized); + } + } + } + + /// Normalize dense `[rows, 128]` BF16 input with 4 values per lane. + /// + /// `cols` must equal 128; it remains a runtime mean-square divisor to + /// preserve rounding. Uses the two BF16 rounding boundaries and the same + /// four-warp launch, readiness, epsilon, and ownership contract as + /// `rmsnorm_bf16`. + #[kernel] + pub fn rmsnorm_bf16_128( + input: &[BFloat16], + weight: &[BFloat16], + output: DisjointSlice, + rows: usize, + cols: usize, + epsilon: f32, + ) { + rmsnorm_bf16_fixed::<4>(input, weight, output, rows, cols, epsilon); + } + + /// Normalize dense `[rows, 1024]` BF16 input with 32 values per lane. + /// + /// `cols` must equal 1024; it remains a runtime mean-square divisor to + /// preserve rounding. Uses the two BF16 rounding boundaries and the same + /// four-warp launch, readiness, epsilon, and ownership contract as + /// `rmsnorm_bf16`. + #[kernel] + pub fn rmsnorm_bf16_1024( + input: &[BFloat16], + weight: &[BFloat16], + output: DisjointSlice, + rows: usize, + cols: usize, + epsilon: f32, + ) { + rmsnorm_bf16_fixed::<32>(input, weight, output, rows, cols, epsilon); + } + + /// Normalize dense `[rows, 2048]` BF16 input with 64 values per lane. + /// + /// `cols` must equal 2048; it remains a runtime mean-square divisor to + /// preserve rounding. Uses the two BF16 rounding boundaries and the same + /// four-warp launch, readiness, epsilon, and ownership contract as + /// `rmsnorm_bf16`. + #[kernel] + pub fn rmsnorm_bf16_2048( + input: &[BFloat16], + weight: &[BFloat16], + output: DisjointSlice, + rows: usize, + cols: usize, + epsilon: f32, + ) { + rmsnorm_bf16_fixed::<64>(input, weight, output, rows, cols, epsilon); + } + + /// Copy BF16 bits through eight-element vector tiles, preserving NaN payloads. + /// + /// Input and output lengths must match and output must not alias input. + /// Complete tiles require 16-byte-aligned bases; scalar tails preserve every + /// bit too. Each unique linear thread index owns one tile, without collectives. + /// The host establishes source readiness and allocation lifetimes. + #[kernel] + pub fn copy_bf16(input: &[BFloat16], mut output: DisjointSlice) { + let tid = thread::index_1d().get(); + let input_tile = cute_rs::Tensor::from_slice(input) + .zipped_divide::<8>() + .slice(tid); + let mut output_tile = cute_rs::TensorMut::from_disjoint_slice(&mut output) + .zipped_divide::<8>() + .slice(tid); + if input_tile.is_full() { + // SAFETY: device buffers are aligned, and each tile spans eight + // adjacent two-byte elements with one writer and matching lengths. + unsafe { output_tile.store(input_tile.load()) }; + } else { + let mut index = input_tile.base(); + while index < input.len() { + unsafe { output_tile.store_at(index, input[index]) }; + index += 1; + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.lock new file mode 100644 index 0000000000..7cc8aa969c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.lock @@ -0,0 +1,805 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "elementwise_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.toml new file mode 100644 index 0000000000..7767fd5731 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "elementwise_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "CuTe-style f32 and f16 elementwise add with vectorized full tiles and scalar tails" + +# Mark as standalone crate (not part of parent workspace). +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs" } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } diff --git a/cuda-oxide/cute-rs-preview/examples/elementwise_cute/README.md b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/README.md new file mode 100644 index 0000000000..b50f667a3b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/README.md @@ -0,0 +1,102 @@ +# Elementwise add with Tensor tiles + +This is the smallest `cute-rs` example: + +```text +C[i] = A[i] + B[i] +``` + +It runs once with `f32` values and once with `f16` values. Each GPU thread owns +one contiguous 16-byte tile. + +`Tensor` is a read-only view. `TensorMut` is a writable view that cannot be +copied or cloned, which makes accidental overlapping output owners harder to +create. + +## The Tensor flow + +`zipped_divide` gives one flat vector two coordinates: + +```text +flat vector +[ 0 1 2 3 | 4 5 6 7 | 8 9 10 11 | ... ] + + zipped_divide::<4>() + │ + ▼ +coordinate = (place inside a tile, tile number) +``` + +The thread number selects the second coordinate. Loading the resulting slice +moves one tile into registers; storing moves the result back to global memory. + +```text +A storage ──► Tensor ──► tiled view ──► slice(thread) + │ + ▼ + load registers + │ +B storage ──► Tensor ──► tiled view ──► slice(thread) + │ + ▼ + load registers + │ + elementwise add + │ + ▼ +C storage ◄── TensorMut ◄── tiled view ◄── store registers +``` + +The Rust reads like the diagram: + +```rust +let g_a = Tensor::from_slice(a).zipped_divide::(); +let g_b = Tensor::from_slice(b).zipped_divide::(); +let g_c = TensorMut::from_disjoint_slice(&mut c).zipped_divide::(); + +let t_a = g_a.slice(tid); +let t_b = g_b.slice(tid); +let mut t_c = g_c.slice(tid); +``` + +The f32 kernel in `src/main.rs` takes its `a` and `c` views through the one-call form, +`local_tile::(tid)`, which is CuTe's `local_tile` for a one-dimensional tiler and +equals `zipped_divide::().slice(tid)`; the f16 kernel keeps the two steps so both +forms stay side by side. + +```rust + +let result = unsafe { t_a.load() } + unsafe { t_b.load() }; +unsafe { t_c.store(result) }; +``` + +## The final partial tile + +The example sizes are intentionally not exact multiples of a tile: + +```text +f32: 1003 values, 4 per tile + +... | 996 997 998 999 | 1000 1001 1002 -- | + └─ scalar tail ────┘ +``` + +Complete tiles use the tile load/add/store path. The final thread uses a small +Rust loop and touches only valid values. The `f16` run does the same thing with +eight values per tile. + +## Run it + +You need the cuda-oxide toolchain, the pinned CUTLASS compiler +(`cargo oxide toolchain install cutlass`) and an `sm_120a` or `sm_100a` GPU. From the +`cuda-oxide/` directory: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir cargo oxide run elementwise_cute --arch sm_120a +``` + +A successful run confirms: + +- all 1,003 `f32` result bit patterns match the host reference; +- all 2,005 `f16` result bit patterns match the host reference; +- both vector-tile paths and their scalar tails are exercised. diff --git a/cuda-oxide/cute-rs-preview/examples/elementwise_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/src/main.rs new file mode 100644 index 0000000000..8126eab3b8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/src/main.rs @@ -0,0 +1,175 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Elementwise addition with a short final tile, in both f32 and f16. +//! +//! Each thread owns one contiguous tile (4 x f32 or 8 x f16, both 16 +//! bytes). The kernel spells the same view flow as CuTeDSL: +//! +//! ```text +//! Tensor -> zipped_divide -> slice -> load / add / store +//! ``` +//! +//! Full tiles use one vector load or store. The final, shorter tile uses an +//! ordinary Rust branch and scalar loop. +//! +//! PASS requires bit-exact correctness for both element types on the GPU. +//! Backend-specific branches add their own code-shape and artifact gates. + +#![feature(f16)] + +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; +use cuda_core::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; + +/// Deliberately not divisible by the tiles: f32 tail of 3, f16 tail of 5. +const N_F32: usize = 1003; +/// FP16 element count, leaving five values in the final eight-value tile. +const N_F16: usize = 2005; +/// Four FP32 values per thread, totaling one 16-byte vector. +const TILE_F32: usize = 4; // 16 bytes +/// Eight FP16 values per thread, totaling one 16-byte vector. +const TILE_F16: usize = 8; // 16 bytes + +/// Device entry points for vector additions with scalar tail handling. +#[cuda_module] +mod kernels { + use super::*; + + /// Add matching FP32 tiles, using scalar stores for the short final tile. + #[kernel] + pub fn add_f32(a: &[f32], b: &[f32], mut c: DisjointSlice) { + let tid = thread::index_1d().get(); + let t_a = cute_rs::Tensor::from_slice(a).local_tile::(tid); + let g_b = cute_rs::Tensor::from_slice(b).zipped_divide::(); + let mut t_c = cute_rs::TensorMut::from_disjoint_slice(&mut c).local_tile::(tid); + + let t_b = g_b.slice(tid); + + if t_a.is_full() { + let a_values = unsafe { t_a.load() }; + let b_values = unsafe { t_b.load() }; + unsafe { t_c.store(a_values + b_values) }; + } else { + // Ragged tail: plain Rust scalar loop, boundary thread only. + let mut k = t_a.base(); + while k < a.len() { + unsafe { t_c.store_at(k, a[k] + b[k]) }; + k += 1; + } + } + } + + /// Add matching FP16 tiles, using scalar stores for the short final tile. + #[kernel] + pub fn add_f16(a: &[f16], b: &[f16], mut c: DisjointSlice) { + let tid = thread::index_1d().get(); + let g_a = cute_rs::Tensor::from_slice(a).zipped_divide::(); + let g_b = cute_rs::Tensor::from_slice(b).zipped_divide::(); + let g_c = cute_rs::TensorMut::from_disjoint_slice(&mut c).zipped_divide::(); + + let t_a = g_a.slice(tid); + let t_b = g_b.slice(tid); + let mut t_c = g_c.slice(tid); + + if t_a.is_full() { + let a_values = unsafe { t_a.load() }; + let b_values = unsafe { t_b.load() }; + unsafe { t_c.store(a_values + b_values) }; + } else { + let mut k = t_a.base(); + while k < a.len() { + unsafe { t_c.store_at(k, a[k] + b[k]) }; + k += 1; + } + } + } +} + +/// Round the work size up to a whole number of nonzero-sized tiles. +fn ceil_div(a: usize, b: usize) -> usize { + a.div_ceil(b) +} + +/// Run both elementwise kernels and require bit-exact output for every value. +fn main() { + println!("=== CuTe-style elementwise add, ragged edges, f32 + f16 ===\n"); + + let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); + let stream = ctx.default_stream(); + let module = kernels::load(&ctx).expect("Failed to load embedded CUDA module"); + + // f32 + let a32: Vec = (0..N_F32).map(|i| i as f32).collect(); + let b32: Vec = (0..N_F32).map(|i| (i * 2) as f32).collect(); + let a32_dev = DeviceBuffer::from_host(&stream, &a32).unwrap(); + let b32_dev = DeviceBuffer::from_host(&stream, &b32).unwrap(); + let mut c32_dev = DeviceBuffer::::zeroed(&stream, N_F32).unwrap(); + let threads32 = ceil_div(N_F32, TILE_F32) as u32; + // SAFETY: each thread owns one disjoint tile (or the scalar tail). + unsafe { + module.add_f32( + &stream, + LaunchConfig::for_num_elems(threads32), + &a32_dev, + &b32_dev, + &mut c32_dev, + ) + } + .expect("f32 kernel launch failed"); + let c32 = c32_dev.to_host_vec(&stream).unwrap(); + for i in 0..N_F32 { + let expected = a32[i] + b32[i]; + assert_eq!( + c32[i].to_bits(), + expected.to_bits(), + "f32 bit mismatch at {i}: expected {expected}, got {}", + c32[i], + ); + } + println!( + "f32 numeric check: all {N_F32} elements correct (tail of {})", + N_F32 % TILE_F32 + ); + + // f16 + let a16: Vec = (0..N_F16) + .map(|i| ((i % 100) as f32 / 8.0) as f16) + .collect(); + let b16: Vec = (0..N_F16).map(|i| ((i % 50) as f32 / 4.0) as f16).collect(); + let a16_dev = DeviceBuffer::from_host(&stream, &a16).unwrap(); + let b16_dev = DeviceBuffer::from_host(&stream, &b16).unwrap(); + let mut c16_dev = DeviceBuffer::::zeroed(&stream, N_F16).unwrap(); + let threads16 = ceil_div(N_F16, TILE_F16) as u32; + // SAFETY: as above. + unsafe { + module.add_f16( + &stream, + LaunchConfig::for_num_elems(threads16), + &a16_dev, + &b16_dev, + &mut c16_dev, + ) + } + .expect("f16 kernel launch failed"); + let c16 = c16_dev.to_host_vec(&stream).unwrap(); + for i in 0..N_F16 { + let expected = a16[i] + b16[i]; + assert_eq!( + c16[i].to_bits(), + expected.to_bits(), + "f16 bit mismatch at {i}: expected {}, got {}", + expected as f32, + c16[i] as f32, + ); + } + println!( + "f16 numeric check: all {N_F16} elements correct (tail of {})", + N_F16 % TILE_F16 + ); + + println!("\n✓ SUCCESS: f32 and f16 elementwise add are bit-exact!"); +} diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.lock new file mode 100644 index 0000000000..e34a1ee3fe --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "ffn_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.toml new file mode 100644 index 0000000000..ce29c2f009 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "ffn_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "Checked BF16 feed-forward composition using typed tensor-core operands" + +[workspace] + +[dependencies] +cute-rs = { path = "../../cute-rs", features = ["host"] } +ctm-rs = { path = "../../ctm-rs" } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core" } diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/README.md b/cuda-oxide/cute-rs-preview/examples/ffn_cute/README.md new file mode 100644 index 0000000000..cb6847c581 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/README.md @@ -0,0 +1,43 @@ +# A BF16 feed-forward block + +This cuda-oxide example composes the bias-free feed-forward core: + +```text + ┌─ gate projection → SiLU ─┐ +input ───────────┤ × → down projection → output + └─ up projection ──────────┘ +``` + +All stored tensors are BF16. Projections accumulate in FP32 and round to BF16. +SiLU rounds before multiplication by the up projection, then the product rounds +again. RMSNorm and residual addition belong to the surrounding decoder layer. + +## The CuTe flow + +Each warp (32 threads) uses `Bf16Mma16x8x16` partitions to compute a 16×8 +output tile, accumulate 16 input values at a time, and store its owned coordinates. Predicated loads supply zeros for +matrix tails. Weights are stored as `[output, input]` matrices. + +The activation uses `Tensor` and `TensorMut` tiles with eight BF16 values per +thread. Typed weight roles distinguish gate, up, and down projections, while an +owned workspace holds the intermediate tensors. The runner waits for pending +input work and completes all four kernels before returning. + +## Run it + +From `cuda-oxide/`, with the pinned toolchain, CUDA toolkit, the CUTLASS compiler +(`cargo oxide toolchain install cutlass`) and an RTX 5090: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run ffn_cute --arch sm_120a +``` + +A successful run prints `PASS FFN` for each shape and a final intermediate/output +count. It checks each stage against an FP64 reference, checks guards and tails, +and exercises the exact SiLU rounding boundary and pending input writes. + +See [numerics and validation](reference/validation.md) for shape contracts, +tolerances, host tests, and sanitizer commands. This is synthetic FFN +correctness evidence; checkpoint inference and model performance remain +separate requirements. diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/reference/validation.md b/cuda-oxide/cute-rs-preview/examples/ffn_cute/reference/validation.md new file mode 100644 index 0000000000..4d932876a0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/reference/validation.md @@ -0,0 +1,72 @@ +# Contracts and validation + +[Back to the example](../README.md). + +## Arithmetic contract + +Standalone cuda-oxide correctness example for the bias-free FFN core +`down(silu(gate(x)) * up(x))`. This composes three BF16 tensor-core projections +and one activation kernel. FP32 projection accumulators round to BF16; SiLU +rounds to BF16 before multiplying the BF16 up projection, and the product rounds +again. RMSNorm and residual addition belong to the surrounding decoder layer. + +## Storage and execution + +The host-only `cute_rs::ffn` API carries gate/up/down roles and input/output +widths in Rust types. Weights borrow existing `DeviceBuffer` allocations +with `[output,input]` orientation. Dynamic shapes validate row strides, storage +capacity, arithmetic overflow and addressable byte extents. Owned workspace +buffers keep mutable intermediates distinct; input and weight borrows remain +live through execution. Role tags cannot identify incorrectly loaded weights. + +The example's synchronous runner checks buffer contexts and launch geometry, +waits for prior work in the owning context, then queues all four kernels on one +stream and waits for completion, including enqueue-error paths. The initial +context wait makes pending writes on other streams safe without adding a new +readiness-tracking framework. It is intentionally conservative for throughput. +Zero rows perform no launches. + +Each full warp computes a 16×8 output tile using typed `Bf16A`/`Bf16B` operands, +FP32 accumulators and K16 MMA steps. Predicated global loads handle row, column +and K tails and padded strides. Warp participation remains an explicit unsafe +contract. Activation uses existing typed 128-bit copies for complete tiles and +scalar copies for tails. This first schedule has no shared-memory pipeline. + +## Host and GPU checks + +From `cuda-oxide/`: + +```bash +cargo test --offline -p cute-rs --features host --all-targets +cargo test --offline -p cute-rs --features host --doc +cargo test --offline --manifest-path cute-rs-preview/examples/ffn_cute/Cargo.toml +CARGO_NET_OFFLINE=true CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run ffn_cute --arch sm_120a +/usr/local/cuda/bin/compute-sanitizer --tool memcheck --error-exitcode 1 \ + cute-rs-preview/examples/ffn_cute/target/release/ffn_cute +/usr/local/cuda/bin/compute-sanitizer --tool synccheck --error-exitcode 1 \ + cute-rs-preview/examples/ffn_cute/target/release/ffn_cute +``` + +## Acceptance criteria and scope + +Synthetic validation covers hidden/intermediate widths 1024/3072 and 2048/6144 +at one and three rows, plus 19/23, 32/17 and 1/1 cases, empty rows, and a +17-row tile tail. Asymmetric data and NaN padding expose weight-orientation and +stride errors. Every intermediate and final output is compared with an +independent FP64 reference with integer BF16 rounding; output guards must stay +untouched. The down projection is also checked on its actual GPU input to +separate projection errors from propagated activation error. + +For these bounded finite fixtures, gate/up/SiLU/product use +`atol=1e-6, rtol=0.008`; down uses `atol=1e-5, rtol=0.008`; the complete +composition uses `atol=2e-5, rtol=0.02`. A separate exact fixture requires product +bits `0x3fbc`, rejecting the incorrect cast-once result `0x3fbd`. A delayed +producer stream regression checks that safe asynchronous input writes complete +before the runner reads them. Host tests cover shape and grid rejection, and +compile-fail tests cover dtype, weight-role, dimension and borrow contracts. + +Validation runs on local RTX 5090 / SM120 using the pinned CUTLASS compiler. +This is synthetic FFN correctness evidence; checkpoint inference, model-level +performance, and execution on other architectures remain separate acceptance +requirements. diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/kernel.rs b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/kernel.rs new file mode 100644 index 0000000000..34c4e37351 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/kernel.rs @@ -0,0 +1,434 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! BF16 linear and activation kernels used by the checked host composition. +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, launch_contract, thread}; +use cute_rs::BFloat16; + +/// Device entries for BF16 projections and activation, shared by FFN and decoder examples. +/// +/// Matrix inputs use element strides; weights store `[output, input]`, while +/// outputs are dense. Host runners must prove bounds, context readiness, distinct +/// mutable storage, and convergent warp participation before launching. +#[cuda_module] +pub mod kernels { + use super::*; + use cute_rs::mma::bf16::{Bf16Mma16x8x16, FragmentC, PackedDecode}; + use cute_rs::numeric::silu_bf16 as silu; + + /// Read `(row, col)` from an element-strided row-major BF16 matrix. + /// + /// Missing logical rows or columns return positive zero for MMA tail padding. + /// In-range coordinates require `data` to cover `row * stride + col`; this + /// helper does not independently validate the allocation or row stride. + #[inline(always)] + fn element( + data: &[BFloat16], + row: usize, + col: usize, + rows: usize, + cols: usize, + stride: usize, + ) -> BFloat16 { + if row < rows && col < cols { + data[row * stride + col] + } else { + BFloat16::from_bits(0) + } + } + + /// Project a short input batch using checked packed BF16 MMA operands. + /// + /// One 32-thread CTA owns eight output columns; K16 MMA steps accumulate in + /// FP32. When enabled, bias is added in FP32 before one BF16 output rounding. + /// + /// # Safety + /// The entire warp must enter with uniform dimensions and the same validated + /// `source`. `width` and `outputs` must match that view. Output must contain + /// every dense `[source rows, outputs]` element, remain distinct from inputs, + /// and have only this grid's partitioned writers. With `HAS_BIAS`, `bias` + /// must contain `outputs` initialized elements. Sources must be ready, and + /// all allocations must remain live until the kernel finishes. + #[inline(always)] + unsafe fn packed_decode_impl( + source: &PackedDecode<'_>, + bias: &[BFloat16], + output: &mut DisjointSlice, + width: usize, + outputs: usize, + ) { + let lane = thread::threadIdx_x() as usize; + let n = thread::blockIdx_x() as usize * 8; + if n >= outputs { + return; + } + let mma = Bf16Mma16x8x16::get_slice(lane); + let a = mma.partition_a(); + let b = mma.partition_b(); + let mut accumulator = FragmentC::ZERO; + let mut k = 0; + while k < width { + // SAFETY: PackedDecode validates source extents/alignment; N and K + // select full tiles. Every lane reaches the same instruction order. + unsafe { + if width - k >= 512 { + // Fetch 32 K16 operand pairs into a typed register tile so + // independent loads can overlap before their consumers. + // The fold visits ascending K: each MMA sees exactly the + // preceding accumulator, preserving FP32 rounding order. + let operands = cute_rs::tensor::RegisterTile::<_, 32>::from_fn(|step| { + ( + a.load_decode_packed(source, k + step * 16), + b.load_decode_packed(source, n, k + step * 16), + ) + }); + accumulator = + operands.fold(accumulator, |acc, (av, bv)| mma.mma_atom_call(acc, av, bv)); + k += 512; + } else { + let av = a.load_decode_packed(source, k); + let bv = b.load_decode_packed(source, n, k); + accumulator = mma.mma_atom_call(accumulator, av, bv); + k += 16; + } + } + } + let mut output = cute_rs::tensor::partition::StridedTensorMut::from_disjoint_slice(output); + mma.partition_c() + .store_decode_with(accumulator, source, |row, col, value| { + let value = if HAS_BIAS { + value + bias[n + col].to_f32() + } else { + value + }; + // SAFETY: C's bijective layout and the checked grid give each lane + // disjoint initialized output coordinates, with one BF16 rounding. + unsafe { + output.store_at(row * outputs + n + col, BFloat16::from_f32(value)); + } + }); + } + + /// Project strided BF16 matrices through typed 16×8×16 warp MMA partitions. + /// + /// Input is `[rows, width]`, weights are `[outputs, width]`, and strides count + /// elements. Missing M/N/K coordinates contribute zero; output is dense + /// `[rows, outputs]`. `PREFETCH` loads four K16 steps ahead without changing + /// accumulation order. Bias, when enabled, is added before BF16 rounding. + /// + /// The enclosing launch must provide exactly `32 * WARPS_PER_BLOCK` threads, + /// complete convergent warps, valid strided source extents, and distinct + /// writable output storage. `outputs` must be positive, grid arithmetic must + /// fit, and `HAS_BIAS` requires at least `outputs` bias values. + #[inline(always)] + #[allow(clippy::too_many_arguments)] + fn linear_bf16_impl< + const HAS_BIAS: bool, + const WARPS_PER_BLOCK: usize, + const PREFETCH: bool, + >( + input: &[BFloat16], + weight: &[BFloat16], + bias: &[BFloat16], + mut output: DisjointSlice, + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) { + let tid = thread::threadIdx_x() as usize; + let lane = tid % 32; + let tile = thread::blockIdx_x() as usize * WARPS_PER_BLOCK + tid / 32; + let columns = outputs.div_ceil(8); + if tile < rows.div_ceil(16) * columns { + let m = tile / columns * 16; + let n = tile % columns * 8; + let mma = Bf16Mma16x8x16::get_slice(lane); + let load_a = |k| { + mma.partition_a() + .load_with(|r, c| element(input, m + r, k + c, rows, width, input_stride)) + }; + let load_b = |k| { + mma.partition_b() + .load_with(|r, c| element(weight, n + r, k + c, outputs, width, weight_stride)) + }; + let mut accumulator = FragmentC::ZERO; + let mut k = 0; + while k < width { + if PREFETCH && width - k >= 64 { + let a0 = load_a(k); + let b0 = load_b(k); + let a1 = load_a(k + 16); + let b1 = load_b(k + 16); + let a2 = load_a(k + 32); + let b2 = load_b(k + 32); + let a3 = load_a(k + 48); + let b3 = load_b(k + 48); + // SAFETY: the full warp uses the atom's typed fragments; + // arithmetic consumes K tiles in unchanged increasing order. + unsafe { + accumulator = mma.mma_atom_call(accumulator, a0, b0); + accumulator = mma.mma_atom_call(accumulator, a1, b1); + accumulator = mma.mma_atom_call(accumulator, a2, b2); + accumulator = mma.mma_atom_call(accumulator, a3, b3); + } + k += 64; + continue; + } + let a = load_a(k); + let b = load_b(k); + // SAFETY: no lane skips the collective; tails are zero-padded. + accumulator = unsafe { mma.mma_atom_call(accumulator, a, b) }; + k += 16; + } + let mut output = + cute_rs::tensor::partition::StridedTensorMut::from_disjoint_slice(&mut output); + mma.partition_c().store_with(accumulator, |r, c, value| { + if m + r < rows && n + c < outputs { + let value = if HAS_BIAS { + value + bias[n + c].to_f32() + } else { + value + }; + // SAFETY: C's bijective layout and the host-checked tile grid + // give disjoint output writes within the validated extent. + unsafe { + output.store_at((m + r) * outputs + n + c, BFloat16::from_f32(value)); + } + } + }); + } + } + + /// Compute a bias-free projection with FP32 accumulation and BF16 output. + /// + /// Launch 128 threads per block, four complete warps, with one warp per 16×8 + /// output tile. The caller validates `[rows, width]` input, `[outputs, width]` + /// weights, element strides, and dense `[rows, outputs]` nonaliasing output. + /// Inputs must be ready and all borrows retained until completion; tails use + /// zero padding without skipping any lane of a participating warp. + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn linear_bf16( + input: &[BFloat16], + weight: &[BFloat16], + output: DisjointSlice, + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) { + // The const-false specialization never reads the placeholder bias slice. + linear_bf16_impl::( + input, + weight, + input, + output, + rows, + width, + outputs, + input_stride, + weight_stride, + ); + } + + /// Add the projection bias in FP32 before rounding the output once to BF16. + /// + /// Uses the same 128-thread launch, strided matrix layout, readiness, and + /// disjoint-output requirements as `linear_bf16`. Bias contains one BF16 + /// value per output column and remains live throughout the projection. + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn linear_bias_bf16( + input: &[BFloat16], + weight: &[BFloat16], + bias: &[BFloat16], + output: DisjointSlice, + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) { + linear_bf16_impl::( + input, + weight, + bias, + output, + rows, + width, + outputs, + input_stride, + weight_stride, + ); + } + + /// The same MMA arithmetic with one warp per CTA, distributing short-row + /// decode projections over more SMs without changing reduction order. + /// + /// The 32-thread launch contract supplies one full warp. Checked packed + /// operands are used when available; otherwise the same element-strided + /// `[rows, width]` / `[outputs, width]` matrices use predicated loads. The + /// caller must validate those fallback extents, dense output capacity, and + /// readiness and ownership exactly as for `linear_bf16`. + #[launch_contract(domain=1, block=(32,1,1), min_compute_capability=(8,0))] + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn decode_linear_bf16( + input: &[BFloat16], + weight: &[BFloat16], + mut output: DisjointSlice, + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) { + if outputs != 0 + && rows <= output.len() / outputs + && PackedDecode::with_validated( + input, + weight, + rows, + width, + outputs, + input_stride, + weight_stride, + |source| { + // SAFETY: the borrowed proof establishes packed operand + // bounds; output capacity is checked above, and this entry's + // launch contract supplies exactly one complete warp. + unsafe { + packed_decode_impl::(source, input, &mut output, width, outputs); + } + }, + ) + { + return; + } + linear_bf16_impl::( + input, + weight, + input, + output, + rows, + width, + outputs, + input_stride, + weight_stride, + ); + } + + /// Decode projection with bias added to FP32 accumulators before BF16 output. + /// + /// Uses the one-warp packed-or-predicated schedule of `decode_linear_bf16`. + /// Bias must contain `outputs` initialized BF16 values; all matrix bounds, + /// source readiness, output nonaliasing, and full-warp participation apply. + #[launch_contract(domain=1, block=(32,1,1), min_compute_capability=(8,0))] + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn decode_linear_bias_bf16( + input: &[BFloat16], + weight: &[BFloat16], + bias: &[BFloat16], + mut output: DisjointSlice, + rows: usize, + width: usize, + outputs: usize, + input_stride: usize, + weight_stride: usize, + ) { + if outputs != 0 + && rows <= output.len() / outputs + && bias.len() >= outputs + && PackedDecode::with_validated( + input, + weight, + rows, + width, + outputs, + input_stride, + weight_stride, + |source| { + // SAFETY: the source proof, output/bias extent checks and + // exact full-warp launch contract establish packed loads, + // disjoint stores, and collective participation. + unsafe { + packed_decode_impl::(source, bias, &mut output, width, outputs); + } + }, + ) + { + return; + } + linear_bf16_impl::( + input, + weight, + bias, + output, + rows, + width, + outputs, + input_stride, + weight_stride, + ); + } + + /// Compute BF16 SiLU and gated products with a distinct rounding after each. + /// + /// The first `length` values of gate/up inputs and both outputs must exist + /// and use distinct output allocations. Each linear thread owns eight BF16 + /// elements; complete tiles require 16-byte-aligned storage and tails use + /// scalar accesses. No collectives occur. The caller establishes readiness + /// and keeps all allocations live until completion. + #[kernel] + pub fn activate_bf16( + gate: &[BFloat16], + up: &[BFloat16], + mut silu_output: DisjointSlice, + mut product: DisjointSlice, + length: usize, + ) { + let tid = thread::index_1d().get(); + let gates = cute_rs::Tensor::from_slice(gate) + .zipped_divide::<8>() + .slice(tid); + let ups = cute_rs::Tensor::from_slice(up) + .zipped_divide::<8>() + .slice(tid); + let mut silus = cute_rs::TensorMut::from_disjoint_slice(&mut silu_output) + .zipped_divide::<8>() + .slice(tid); + let mut products = cute_rs::TensorMut::from_disjoint_slice(&mut product) + .zipped_divide::<8>() + .slice(tid); + let base = gates.base(); + if base < length && length - base >= 8 { + // SAFETY: checked buffers contain `length` values; bases are + // aligned multiples of 8, and each thread owns one disjoint tile. + let values = unsafe { gates.load() }.map(silu); + let multiplied = values.zip_map(unsafe { ups.load() }, |s, u| { + BFloat16::from_f32(s.to_f32() * u.to_f32()) + }); + unsafe { + silus.store(values); + products.store(multiplied); + } + } else { + let mut index = base; + while index < length { + let s = silu(gate[index]); + let value = BFloat16::from_f32(s.to_f32() * up[index].to_f32()); + unsafe { + silus.store_at(index, s); + products.store_at(index, value); + } + index += 1; + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/main.rs new file mode 100644 index 0000000000..96879e5401 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/main.rs @@ -0,0 +1,217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Checked BF16 feed-forward composition and intermediate numerical validation. +/// BF16 matrix projections and tiled SiLU/product kernels. +mod kernel; +/// Checked synchronous FFN execution and owned intermediate storage. +mod operator; +/// Independent host matrix and activation oracles with integer BF16 rounding. +mod reference; + +use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; +use cute_rs::{ + BFloat16, + ffn::{Down, FfnShape, FfnWeights, Gate, ProjectionWeights, RowMajorShape, Up}, +}; +use operator::{GUARD, Result, Runner, Workspace}; + +/// Asymmetric bounded dyadic values keep the projection reference well-conditioned. +fn matrix(shape: RowMajorShape, seed: usize) -> Vec { + let mut data = vec![BFloat16::from_bits(GUARD); shape.required_len()]; + for row in 0..shape.rows() { + for col in 0..shape.cols() { + let code = ((row * 37 + col * 13 + (row ^ col) * 3 + seed * 11) % 17) as i32 - 8; + let scale = if seed == 0 { 8.0 } else { 128.0 }; + data[row * shape.row_stride() + col] = reference::round(code as f32 / scale); + } + } + data +} + +/// Copy every allocated element and check that no logical write touched the guard. +fn read_output( + stream: &CudaStream, + data: &DeviceBuffer, + logical: usize, +) -> Result> { + let mut actual = vec![BFloat16::from_bits(0); data.len()]; + data.copy_to_host(stream, &mut actual)?; + assert!( + actual[logical..].iter().all(|v| v.to_bits() == GUARD), + "output guard overwritten" + ); + actual.truncate(logical); + Ok(actual) +} + +/// Compare independent reference values, allowing one BF16 output quantum plus cancellation error. +fn verify(label: &str, actual: &[BFloat16], expected: &[BFloat16], atol: f64, rtol: f64) -> f64 { + assert_eq!(actual.len(), expected.len()); + let mut maximum = 0.0_f64; + for (i, (&a, &e)) in actual.iter().zip(expected).enumerate() { + let a = reference::widen(a); + let e = reference::widen(e); + let error = (a - e).abs(); + assert!( + a.is_finite() && e.is_finite() && error <= atol + rtol * e.abs(), + "{label} index={i}: actual={a} expected={e} error={error}" + ); + maximum = maximum.max(error); + } + maximum +} + +/// Allocate checked weights once and execute all requested batches at the same model widths. +fn validate( + runner: &Runner, + stream: &CudaStream, + rows: &[usize], + padding: usize, +) -> Result { + let gate_shape = RowMajorShape::new(I, H, H + padding)?; + let up_shape = RowMajorShape::new(I, H, H + padding + 1)?; + let down_shape = RowMajorShape::new(H, I, I + padding + 2)?; + let gate = matrix(gate_shape, 1); + let up = matrix(up_shape, 2); + let down = matrix(down_shape, 3); + let gate_device = DeviceBuffer::from_host(stream, &gate)?; + let up_device = DeviceBuffer::from_host(stream, &up)?; + let down_device = DeviceBuffer::from_host(stream, &down)?; + let weights = FfnWeights::::new( + ProjectionWeights::::new(&gate_device, gate_shape.row_stride())?, + ProjectionWeights::::new(&up_device, up_shape.row_stride())?, + ProjectionWeights::::new(&down_device, down_shape.row_stride())?, + ); + // A safe checked constructor must reject an overlapping stride before any launch. + assert!(ProjectionWeights::::new(&gate_device, H - 1).is_err()); + let mut checked = 0; + for &rows in rows { + let shape = FfnShape::::new(rows, H + padding)?; + let input = matrix(shape.input(), 0); + let input_device = DeviceBuffer::from_host(stream, &input)?; + let mut workspace = Workspace::new(stream, shape)?; + runner.run(shape, &input_device, &weights, &mut workspace)?; + let gates = read_output(stream, &workspace.gate, shape.intermediate().required_len())?; + let ups = read_output(stream, &workspace.up, shape.intermediate().required_len())?; + let silu = read_output(stream, &workspace.silu, shape.intermediate().required_len())?; + let product = read_output( + stream, + &workspace.product, + shape.intermediate().required_len(), + )?; + let output = read_output(stream, &workspace.output, shape.hidden().required_len())?; + let expected_gate = reference::linear( + &input, + &gate, + rows, + H, + I, + shape.input().row_stride(), + gate_shape.row_stride(), + ); + let expected_up = reference::linear( + &input, + &up, + rows, + H, + I, + shape.input().row_stride(), + up_shape.row_stride(), + ); + let (expected_silu, expected_product) = reference::activate(&expected_gate, &expected_up); + let expected_output = reference::linear( + &expected_product, + &down, + rows, + I, + H, + I, + down_shape.row_stride(), + ); + let mut maximum = 0.0_f64; + for (label, actual, expected) in [ + ("gate", &gates, &expected_gate), + ("up", &ups, &expected_up), + ("silu", &silu, &expected_silu), + ("product", &product, &expected_product), + ] { + maximum = maximum.max(verify(label, actual, expected, 1e-6, 0.008)); + } + // Check down independently on its actual input as well as the complete + // composition, so propagated activation error cannot hide a bad GEMM. + let local_down = reference::linear(&product, &down, rows, I, H, I, down_shape.row_stride()); + maximum = maximum.max(verify("down", &output, &local_down, 1e-5, 0.008)); + maximum = maximum.max(verify("end-to-end", &output, &expected_output, 2e-5, 0.02)); + checked += 4 * rows * I + rows * H; + println!( + "PASS FFN H={H} I={I} rows={rows} input_stride={} max_abs_error={maximum:.8e}", + shape.input().row_stride() + ); + } + Ok(checked) +} + +/// Safe asynchronous buffer writes can precede a call on a different stream. +fn check_pending_input(runner: &Runner, stream: &CudaStream) -> Result<()> { + use std::sync::mpsc; + use std::time::Duration; + + let producer = stream.context().new_stream()?; + let one = [BFloat16::from_bits(0x3f80)]; + let mut input = DeviceBuffer::from_host(stream, &one)?; + let weight = DeviceBuffer::from_host(stream, &one)?; + let weights = FfnWeights::<1, 1>::new( + ProjectionWeights::::new(&weight, 1)?, + ProjectionWeights::::new(&weight, 1)?, + ProjectionWeights::::new(&weight, 1)?, + ); + let shape = FfnShape::new(1, 1)?; + let mut workspace = Workspace::new(stream, shape)?; + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + producer.launch_host_function(move || { + let _ = entered_tx.send(()); + // No CUDA calls in the callback. Bound the wait even on a test failure. + let _ = release_rx.recv_timeout(Duration::from_secs(10)); + })?; + entered_rx.recv_timeout(Duration::from_secs(10))?; + input.zero_async(&producer)?; + assert!(!producer.query()?, "producer must still be pending"); + let release = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(250)); + let _ = release_tx.send(()); + }); + let result = runner.run(shape, &input, &weights, &mut workspace); + let completion = producer.synchronize(); + release.join().expect("producer release thread panicked"); + result?; + completion?; + let output = read_output(stream, &workspace.output, 1)?; + assert_eq!( + output[0].to_bits(), + 0, + "FFN read input before its producer completed" + ); + println!("PASS: pending input on a separate stream is ordered before FFN"); + Ok(()) +} + +/// Validate every FFN stage, guarded output, and producer ordering on CUDA device 0. +fn main() -> Result<()> { + let context = CudaContext::new(0)?; + let stream = context.default_stream(); + let runner = Runner::new(stream.clone())?; + check_pending_input(&runner, &stream)?; + let mut checked = validate::<19, 23>(&runner, &stream, &[0, 1, 17], 5)?; + checked += validate::<32, 17>(&runner, &stream, &[5], 3)?; + checked += validate::<1, 1>(&runner, &stream, &[3], 0)?; + checked += validate::<1024, 3072>(&runner, &stream, &[1, 3], 0)?; + checked += validate::<2048, 6144>(&runner, &stream, &[1, 3], 0)?; + runner.check_rounding()?; + println!( + "PASS: {checked} BF16 FFN intermediate/output values; guards, tails, empty rows and exact SiLU rounding fixture passed" + ); + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/operator.rs b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/operator.rs new file mode 100644 index 0000000000..7252cdc432 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/operator.rs @@ -0,0 +1,265 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Checked composition using existing buffer ownership and typed weight roles. +use crate::kernel::kernels; +use cuda_core::LaunchConfig; +use cuda_core::{CudaStream, DeviceBuffer}; +use cute_rs::{ + BFloat16, + ffn::{FfnShape, FfnWeights}, +}; +use std::sync::Arc; + +/// Host result carrying CUDA failures and rejected shape or context contracts. +pub type Result = std::result::Result>; +/// BF16 quiet-NaN bit pattern used to detect writes outside logical tensors. +pub const GUARD: u16 = 0x7fc1; +/// Number of extra BF16 sentinel elements appended to each workspace buffer. +pub const GUARD_LEN: usize = 8; + +/// Reusable owned intermediate buffers; const widths prevent cross-model reuse. +pub struct Workspace { + /// Batch row count that a runner must match when reusing this workspace. + rows: usize, + /// Dense `[rows, I]` gate projection, rounded to BF16, plus trailing guards. + pub gate: DeviceBuffer, + /// Dense `[rows, I]` up projection, rounded to BF16, plus trailing guards. + pub up: DeviceBuffer, + /// Dense `[rows, I]` BF16 SiLU values before multiplication, plus guards. + pub silu: DeviceBuffer, + /// Dense `[rows, I]` BF16 SiLU/up products used by the down projection, plus guards. + pub product: DeviceBuffer, + /// Dense `[rows, H]` BF16 down-projection result, plus trailing guards. + pub output: DeviceBuffer, +} +impl Workspace { + /// Allocate separate sentinel-filled buffers for this batch and the const widths. + pub fn new(stream: &CudaStream, shape: FfnShape) -> Result { + let alloc = |len: usize| -> Result> { + let total = len.checked_add(GUARD_LEN).ok_or("guard extent overflow")?; + Ok(DeviceBuffer::from_host( + stream, + &vec![BFloat16::from_bits(GUARD); total], + )?) + }; + Ok(Self { + rows: shape.rows(), + gate: alloc(shape.intermediate().required_len())?, + up: alloc(shape.intermediate().required_len())?, + silu: alloc(shape.intermediate().required_len())?, + product: alloc(shape.intermediate().required_len())?, + output: alloc(shape.hidden().required_len())?, + }) + } +} + +/// Return a 128-thread grid with four warps, each owning one 16×8 output tile. +/// +/// Dimensions are element counts. Zero rows return `None`; zero output width, +/// tile-count overflow, and grids beyond CUDA's signed grid-x limit are errors. +/// This check does not validate matrix allocation lengths or strides. +pub fn linear_grid(rows: usize, outputs: usize) -> Result> { + if outputs == 0 { + return Err("linear output width must be positive".into()); + } + if rows == 0 { + return Ok(None); + } + let tiles = rows + .div_ceil(16) + .checked_mul(outputs.div_ceil(8)) + .ok_or("tile count overflow")?; + let blocks = tiles.div_ceil(4); + if blocks > i32::MAX as usize { + return Err("linear grid exceeds CUDA grid-x range".into()); + } + Ok(Some(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + })) +} + +/// The module and stream are bound to the same context at construction. +pub struct Runner { + /// Kernel bundle loaded in the same CUDA context as the execution stream. + module: kernels::LoadedModule, + /// Stream ordering all projections and activation; each run waits for completion. + stream: Arc, +} +impl Runner { + /// Load this example's kernels and retain the stream used for synchronous execution. + pub fn new(stream: Arc) -> Result { + Ok(Self { + // SAFETY: the loaded artifact is generated from this example's + // kernel source, including the checked one-warp decode entries. + module: unsafe { kernels::load(stream.context())? }, + stream, + }) + } + /// Execute synchronously, ordering prior work in the owning context and + /// retaining all buffer borrows through completion, including enqueue errors. + /// + /// Input uses the shape's row stride; weights retain `[output, input]` + /// orientation. Intermediates and output are dense, with BF16 rounding after + /// each projection, after SiLU, and after the SiLU/up product. Shape, buffer + /// capacity, context, and launch errors are rejected before execution. + /// A validated zero-row batch launches nothing. + pub fn run( + &self, + shape: FfnShape, + input: &DeviceBuffer, + weights: &FfnWeights<'_, H, I>, + workspace: &mut Workspace, + ) -> Result<()> { + shape.input().validate_len(input.len())?; + if workspace.rows != shape.rows() { + return Err("workspace batch does not match operation".into()); + } + for buffer in [ + &workspace.gate, + &workspace.up, + &workspace.silu, + &workspace.product, + ] { + shape.intermediate().validate_len(buffer.len())?; + } + shape.hidden().validate_len(workspace.output.len())?; + for buffer in [ + input, + weights.gate().buffer(), + weights.up().buffer(), + weights.down().buffer(), + &workspace.gate, + &workspace.up, + &workspace.silu, + &workspace.product, + &workspace.output, + ] { + if !Arc::ptr_eq(buffer.context(), self.stream.context()) { + return Err("FFN buffers must share the runner context".into()); + } + } + let Some(gate_grid) = linear_grid(shape.rows(), I)? else { + return Ok(()); + }; + let down_grid = linear_grid(shape.rows(), H)?.ok_or("missing down grid")?; + let length = shape.intermediate().required_len(); + let activation_tiles = length.div_ceil(8); + let activation_blocks = activation_tiles.div_ceil(128); + if activation_blocks > i32::MAX as usize { + return Err("activation grid exceeds CUDA grid-x range".into()); + } + let activation_grid = LaunchConfig { + grid_dim: (activation_blocks as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + }; + // A buffer can have pending safe asynchronous writes on another stream. + // This correctness-first runner has no readiness tracking, so establish + // readiness for every borrowed buffer before launching any readers. + self.stream.context().synchronize()?; + // SAFETY: checked extents, typed weight roles, distinct owned workspace + // allocations and complete warps. All launches use one stream/context; + // the unconditional synchronization below retains borrows even on error. + let result = (|| -> Result<()> { + unsafe { + self.module.linear_bf16( + &self.stream, + gate_grid, + input, + weights.gate().buffer(), + &mut workspace.gate, + shape.rows(), + H, + I, + shape.input().row_stride(), + weights.gate().shape().row_stride(), + )?; + self.module.linear_bf16( + &self.stream, + gate_grid, + input, + weights.up().buffer(), + &mut workspace.up, + shape.rows(), + H, + I, + shape.input().row_stride(), + weights.up().shape().row_stride(), + )?; + self.module.activate_bf16( + &self.stream, + activation_grid, + &workspace.gate, + &workspace.up, + &mut workspace.silu, + &mut workspace.product, + length, + )?; + self.module.linear_bf16( + &self.stream, + down_grid, + &workspace.product, + weights.down().buffer(), + &mut workspace.output, + shape.rows(), + I, + H, + I, + weights.down().shape().row_stride(), + )?; + } + Ok(()) + })(); + let completion = self.stream.synchronize(); + result?; + completion?; + Ok(()) + } + /// Exercise the exact SiLU rounding discriminator through the real kernel. + pub fn check_rounding(&self) -> Result<()> { + let gate = DeviceBuffer::from_host(&self.stream, &[BFloat16::from_bits(0x3f80)])?; + let up = DeviceBuffer::from_host(&self.stream, &[BFloat16::from_bits(0x4001)])?; + let mut silu = DeviceBuffer::from_host(&self.stream, &[BFloat16::from_bits(GUARD)])?; + let mut product = DeviceBuffer::from_host(&self.stream, &[BFloat16::from_bits(GUARD)])?; + unsafe { + self.module.activate_bf16( + &self.stream, + LaunchConfig::for_num_elems(1), + &gate, + &up, + &mut silu, + &mut product, + 1, + )?; + } + let mut actual = [BFloat16::from_bits(0)]; + silu.copy_to_host(&self.stream, &mut actual)?; + assert_eq!(actual[0].to_bits(), 0x3f3b); + product.copy_to_host(&self.stream, &mut actual)?; + assert_eq!( + actual[0].to_bits(), + 0x3fbc, + "cast-once would produce 0x3fbd" + ); + Ok(()) + } +} + +/// CPU checks of projection tile counts and CUDA grid overflow rejection. +#[cfg(test)] +mod tests { + use super::*; + /// Check partial 16×8 tiles, empty batches, invalid widths, and overflow rejection. + #[test] + fn launch_geometry_handles_tails_and_rejects_overflow() { + assert!(linear_grid(0, 1).unwrap().is_none()); + assert!(linear_grid(1, 0).is_err()); + assert_eq!(linear_grid(17, 9).unwrap().unwrap().grid_dim, (1, 1, 1)); + assert_eq!(linear_grid(33, 17).unwrap().unwrap().grid_dim, (3, 1, 1)); + assert!(linear_grid(usize::MAX, usize::MAX).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/reference.rs b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/reference.rs new file mode 100644 index 0000000000..667243fef2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/src/reference.rs @@ -0,0 +1,113 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent FP64 matrix arithmetic and BF16 rounding for synthetic fixtures. +use cute_rs::BFloat16; + +/// Integer RNE oracle independent of the device conversion helper. +pub fn round(value: f32) -> BFloat16 { + let bits = value.to_bits(); + let high = (bits >> 16) as u16; + if value.is_nan() { + return BFloat16::from_bits(high | 0x40); + } + let low = bits & 0xffff; + BFloat16::from_bits( + high.wrapping_add(u16::from(low > 0x8000 || (low == 0x8000 && high & 1 != 0))), + ) +} +/// Widen stored bits independently from the device implementation. +pub fn widen(value: BFloat16) -> f64 { + f64::from(f32::from_bits(u32::from(value.to_bits()) << 16)) +} +/// Compute X times transposed stored weights using sequential FP64 dot products. +/// +/// Input is `[rows, input_width]`; weights are `[output_width, input_width]`. +/// Both strides count BF16 elements. The caller supplies buffers covering every +/// indexed row; the returned `[rows, output_width]` matrix is dense. Each dot +/// product narrows through FP32 before integer round-to-nearest-even BF16. +#[allow(clippy::too_many_arguments)] +pub fn linear( + input: &[BFloat16], + weight: &[BFloat16], + rows: usize, + input_width: usize, + output_width: usize, + input_stride: usize, + weight_stride: usize, +) -> Vec { + let mut result = Vec::with_capacity(rows * output_width); + for row in 0..rows { + for out in 0..output_width { + let mut sum = 0.0_f64; + for k in 0..input_width { + sum += + widen(input[row * input_stride + k]) * widen(weight[out * weight_stride + k]); + } + result.push(round(sum as f32)); + } + } + result +} +/// Compute SiLU in FP64; preserve its BF16 output before the BF16 product. +pub fn activate(gate: &[BFloat16], up: &[BFloat16]) -> (Vec, Vec) { + assert_eq!(gate.len(), up.len()); + let silu: Vec<_> = gate + .iter() + .map(|&g| { + let g = widen(g); + let s = if g >= 0.0 { + 1.0 / (1.0 + (-g).exp()) + } else { + let e = g.exp(); + e / (1.0 + e) + }; + round((g * s) as f32) + }) + .collect(); + let product = silu + .iter() + .zip(up) + .map(|(&s, &u)| round((widen(s) * widen(u)) as f32)) + .collect(); + (silu, product) +} + +/// CPU fixtures exposing transposed weights, BF16 ties, and activation rounding. +#[cfg(test)] +mod tests { + use super::*; + /// Check `[output, input]` orientation while ignoring NaN-padded row strides. + #[test] + fn asymmetric_output_by_input_weights_and_padding_are_observable() { + let x = [1., 2., 3., f32::NAN, 4., 5., 6., f32::NAN].map(round); + let w = [1., -2., 3., f32::NAN, 2., 1., -1., f32::NAN].map(round); + assert_eq!( + linear(&x, &w, 2, 3, 2, 4, 4) + .into_iter() + .map(widen) + .collect::>(), + [6., 1., 12., 7.] + ); + assert!(linear(&[], &w, 0, 3, 2, 3, 4).is_empty()); + } + /// Check even-tie BF16 rounding, signed zero, and NaN preservation. + #[test] + fn independent_rounding_covers_ties_and_special_values() { + assert_eq!(round(f32::from_bits(0x3f80_8000)).to_bits(), 0x3f80); + assert_eq!(round(f32::from_bits(0x3f81_8000)).to_bits(), 0x3f82); + assert_eq!(round(-0.0).to_bits(), 0x8000); + assert!(widen(round(f32::from_bits(0x7f80_0001))).is_nan()); + } + /// Distinguish the BF16 SiLU intermediate from rounding the full product only once. + #[test] + fn activation_keeps_the_silu_rounding_boundary() { + let (silu, product) = activate(&[round(1.0)], &[BFloat16::from_bits(0x4001)]); + assert_eq!(silu[0].to_bits(), 0x3f3b); + assert_eq!(product[0].to_bits(), 0x3fbc); + let fused = + round((1.0_f64 / (1.0 + (-1.0_f64).exp()) * widen(BFloat16::from_bits(0x4001))) as f32); + assert_ne!(product[0].to_bits(), fused.to_bits()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.lock new file mode 100644 index 0000000000..74d4dad6af --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "fmha_decode_sm120a_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", + "serde_json", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.toml new file mode 100644 index 0000000000..1879dcad0a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "fmha_decode_sm120a_cute" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "SM120a: Dedicated split-KV FP16 attention decode with typed CuTe copies" + +[workspace] + +[dependencies] +cute-rs = { path = "../../cute-rs" } +ctm-rs = { path = "../../ctm-rs" } +# CTM reexports kernel macros, whose expansions still name ::cuda_device. +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } +serde_json = "1" diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/README.md b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/README.md new file mode 100644 index 0000000000..6aeeab2f70 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/README.md @@ -0,0 +1,70 @@ +# Single-token decode with cute-rs — SM120a (RTX 5090) + +This cuda-oxide example computes attention for one query per head against an +existing FP16 KV cache. A dedicated kernel processes 64 cache entries per block; +a second kernel combines the partial softmax results and writes FP16 output. + +```text +Q/K/V tensor views → typed copies → FP32 dot products and local softmax + → split states → merge kernel → FP16 output +``` + +- Batch **1**, head dimension **64**, contiguous BSHD tensors, FP32 accumulation. +- Any positive cache length is supported, including partial final tiles. +- `--specialization auto` (default) uses static `(Sk,H)` shapes **(64,8), (65,3), + (1024,8), (4097,8)** when available; otherwise it uses runtime dimensions. + `static` requires one of those shapes; `runtime` keeps dimensions dynamic. +- The cache contains only visible tokens. Cache management, paging and grouped-query + attention are outside this example. + +From `cuda-oxide/`, once per checkout: + +```bash +cargo oxide toolchain install cutlass +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Compare runtime and static specialization on RTX 5090: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ +CUTE_DSL_ARCH=sm_120a CUTE_DSL_COMPILER_OPT=--opt-level=2 \ + cargo oxide run fmha_decode_sm120a_cute --arch sm_120a -- \ + --all-cases --warmup 100 --iters 31 +``` + +Both implementations check every output against a CPU FP64 oracle. The command +also verifies matching input hashes and output parity, then prints two rows of +CUDA-event medians at the default **Sk=1024, H=8** shape. Each interval covers +the partial and merge kernels. `--all-cases` requires a listed static shape and +overrides `--specialization`; attention is noncausal and single-token. + +For one case, replace `--all-cases` with `--specialization static` or `runtime`. +Add `--json PATH` to save the comparison, `--iters 0` to verify without timing, +or `--rust-only` for a standalone Rust report. `--all-cases` cannot be combined +with `--rust-only`. Pass `--help` after `--` for all options. + +Use `--python PATH` to select another interpreter and `--device N` to select +another visible GPU; device 0 is the default. +The CuTeDSL reference is independently authored here using the same split-KV +algorithm; it is **not an upstream CUTLASS decode kernel**. + +RTX 5090 results at `a0d5776b2` (2026-10-18), **B=1, Sq=1, Sk=1024, H=8, D=64**, +CUTLASS 4.7 at optimization level 2 on both sides. Each value is the median of +**31 single-invocation samples after 100 warmups**, with reused buffers. +Both implementations passed correctness checks. Times are microseconds: + +| Shapes | CUDA events: Rust (cute-rs) | CUDA events: CuTeDSL | Nsight kernel sum: Rust (cute-rs) | Nsight kernel sum: CuTeDSL | +| --- | ---: | ---: | ---: | ---: | +| Runtime | 4.768 | 11.808 | 3.392 | 3.520 | +| Static | 4.416 | 11.808 | 2.880 | 2.656 | + +Both metrics include the partial and merge kernels. CUDA events include gaps +during host submission; Nsight Systems sums the two GPU execution durations, +excluding those gaps but including device stalls. These are separate runs; +subtracting columns does not isolate host overhead. Clocks were unlocked and +CuTeDSL ran first in each case. This single run per method does not establish +a stable speed ranking. + +- [CuTeDSL reference, setup and validation commands](reference/README.md) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/README.md b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/README.md new file mode 100644 index 0000000000..90dcf19022 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/README.md @@ -0,0 +1,96 @@ +# Dedicated CuTeDSL decode reference + +`attention_decode.py` is authored in this repository using the public CUTLASS +CuTeDSL API. It is **not an upstream CUTLASS kernel**. It implements the same +one-query split-KV schedule as the Rust example: 64 cache rows per 128-thread +CTA, FP32 dot products and stable softmax, FP32 split states, and a second +32-thread-per-head merge kernel producing FP16 output. + +`launch.py` owns the decode configuration, allocations and both launches. +It imports only deterministic fixture and CPU-oracle helpers from the prefill +reference. It does not import or call the prefill GPU kernel. Cache lengths +need not be multiples of 64; invalid keys are masked out of softmax. + +`--specialization auto|runtime|static` matches the Rust dispatch. Auto specializes +`(Sk,H) = (64,8), (65,3), (1024,8), (4097,8)` and otherwise keeps dimensions +runtime. Static mode rejects other shapes; runtime mode also keeps the KV +stride and merge bounds dynamic. The comparison requires the same effective +mode on both sides, so neither compiler receives an extra shape advantage. + +Use `nvidia-cutlass-dsl[cu13]==4.7.0` with a compatible CUDA-enabled PyTorch +build and driver. The shared [`requirements.txt`](../../../requirements.txt) +pins the DSL; dependencies retain their own licenses. All code in this +directory has the indicated first-party Apache-2.0 license. + +Prepare the shared Python environment once for both prefill and decode, +from `cuda-oxide/`: + +```bash +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Select a CUDA-enabled PyTorch build that supports SM120 and matches your CUDA +runtime. The normal `cargo oxide run` command finds the shared environment +automatically; no activation or Python environment-variable export is needed. +Python selection is `--python PATH` after `--`, `CUDA_OXIDE_CUTE_PYTHON`, the +example's `.venv/bin/python`, the shared `cute-rs-preview/.venv/bin/python`, then `python3`. +See the [example README](../README.md) for direct-event comparisons. +The main command with `--all-cases --warmup 100 --iters 31` prints runtime +and static direct CUDA-event medians for the selected `(Sk,H)` shape. Decode +has only noncausal single-token attention, so there are two rows. This mode +requires one of the static shapes above and overrides `--specialization`. +Each event interval covers both kernels, their gap and possible host launch +submission delays. It differs from Nsight's sum of kernel-active durations +and from the span between the first kernel's start and the last kernel's end. +Comparing those metrics does not establish a regression. + +`compare.py` retains the separate Nsight Systems profiling workflow and +invokes the Rust binary with `--rust-only`. With Nsight Systems installed, +reproduce the archived runtime/static measurement protocol from `cuda-oxide/`: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build --arch sm_120a --device-codegen-crate fmha_decode_sm120a_cute \ + --cargo-target-dir /tmp/cuda-oxide-fmha-decode -- \ + --manifest-path cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.toml --release + +for specialization in runtime static; do + CUTE_DSL_ARCH=sm_120a cute-rs-preview/.venv/bin/python \ + cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/compare.py \ + --rust-binary /tmp/cuda-oxide-fmha-decode/release/fmha_decode_sm120a_cute \ + --seq-k 1024 --heads 8 --specialization "$specialization" --noncausal \ + --warmup 100 --samples 31 --rounds 3 \ + --output-dir "/tmp/decode-nsys-$specialization" +done +``` + +The harness reports the median of round medians for active sum and sequence +span separately and retains raw evidence in fresh directories. Hardware state, +clocks and software versions can affect new results; reproducing the protocol +does not guarantee the archived numbers. + +From this directory: + +```bash +CUTE_DSL_ARCH=sm_120a ../../../.venv/bin/python launch.py \ + --seq-k 65 --heads 3 --warmup 0 --iters 0 +../../../.venv/bin/python -m unittest discover -s . -p 'test_*.py' +``` + +`benchmark.py` accepts the same options. Its event measurements include the +whole partial-plus-merge sequence. `compare.py` runs both implementations +serially under Nsight Systems and distinguishes summed kernel-active time +from the elapsed two-kernel span. `validate.py` covers short, ragged and long +caches, multiple heads, uniform attention and large logits. + +Each implementation checks every output against an independent FP64 stable +softmax with `abs(actual-expected) <= 2e-4 + 2e-3 * abs(expected)`, poisons +the output and partial buffers before correctness launches, and requires a +second launch to reproduce the output bit for bit. From the example directory, +with the shared environment's interpreter as `python`: + +```bash +CUTE_DSL_ARCH=sm_120a python reference/validate.py \ + --binary /path/to/fmha_decode_sm120a_cute --reference --output-dir /tmp/decode-validation +``` diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/attention_decode.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/attention_decode.py new file mode 100644 index 0000000000..8157a247ea --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/attention_decode.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Independently authored CuTeDSL counterpart of the cuda-oxide decode example. + +This is not an upstream CUTLASS attention kernel. It implements the same +single-query, 64-row split-KV schedule as the Rust example using public CuTeDSL +copy/layout primitives and scalar FP32 arithmetic. No query padding, matrix +multiply instructions, or tensor memory are used. +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync +from cutlass.utils import SmemAllocator + + +HEAD_DIM = 64 +KV_CHUNK = 64 +PARTIAL_THREADS = 128 +MERGE_THREADS = 32 +PARTIAL_SHARED_BYTES = 17696 +SCALE_LOG2 = 0.125 * 1.4426950408889634 + + +class AttentionDecode: + """FP16 BSHD MHA with B=Sq=1, D=64 and a positive dense KV-cache length.""" + + def __init__(self, static_seq_k=0, static_heads=0): + if (static_seq_k == 0) != (static_heads == 0) or min(static_seq_k, static_heads) < 0: + raise ValueError("static dimensions must both be positive, or both zero for runtime") + self.static_seq_k = static_seq_k + self.static_heads = static_heads + + @cute.jit + def __call__( + self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, + partial: cute.Tensor, out: cute.Tensor, + seq_k: cutlass.Int32, heads: cutlass.Int32, splits: cutlass.Int32, + stream: cuda.CUstream, + ): + shared_layout = cute.make_composed_layout( + cute.make_swizzle(3, 3, 3), 0, + cute.make_layout((64, 64), stride=(64, 1)), + ) + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float16, num_bits_per_copy=128, + ) + copy = cute.make_tiled_copy_tv( + copy_atom, cute.make_layout((16, 8), stride=(8, 1)), + cute.make_layout((1, 8), stride=(8, 1)), + ) + self.fmha_decode_partial(q, k, v, partial, seq_k, heads, shared_layout, copy).launch( + grid=(heads * splits, 1, 1), + block=(PARTIAL_THREADS, 1, 1), stream=stream, + ) + self.fmha_decode_merge(partial, out, heads, splits).launch( + grid=(heads, 1, 1), block=(MERGE_THREADS, 1, 1), stream=stream, + ) + + @cute.kernel + def fmha_decode_partial( + self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, + partial: cute.Tensor, seq_k: cutlass.Int32, heads: cutlass.Int32, + shared_layout: cute.ComposedLayout, copy: cute.TiledCopy, + ): + tid, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + lane = tid % 32 + warp = tid // 32 + if cutlass.const_expr(self.static_seq_k > 0): + seq_k = self.static_seq_k + heads = self.static_heads + splits = (self.static_seq_k + KV_CHUNK - 1) // KV_CHUNK + kv_stride = self.static_heads * HEAD_DIM + state_stride = splits * 66 + else: + splits = (seq_k + KV_CHUNK - 1) // KV_CHUNK + kv_stride = cute.assume(cutlass.Int64(heads) * HEAD_DIM, divby=8) + state_stride = cutlass.Int64(splits) * 66 + head = block // splits + split = block % splits + start = split * KV_CHUNK + + # The selected mode matches Rust: specialize both dimensions together, + # or preserve runtime dimensions, strides and split bounds. + qo_layout = cute.make_layout((1, 1, heads, 64), stride=(0, 0, 64, 1)) + kv_layout = cute.make_layout( + (1, seq_k, heads, 64), + stride=(0, kv_stride, 64, 1), + ) + state_layout = cute.make_layout( + (heads, splits, 66), stride=(state_stride, 66, 1), + ) + q = cute.make_tensor(q.iterator, qo_layout) + k = cute.make_tensor(k.iterator, kv_layout) + v = cute.make_tensor(v.iterator, kv_layout) + partial = cute.make_tensor(partial.iterator, state_layout) + + smem = SmemAllocator() + sk = smem.allocate_tensor(cutlass.Float16, shared_layout, 1024) + sv = smem.allocate_tensor(cutlass.Float16, shared_layout, 1024) + scores = smem.allocate_tensor(cutlass.Float32, cute.make_layout(64), 4) + warp_max = smem.allocate_tensor(cutlass.Float32, cute.make_layout(4), 4) + warp_sum = smem.allocate_tensor(cutlass.Float32, cute.make_layout(4), 4) + warp_output = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((4, 64), stride=(64, 1)), 4, + ) + + gk = cute.local_tile(k[0, None, head, None], (64, 64), (split, 0)) + gv = cute.local_tile(v[0, None, head, None], (64, 64), (split, 0)) + thr_copy = copy.get_slice(tid) + tgk = thr_copy.partition_S(gk) + tgv = thr_copy.partition_S(gv) + tsk = thr_copy.partition_D(sk) + tsv = thr_copy.partition_D(sv) + # Four 16-row groups, each thread moving one aligned 16-byte vector. + # Invalid rows are zeroed; their global pointers are never read. + for group in cutlass.range_constexpr(4): + if start + group * 16 + tid // 8 < seq_k: + cute.copy(copy, tgk[None, group, None], tsk[None, group, None]) + cute.copy(copy, tgv[None, group, None], tsv[None, group, None]) + else: + tsk[None, group, None].fill(0.0) + tsv[None, group, None].fill(0.0) + cute.arch.cp_async_commit_group() + + rq = cute.make_rmem_tensor(cute.make_layout(16), cutlass.Float32) + for d in cutlass.range_constexpr(16): + rq[d] = q[0, 0, head, (tid % 4) * 16 + d].to(cutlass.Float32) + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # Four threads cooperate on one key; two passes cover the 64-row tile. + for key_pass in cutlass.range_constexpr(2): + row = tid // 4 + key_pass * 32 + dot = cutlass.Float32(0.0) + for d in cutlass.range_constexpr(16): + kval = sk[row, (tid % 4) * 16 + d].to(cutlass.Float32) + dot = dot + rq[d] * kval + dot = dot + cute.arch.shuffle_sync_bfly(dot, offset=1) + dot = dot + cute.arch.shuffle_sync_bfly(dot, offset=2) + if tid % 4 == 0: + if start + row < seq_k: + scores[row] = dot + else: + scores[row] = cutlass.Float32(float("-inf")) + cute.arch.sync_threads() + + score = cutlass.Float32(float("-inf")) + if tid < KV_CHUNK: + score = scores[tid] + maximum = score + for offset in cutlass.range_constexpr(5): + maximum = cute.arch.fmax( + maximum, cute.arch.shuffle_sync_bfly(maximum, offset=16 >> offset), + ) + if lane == 0: + warp_max[warp] = maximum + cute.arch.sync_threads() + maximum = cute.arch.fmax( + cute.arch.fmax(warp_max[0], warp_max[1]), + cute.arch.fmax(warp_max[2], warp_max[3]), + ) + weight = cute.math.exp2((score - maximum) * SCALE_LOG2, fastmath=True) + if tid < KV_CHUNK: + scores[tid] = weight + denominator = weight + for offset in cutlass.range_constexpr(5): + denominator = denominator + cute.arch.shuffle_sync_bfly( + denominator, offset=16 >> offset, + ) + if lane == 0: + warp_sum[warp] = denominator + cute.arch.sync_threads() + denominator = warp_sum[0] + for w in cutlass.range_constexpr(1, 4): + denominator = denominator + warp_sum[w] + + # Each warp accumulates 16 cache rows; each lane owns two dimensions. + p0 = cutlass.Float32(0.0) + p1 = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(16): + row = warp * 16 + r + p = scores[row] + p0 = p0 + p * sv[row, lane * 2].to(cutlass.Float32) + p1 = p1 + p * sv[row, lane * 2 + 1].to(cutlass.Float32) + warp_output[warp, lane * 2] = p0 + warp_output[warp, lane * 2 + 1] = p1 + cute.arch.sync_threads() + if warp == 0: + p0 = warp_output[0, lane * 2] + p1 = warp_output[0, lane * 2 + 1] + for w in cutlass.range_constexpr(1, 4): + p0 = p0 + warp_output[w, lane * 2] + p1 = p1 + warp_output[w, lane * 2 + 1] + partial[head, split, lane * 2 + 2] = p0 + partial[head, split, lane * 2 + 3] = p1 + if lane == 0: + partial[head, split, 0] = maximum + partial[head, split, 1] = denominator + + @cute.kernel + def fmha_decode_merge( + self, partial: cute.Tensor, out: cute.Tensor, + heads: cutlass.Int32, splits: cutlass.Int32, + ): + lane, _, _ = cute.arch.thread_idx() + head, _, _ = cute.arch.block_idx() + if cutlass.const_expr(self.static_seq_k > 0): + heads = self.static_heads + splits = (self.static_seq_k + KV_CHUNK - 1) // KV_CHUNK + state_stride = splits * 66 + else: + state_stride = cutlass.Int64(splits) * 66 + state_layout = cute.make_layout( + (heads, splits, 66), stride=(state_stride, 66, 1), + ) + partial = cute.make_tensor(partial.iterator, state_layout) + out = cute.make_tensor( + out.iterator, cute.make_layout((1, 1, heads, 64), stride=(0, 0, 64, 1)), + ) + maximum = cutlass.Float32(float("-inf")) + for split in cutlass.range(splits): + maximum = cute.arch.fmax(maximum, partial[head, split, 0]) + denominator = cutlass.Float32(0.0) + p0 = cutlass.Float32(0.0) + p1 = cutlass.Float32(0.0) + for split in cutlass.range(splits): + weight = cute.math.exp2( + (partial[head, split, 0] - maximum) * SCALE_LOG2, fastmath=True, + ) + denominator = denominator + partial[head, split, 1] * weight + p0 = p0 + partial[head, split, lane * 2 + 2] * weight + p1 = p1 + partial[head, split, lane * 2 + 3] * weight + inverse = cute.arch.rcp_approx(denominator) + out[0, 0, head, lane * 2] = (p0 * inverse).to(cutlass.Float16) + out[0, 0, head, lane * 2 + 1] = (p1 * inverse).to(cutlass.Float16) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/benchmark.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/benchmark.py new file mode 100644 index 0000000000..628ae938e8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/benchmark.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checked direct-event timing for the dedicated split-KV decode kernel sequence.""" + +from launch import parse_args, run + + +if __name__ == "__main__": + run(parse_args()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/compare.py new file mode 100644 index 0000000000..14c2a12ea5 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/compare.py @@ -0,0 +1,446 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compare dedicated split-KV decode sequences with serial Nsight Systems runs. + +The CuTeDSL comparator is an independently authored cuda-oxide example, not an +official upstream decode benchmark. Both implementations launch partial then +merge on one stream. Report active duration (sum of both kernels) and wall span +(first start through last end) separately. Select the GPU externally through +CUDA_VISIBLE_DEVICES. Defaults: 100 warmups, 31 samples, three alternating rounds. +""" + +import argparse +from collections import Counter +from datetime import datetime, timezone +import importlib.util +import json +import math +import os +from pathlib import Path +import shutil +import sqlite3 +import statistics +import subprocess +import sys +import tempfile +import time + + +HERE = Path(__file__).resolve().parent +KERNEL_SEQUENCE = ["fmha_decode_partial", "fmha_decode_merge"] +CORRECTNESS_LAUNCHES = 2 +KV_CHUNK = 64 +STATIC_PORTFOLIO = ((64, 8), (65, 3), (1024, 8), (4097, 8)) + + +def load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Reuse only fixture and process/provenance helpers. Isolate the prefill +# comparator's absolute launch import so importing this module cannot silently +# select the decode launch module or alter a caller's existing launch module. +_prefill = HERE.parent.parent / "fmha_prefill_sm120a_cute/reference" +fixture = load_module("decode_compare_fixture", _prefill / "launch.py") +_previous_launch = sys.modules.get("launch") +try: + sys.modules["launch"] = fixture + common = load_module("decode_compare_common", _prefill / "compare.py") +finally: + if _previous_launch is None: + sys.modules.pop("launch", None) + else: + sys.modules["launch"] = _previous_launch + +expected_input_hashes = common.expected_input_hashes +file_identity = common.file_identity +hash_f16_bytes = common.hash_f16_bytes +reject_nonfinite_json = common.reject_nonfinite_json +require_fields = common.require_fields +run_command = common.run_command + + +def finite_json_float(value): + number = float(value) + if not math.isfinite(number): + raise ValueError(f"nonfinite JSON number: {value}") + return number + + +def cpu_oracle(args): + return fixture.reference_cpu( + *fixture.make_inputs(1, args.seq_k, args.heads, args.input_scale, args.zero_query), + 1, args.seq_k, args.heads, False) + + +def effective_specialization(seq_k, heads, requested): + if requested not in ("auto", "runtime", "static"): + raise ValueError(f"unknown specialization mode: {requested!r}") + supported = (seq_k, heads) in STATIC_PORTFOLIO + if requested == "static" and not supported: + raise ValueError(f"static specialization is unavailable for (seq_k, heads)=({seq_k}, {heads})") + return "static" if supported and requested != "runtime" else "runtime" + + +def validate_report(path, output_path, args, started_ns, input_hashes, expected_output=None): + """Require fresh, full CPU/replay evidence bound to exactly these FP16 bytes.""" + for filename in (path, output_path): + if filename.stat().st_mtime_ns < started_ns: + raise ValueError(f"stale runner evidence: {filename}") + result = json.loads(path.read_text(), parse_constant=reject_nonfinite_json, parse_float=finite_json_float) + data = output_path.read_bytes() + require_fields(result, { + "workload": "fmha-decode", "shape": [1, 1, args.seq_k, args.heads, 64], + "causal": False, "input_scale": args.input_scale, "zero_query": args.zero_query, + "verification": "cpu-pass", "replay_bitwise_equal": True, "input_hashes": input_hashes, + "metric_unit": "us", + }, "report") + events = result.get("samples_us") + if (not isinstance(events, list) or len(events) != args.samples + or any(type(value) not in (float, int) or not math.isfinite(value) or value <= 0 for value in events)): + raise ValueError("runner must report one positive finite event sample per measured kernel sequence") + metric = result.get("metric_value") + if (events and (type(metric) not in (float, int) or metric != statistics.median(events)) + or not events and metric is not None): + raise ValueError("runner event metric must equal the sample median, or null when no timing was requested") + count = args.heads * 64 + validation = result.get("validation") + require_fields(validation, {"elements": count, "atol": fixture.ATOL, "rtol": fixture.RTOL, + "oracle": "cpu-fp64-stable-softmax"}, "validation") + for field in ("max_abs_error", "max_tolerance_ratio"): + value = validation.get(field) + if type(value) not in (float, int) or not math.isfinite(value) or value < 0: + raise ValueError(f"invalid validation.{field}: {value!r}") + if validation["max_tolerance_ratio"] > 1: + raise ValueError("CPU verification tolerance exceeded") + provenance = result.get("provenance") + require_fields(provenance, { + "device": args.device, "warmup": args.warmup, "iters": args.samples, "dtype": "float16", + "tool": "cuda-events-kernel-sequence", "timing_mode": "direct", + "algorithm": "single-query-split-kv", "kv_chunk": KV_CHUNK, + "split_count": (args.seq_k + KV_CHUNK - 1) // KV_CHUNK, + "kernel_count": 2, "kernel_sequence": KERNEL_SEQUENCE, + "shape_specialization": effective_specialization(args.seq_k, args.heads, args.specialization), + }, "provenance") + capability = provenance.get("compute_capability") + if (not isinstance(provenance.get("gpu"), str) or not provenance["gpu"] + or not isinstance(capability, list) or len(capability) != 2 + or any(type(value) is not int or value < 0 for value in capability)): + raise ValueError("missing GPU identity/compute capability in runner report") + if len(data) != count * 2: + raise ValueError(f"output contains {len(data)} bytes; expected {count * 2}") + output = fixture.decode_f16(data) + if not all(map(math.isfinite, output)): + raise ValueError("nonfinite FP16 output") + if result.get("output_hash") != hash_f16_bytes(data): + raise ValueError("FP16 output does not match the checked report's output hash") + fixture.verify(output, cpu_oracle(args) if expected_output is None else expected_output) + return result, output + + +def sequence_samples(path, warmup, samples, *, expected_names=None, name_contains=None): + """Resolve two kernel identities and accept only complete serial sequences. + +The trace must expose stream/device identity. Every stage has exactly two +correctness launches plus warmups and samples. Auxiliary initialization work +may occur outside a pair, but never within a pair or the measured interval. +""" + if bool(expected_names) == bool(name_contains): + raise ValueError("provide two exact kernel names or two name resolvers") + selectors = expected_names or name_contains + if len(selectors) != 2 or any(not isinstance(name, str) or not name.strip() for name in selectors): + raise ValueError("two nonempty kernel selectors are required") + if warmup < 0 or samples <= 0: + raise ValueError("warmup must be nonnegative and samples positive") + try: + with sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True) as db: + columns = {row[1] for row in db.execute("PRAGMA table_info(CUPTI_ACTIVITY_KIND_KERNEL)")} + if not {"start", "end", "shortName", "deviceId", "streamId"} <= columns: + raise ValueError("kernel trace lacks timestamps, names, deviceId or streamId") + full_names = "demangledName" in columns + name_column = "d.value" if full_names else "s.value" + name_join = "LEFT JOIN StringIds AS d ON d.id=k.demangledName " if full_names else "" + context_column = "k.contextId" if "contextId" in columns else "0" + rows = db.execute( + "SELECT k.start, k.end, s.value, " + name_column + ", k.deviceId, k.streamId, " + + context_column + " FROM CUPTI_ACTIVITY_KIND_KERNEL AS k " + "LEFT JOIN StringIds AS s ON s.id=k.shortName " + name_join + "ORDER BY k.start, k.end" + ).fetchall() + gpu_metadata = [] + if db.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='TARGET_INFO_GPU'").fetchone(): + cursor = db.execute("SELECT * FROM TARGET_INFO_GPU") + fields = [column[0] for column in cursor.description] + gpu_metadata = [dict(zip(fields, [value.hex() if isinstance(value, bytes) else value + for value in row])) for row in cursor] + except sqlite3.Error as error: + raise ValueError(f"cannot read kernel trace {path}: {error}") from error + for start, end, short, full, *_ in rows: + if any(not isinstance(name, str) or not name for name in (short, full)): + raise ValueError("unresolved kernel shortName/demangledName in trace") + if type(start) is not int or type(end) is not int or end <= start: + raise ValueError(f"invalid kernel timestamps: {start!r}, {end!r}") + identities = [] + for selector in selectors: + matches = sorted({row[3] for row in rows if + (selector in row[3] if name_contains else selector in (row[2], row[3]))}) + if len(matches) != 1: + raise ValueError(f"expected one matching kernel name for {selector!r}, found {matches!r}") + identities.append(matches[0]) + if len(set(identities)) != 2: + raise ValueError("partial and merge must resolve to distinct kernels") + selected = [row for row in rows if row[3] in identities] + expected = CORRECTNESS_LAUNCHES + warmup + samples + counts = Counter(row[3] for row in selected) + if any(counts[name] != expected for name in identities): + raise ValueError(f"expected {expected} launches of each kernel, found {dict(counts)!r}") + for field, index in (("device", 4), ("stream", 5), ("context", 6)): + values = {row[index] for row in selected} + if len(values) != 1 or any(type(value) is not int or value < 0 for value in values): + raise ValueError(f"selected sequence has invalid or mismatched {field} identities: {values!r}") + if [row[3] for row in selected] != identities * expected: + raise ValueError("kernel sequence must alternate partial then merge") + if any(previous[1] > following[0] for previous, following in zip(selected, selected[1:])): + raise ValueError("kernel sequence launches overlap; serial timing is required") + pairs = list(zip(selected[::2], selected[1::2])) + auxiliary = [row for row in rows if row[3] not in identities] + split = CORRECTNESS_LAUNCHES + warmup + intervals = [(partial[0], merge[1]) for partial, merge in pairs] + intervals.append((pairs[split][0][0], pairs[-1][1][1])) + for row in auxiliary: + if any(row[0] < end and row[1] > start for start, end in intervals): + raise ValueError(f"auxiliary kernel {row[3]!r} intersects a decode sequence or measured interval") + active = [(partial[1] - partial[0] + merge[1] - merge[0]) / 1000 for partial, merge in pairs] + wall = [(merge[1] - partial[0]) / 1000 for partial, merge in pairs] + result = { + "kernel_names": identities, "kernel_name_field": "demangledName" if full_names else "shortName", + "kernel_launches": dict(counts), "sequence_count": expected, + "trace_device_id": selected[0][4], "trace_stream_id": selected[0][5], + "trace_context_id": selected[0][6] if "contextId" in columns else None, + "trace_gpu_metadata": gpu_metadata, + "selected_kernel_timestamps_ns": [{"name": row[3], "start": row[0], "end": row[1]} for row in selected], + "auxiliary_kernel_counts": dict(sorted(Counter(row[3] for row in auxiliary).items())), + "individual_kernel_samples_us": {name: [(row[1] - row[0]) / 1000 for row in selected + if row[3] == name][split:] for name in identities}, + } + for metric, values in (("active", active), ("wall", wall)): + result[f"correctness_{metric}_samples_us"] = values[:CORRECTNESS_LAUNCHES] + result[f"warmup_{metric}_samples_us"] = values[CORRECTNESS_LAUNCHES:split] + result[f"{metric}_samples_us"] = values[split:] + return result + + +def aggregate_rounds(entries): + return {label: {metric: common.aggregate_rounds([ + {"label": label, "samples_us": entry[f"{metric}_samples_us"], + "median_us": statistics.median(entry[f"{metric}_samples_us"])} + for entry in entries if entry["label"] == label])[label] + for metric in ("active", "wall")} for label in dict.fromkeys(entry["label"] for entry in entries)} + + +def parse_arguments(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rust-binary", type=Path, required=True) + parser.add_argument("--rust-build-manifest", type=Path, + help="optional build JSON containing matching binary_sha256") + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--reference-runner", type=Path, default=HERE / "launch.py") + parser.add_argument("--rust-kernel-names", nargs=2, metavar=("PARTIAL", "MERGE"), + help="exact trace names; defaults resolve unique fmha_decode_partial/merge substrings") + parser.add_argument("--reference-kernel-names", nargs=2, metavar=("PARTIAL", "MERGE"), + help="exact trace names; defaults resolve fmha_decode_partial/merge substrings") + parser.add_argument("--nsys", default=shutil.which("nsys") or "nsys") + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seq-q", type=int, default=1, choices=(1,)) + parser.add_argument("--seq-k", type=int, default=1024) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--specialization", choices=("auto", "runtime", "static"), default="auto", + help="match both implementations' shape specialization; auto uses the shared static portfolio") + parser.add_argument("--noncausal", dest="causal", action="store_false", default=False) + parser.add_argument("--input-scale", type=int, choices=(1, 4, 16), default=1) + parser.add_argument("--zero-query", action="store_true") + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--samples", "--iters", type=int, default=31) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args(argv) + if min(args.seq_k, args.heads, args.samples, args.rounds) <= 0: + parser.error("KV length, heads, samples and rounds must be positive") + if args.seq_k * args.heads * 64 > 0xFFFFFFFF: + parser.error("tensor element count must fit u32 indexing") + if args.device < 0 or args.warmup < 0 or not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("device/warmup must be nonnegative and timeout finite and positive") + try: + effective_specialization(args.seq_k, args.heads, args.specialization) + except ValueError as error: + parser.error(str(error)) + if any(not value.strip() for values in (args.rust_kernel_names or [], args.reference_kernel_names or []) + for value in values): + parser.error("kernel names must be nonempty") + for field in ("rust_binary", "reference_runner", "rust_build_manifest"): + path = getattr(args, field) + if path is not None: + if not path.is_file() or field == "rust_binary" and not os.access(path, os.X_OK): + parser.error(f"{field} must be an existing {'executable ' if field == 'rust_binary' else ''}file") + setattr(args, field, path.resolve()) + return args + + +def main(argv=None): + args = parse_arguments(argv) + args.output_dir.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="decode-", dir=args.output_dir.resolve())) + environment = os.environ.copy() + report = { + "status": "running", "metric_unit": "us", "workload": "fmha-decode", + "metrics": {"active": "sum of partial and merge CUPTI kernel durations", + "wall": "merge end minus partial start, including inter-kernel gap"}, + "duration_source": "CUPTI_ACTIVITY_KIND_KERNEL timestamps (nanoseconds / 1000)", + "reference_origin": "independently authored cuda-oxide CuTeDSL decode example; not an official upstream decode baseline", + "started_ns": time.time_ns(), "started_utc": datetime.now(timezone.utc).isoformat(), + "parameters": {key: getattr(args, key) for key in ("seq_q", "seq_k", "heads", "causal", + "input_scale", "zero_query", "device", "specialization", "warmup", "samples", "rounds", "timeout")}, + "shape_specialization": effective_specialization(args.seq_k, args.heads, args.specialization), + "static_portfolio": [list(shape) for shape in STATIC_PORTFOLIO], + "launch_contract": {"correctness_sequences": 2, "warmup_sequences": args.warmup, + "measured_sequences": args.samples, "kernels_per_sequence": 2, + "kernel_sequence": KERNEL_SEQUENCE}, + "environment": {key: environment.get(key) for key in ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", + "CUTE_DSL_ARCH", "CUDA_HOME", "CUDA_PATH")}, + "aggregation": "median per round, then median of round medians, separately for active and wall", + "runner_event_times": "retained in runner_report; excluded from Nsight timing summaries", + "execution_order": [["rust", "authored-cute-dsl"] if i % 2 == 0 else ["authored-cute-dsl", "rust"] + for i in range(args.rounds)], + "provenance": {}, "commands": [], "runs": [], + } + + def save(): + temporary = directory / "report.json.tmp" + temporary.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + temporary.replace(directory / "report.json") + + def execute(command, log_path): + entry = {"command": command, "log": str(log_path), "status": "running", "started_ns": time.time_ns()} + report["commands"].append(entry) + save() + try: + entry["exit_code"] = run_command(command, args.timeout, environment, log_path) + if entry["exit_code"]: + raise ValueError(f"command exited {entry['exit_code']}; see {log_path}") + entry["status"] = "passed" + except (OSError, ValueError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + entry.update(status="failed", error=str(error) or "interrupted") + raise + finally: + entry["finished_ns"] = time.time_ns() + save() + + print(f"Evidence: {directory}", flush=True) + active = None + save() + try: + provenance = report["provenance"] + provenance["rust_binary"] = file_identity(args.rust_binary) + provenance["reference_runner"] = file_identity(args.reference_runner) + provenance["source_snapshot"] = common.source_snapshot(args.reference_runner) + provenance["rust_build"] = {"status": "unavailable; current source pins do not prove the binary's build toolchain"} + if args.rust_build_manifest: + manifest = json.loads(args.rust_build_manifest.read_text(), parse_constant=reject_nonfinite_json, + parse_float=finite_json_float) + if not isinstance(manifest, dict) or manifest.get("binary_sha256") != provenance["rust_binary"]["sha256"]: + raise ValueError("Rust build manifest must contain binary_sha256 matching --rust-binary") + provenance["rust_build"] = {"file": file_identity(args.rust_build_manifest), "manifest": manifest} + for key, command in ( + ("nsys_version", [args.nsys, "--version"]), ("python_version", [args.python, "--version"]), + ("git_head", ["git", "-C", str(HERE), "rev-parse", "HEAD"]), + ("git_status", ["git", "-C", str(HERE), "status", "--short"]), + ): + log = directory / f"{key}.log" + execute(command, log) + provenance[key] = log.read_text().strip() + report["expected_input_hashes"] = expected_input_hashes(args) + expected_output = cpu_oracle(args) + cli = ["--seq-q", "1", "--seq-k", str(args.seq_k), "--heads", str(args.heads), "--noncausal", + "--input-scale", str(args.input_scale), "--device", str(args.device), + "--specialization", args.specialization, + "--warmup", str(args.warmup), "--iters", str(args.samples)] + if args.zero_query: + cli.append("--zero-query") + implementations = {"rust": [str(args.rust_binary), "--rust-only"], "authored-cute-dsl": [args.python, str(args.reference_runner)]} + output_hashes, kernel_names, baseline = {}, {}, None + for round_index, order in enumerate(report["execution_order"], start=1): + outputs = {} + for label in order: + prefix = directory / f"round-{round_index:02d}-{label}" + result_path, output_path = prefix.with_suffix(".json"), prefix.with_suffix(".f16") + trace_path, database_path = prefix.with_suffix(".nsys-rep"), prefix.with_suffix(".sqlite") + profile = [args.nsys, "profile", "--trace=cuda,nvtx", "--sample=none", "--cpuctxsw=none", + "--output", str(prefix), *implementations[label], *cli, + "--json", str(result_path), "--output", str(output_path)] + export = [args.nsys, "export", "--type=sqlite", "--output", str(database_path), str(trace_path)] + active = {"round": round_index, "label": label, "status": "running", "started_ns": time.time_ns(), + "correctness_report": str(result_path), "output": str(output_path), + "raw_trace": str(trace_path), "sqlite": str(database_path)} + report["runs"].append(active) + save() + if any(path.exists() for path in (result_path, output_path, trace_path, database_path)): + raise ValueError("refusing to reuse existing profiling evidence") + print(f"Round {round_index}/{args.rounds}: {label}", flush=True) + execute(profile, prefix.with_suffix(".profile.log")) + checked, output = validate_report(result_path, output_path, args, active["started_ns"], + report["expected_input_hashes"], expected_output) + require_fields(checked, {"implementation": "cute-rs" if label == "rust" else "cutlass-cute-dsl"}, + "report") + active["runner_report"] = checked + active["independent_cpu_validation"] = fixture.verify(output, expected_output) + outputs[label] = output + gpu = {key: checked["provenance"][key] for key in ("gpu", "compute_capability", "device")} + if baseline is None: + baseline = gpu + provenance["gpu"] = gpu + elif gpu != baseline: + raise ValueError("GPU identity changed between runners/rounds") + if label in output_hashes and output_hashes[label] != checked["output_hash"]: + raise ValueError(f"{label} output changed across rounds") + output_hashes[label] = checked["output_hash"] + if not trace_path.is_file(): + raise ValueError(f"missing raw Nsight Systems trace: {trace_path}") + execute(export, prefix.with_suffix(".export.log")) + names = args.rust_kernel_names if label == "rust" else args.reference_kernel_names + trace_data = sequence_samples(database_path, args.warmup, args.samples, expected_names=names, + name_contains=None if names else KERNEL_SEQUENCE) + if label in kernel_names and kernel_names[label] != trace_data["kernel_names"]: + raise ValueError(f"{label} kernel names changed across rounds") + kernel_names[label] = trace_data["kernel_names"] + if "trace_device_id" not in provenance: + provenance["trace_device_id"] = trace_data["trace_device_id"] + elif trace_data["trace_device_id"] != provenance["trace_device_id"]: + raise ValueError("Nsight kernel device changed between runners/rounds") + active.update(trace_data, status="passed", finished_ns=time.time_ns()) + save() + report.setdefault("cross_checks", []).append({"round": round_index, "status": "passed", + **fixture.verify(outputs["rust"], outputs["authored-cute-dsl"])}) + save() + report.update(status="passed", summary=aggregate_rounds(report["runs"])) + except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + report.update(status="failed", error=str(error) or "interrupted") + if active is not None and active["status"] == "running": + active.update(status="failed", error=report["error"], finished_ns=time.time_ns()) + print(f"FAIL: {report['error']}", file=sys.stderr) + finally: + report.update(finished_ns=time.time_ns(), finished_utc=datetime.now(timezone.utc).isoformat()) + save() + if report["status"] == "passed": + for label, metrics in report["summary"].items(): + print(f"{label}: " + ", ".join(f"{name} {value['median_of_round_medians_us']:.3f} us" + for name, value in metrics.items()) + " (median of round medians)") + print(f"Comparison report: {directory / 'report.json'}") + return 0 if report["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/launch.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/launch.py new file mode 100644 index 0000000000..6d0567f082 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/launch.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Check and time the independently authored CuTeDSL split-KV decode counterpart.""" + +import argparse +import hashlib +import importlib.util +import json +import math +from pathlib import Path +import statistics + +_shared_path = Path(__file__).resolve().parents[2] / "fmha_prefill_sm120a_cute/reference/launch.py" +_spec = importlib.util.spec_from_file_location("fmha_decode_fixture", _shared_path) +_shared = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_shared) + +# Keep fixture and oracle shared; the decode launch sequence is independent. +make_inputs = _shared.make_inputs +reference_cpu = _shared.reference_cpu +verify = _shared.verify +decode_f16 = _shared.decode_f16 +encode_f16 = _shared.encode_f16 +bits_hash = _shared.bits_hash +HEAD_DIM = 64 +KV_CHUNK = 64 +PARTIAL_SHARED_BYTES = 17696 +STATIC_SHAPES = ((64, 8), (65, 3), (1024, 8), (4097, 8)) + + +def split_count(seq_k): + if seq_k <= 0: + raise ValueError("KV length must be positive") + return (seq_k + KV_CHUNK - 1) // KV_CHUNK + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--seq-q", type=int, default=1) + parser.add_argument("--seq-k", type=int, default=1024) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--specialization", choices=("auto", "runtime", "static"), default="auto", + help="auto specializes the same fixed shape portfolio as Rust") + masks = parser.add_mutually_exclusive_group() + masks.add_argument("--causal", action="store_true", dest="causal") + masks.add_argument("--noncausal", action="store_false", dest="causal") + parser.set_defaults(causal=False) + parser.add_argument("--input-scale", type=int, choices=(1, 4, 16), default=1) + parser.add_argument("--zero-query", action="store_true") + parser.add_argument("--device", type=int, default=0, + help="ordinal after CUDA_VISIBLE_DEVICES remapping") + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--iters", type=int, default=31, + help="zero selects verification without timing") + parser.add_argument("--output", type=Path, help="write checked FP16 BSHD output") + parser.add_argument("--json", type=Path) + args = parser.parse_args(argv) + if args.seq_q != 1 or args.causal: + parser.error("decode requires --seq-q 1 and --noncausal") + if args.seq_k <= 0 or args.heads <= 0: + parser.error("KV length and heads must be positive") + if args.seq_k * args.heads * HEAD_DIM > 0xFFFFFFFF: + parser.error("tensor element count must fit u32 indexing") + if args.device < 0 or args.warmup < 0 or args.iters < 0: + parser.error("device, warmup and iters must be nonnegative") + static_available = (args.seq_k, args.heads) in STATIC_SHAPES + if args.specialization == "static" and not static_available: + parser.error("static specialization requires (seq-k, heads) in " + str(STATIC_SHAPES)) + args.shape_specialization = ( + "static" if static_available and args.specialization != "runtime" else "runtime" + ) + args.workload = "fmha-decode" + return args + + +def run(args): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + import torch + + from attention_decode import AttentionDecode + + q_values, k_values, v_values = make_inputs( + 1, args.seq_k, args.heads, args.input_scale, args.zero_query, + ) + expected = reference_cpu(q_values, k_values, v_values, 1, args.seq_k, + args.heads, False, use_torch=True) + torch.cuda.set_device(args.device) + device = torch.device("cuda", args.device) + tensors = [torch.tensor(values, dtype=torch.float16, device=device) + .reshape(1, length, args.heads, HEAD_DIM) + for values, length in ((q_values, 1), (k_values, args.seq_k), (v_values, args.seq_k))] + splits = split_count(args.seq_k) + partial = torch.full((args.heads, splits, 66), float("nan"), dtype=torch.float32, device=device) + output = torch.full((1, 1, args.heads, HEAD_DIM), float("nan"), dtype=torch.float16, device=device) + tensors.extend((partial, output)) + # Flat dynamic views prevent accidental shape specialization. The selected + # kernel mode reconstructs BSHD/state views from constants or runtime values. + views = [from_dlpack(tensor.reshape(-1), assumed_align=16).mark_layout_dynamic(leading_dim=0) + for tensor in tensors] + dimensions = (cutlass.Int32(args.seq_k), cutlass.Int32(args.heads), cutlass.Int32(splits)) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + static_shape = (args.seq_k, args.heads) if args.shape_specialization == "static" else (0, 0) + compiled = cute.compile(AttentionDecode(*static_shape), *views, *dimensions, stream) + + def invoke(): + compiled(*views, *dimensions, stream) + + def read_output(): + return output.detach().cpu().reshape(-1).tolist() + + invoke() + checked = read_output() + verification = verify(checked, expected) + # Poison both final and intermediate storage to expose stale split state. + output.fill_(float("nan")) + partial.fill_(float("nan")) + invoke() + replay = read_output() + verify(replay, expected) + if encode_f16(replay) != encode_f16(checked): + raise ValueError("output changed across identical launches") + for _ in range(args.warmup if args.iters else 0): + invoke() + torch.cuda.synchronize() + events = [(torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(args.iters)] + samples = [] + for start, end in events: + start.record() + invoke() + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) * 1000.0 + if not math.isfinite(elapsed) or elapsed <= 0: + raise ValueError(f"invalid event timing: {elapsed}") + samples.append(elapsed) + after = read_output() + verify(after, expected) + if encode_f16(after) != encode_f16(checked): + raise ValueError("output changed during benchmark") + if args.output: + args.output.write_bytes(encode_f16(after)) + kernel_source = Path(__file__).with_name("attention_decode.py") + result = { + "implementation": "cutlass-cute-dsl", "workload": args.workload, + "shape": [1, 1, args.seq_k, args.heads, HEAD_DIM], + "causal": False, "input_scale": args.input_scale, "zero_query": args.zero_query, + "input_hashes": {name: bits_hash(values) for name, values in zip("qkv", (q_values, k_values, v_values))}, + "output_hash": bits_hash(after), "verification": "cpu-pass", "validation": verification, + "replay_bitwise_equal": True, "metric_value": statistics.median(samples) if samples else None, + "metric_unit": "us", "samples_us": samples, + "provenance": { + "gpu": torch.cuda.get_device_name(args.device), + "compute_capability": list(torch.cuda.get_device_capability(args.device)), + "device": args.device, "tool": "cuda-events-kernel-sequence", "timing_mode": "direct", + "warmup": args.warmup if args.iters else 0, "iters": args.iters, + "algorithm": "single-query-split-kv", "kv_chunk": KV_CHUNK, "split_count": splits, + "kernel_count": 2, "kernel_sequence": ["fmha_decode_partial", "fmha_decode_merge"], + "grids": [[args.heads * splits, 1, 1], [args.heads, 1, 1]], + "threads": [128, 32], "shared_memory_bytes": [PARTIAL_SHARED_BYTES, 0], "dtype": "float16", + "cutlass_version": cutlass.__version__, "torch_version": torch.__version__, + "source_origin": "independently-authored-cuda-oxide-example", + "shape_specialization": args.shape_specialization, + "specialization_request": args.specialization, + "source_file": kernel_source.name, + "source_sha256": hashlib.sha256(kernel_source.read_bytes()).hexdigest(), + }, + } + text = json.dumps(result, indent=2, allow_nan=False) + "\n" + if args.json: + args.json.write_text(text) + print(text, end="") + return result + + +if __name__ == "__main__": + run(parse_args()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/requirements.txt b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/requirements.txt new file mode 100644 index 0000000000..b2cb40f71c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/requirements.txt @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# This comparison uses the shared environment pinned in cute-rs-preview/requirements.txt. +-r ../../../requirements.txt diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_compare.py new file mode 100644 index 0000000000..9e9e7b8fe3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_compare.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU acceptance tests for ordered split-KV decode profiling evidence.""" + +import contextlib +import importlib.util +import io +import json +import os +from pathlib import Path +import sqlite3 +import tempfile +import time +from types import SimpleNamespace +import unittest +from unittest import mock + +HERE = Path(__file__).resolve().parent + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +fixture = load("decode_comparison_fixture", HERE.parent.parent / "fmha_prefill_sm120a_cute/reference/launch.py") +compare = load("decode_comparison", HERE / "compare.py") if (HERE / "compare.py").exists() else None +PARTIAL, MERGE = "fmha_decode_partial", "fmha_decode_merge" + + +def trace(path, rows, *, full_names=None): + """Rows are (name, start, end, device, stream), deliberately inserted backwards.""" + with sqlite3.connect(path) as db: + db.execute("CREATE TABLE StringIds (id INTEGER PRIMARY KEY, value TEXT)") + db.execute("CREATE TABLE CUPTI_ACTIVITY_KIND_KERNEL " + "(start INTEGER, end INTEGER, shortName INTEGER, deviceId INTEGER, streamId INTEGER)") + names = {name: i for i, name in enumerate(dict.fromkeys(row[0] for row in rows))} + db.executemany("INSERT INTO StringIds VALUES (?, ?)", [(i, name) for name, i in names.items()]) + db.executemany("INSERT INTO CUPTI_ACTIVITY_KIND_KERNEL VALUES (?, ?, ?, ?, ?)", + [(start, end, names[name], device, stream) + for name, start, end, device, stream in reversed(rows)]) + if full_names: + db.execute("ALTER TABLE CUPTI_ACTIVITY_KIND_KERNEL ADD COLUMN demangledName INTEGER") + full = {name: len(names) + i for i, name in enumerate(dict.fromkeys(full_names))} + db.executemany("INSERT INTO StringIds VALUES (?, ?)", [(i, name) for name, i in full.items()]) + db.executemany("UPDATE CUPTI_ACTIVITY_KIND_KERNEL SET demangledName=? WHERE start=?", + [(full[name], row[1]) for name, row in zip(full_names, rows)]) + + +def sequences(): + return [row for index, duration in enumerate((900, 800, 700, 600, 31, 11, 21)) + for row in ((PARTIAL, index * 10000000, index * 10000000 + duration * 1000, 0, 7), + (MERGE, index * 10000000 + duration * 1000 + 4000, + index * 10000000 + duration * 1000 + 9000, 0, 7))] + + +def options(): + return SimpleNamespace(seq_q=1, seq_k=65, heads=1, causal=False, input_scale=1, + zero_query=False, warmup=2, samples=3, device=0, specialization="auto") + + +class ComparisonTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(compare, "dedicated decode comparison is not implemented") + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + self.args = options() + self.hashes = compare.expected_input_hashes(self.args) + self.output = fixture.decode_f16(fixture.encode_f16(fixture.reference_cpu( + *fixture.make_inputs(1, 65, 1), 1, 65, 1, False))) + + def report(self): + return { + "implementation": "cute-rs", "workload": "fmha-decode", + "shape": [1, 1, 65, 1, 64], "causal": False, "input_scale": 1, + "zero_query": False, "input_hashes": self.hashes, + "output_hash": fixture.bits_hash(self.output), "verification": "cpu-pass", + "validation": {"elements": 64, "max_abs_error": 0.0, + "max_tolerance_ratio": 0.0, "atol": fixture.ATOL, "rtol": fixture.RTOL, + "oracle": "cpu-fp64-stable-softmax"}, + "replay_bitwise_equal": True, "metric_value": 999.0, "metric_unit": "us", + "samples_us": [999.0] * 3, + "provenance": {"gpu": "test SM120 device", "compute_capability": [12, 0], + "device": 0, "warmup": 2, "iters": 3, "dtype": "float16", + "tool": "cuda-events-kernel-sequence", "timing_mode": "direct", + "algorithm": "single-query-split-kv", "kv_chunk": 64, + "split_count": 2, "kernel_count": 2, + "shape_specialization": "runtime", + "kernel_sequence": [PARTIAL, MERGE]}, + } + + def write_evidence(self, report=None, values=None): + result, output = self.directory / "runner.json", self.directory / "output.f16" + result.write_text(json.dumps(self.report() if report is None else report)) + output.write_bytes(fixture.encode_f16(self.output if values is None else values)) + return result, output + + def validate(self, paths, started=0): + return compare.validate_report(*paths, self.args, started, self.hashes) + + def sample(self, path, **selectors): + return compare.sequence_samples(path, 2, 3, **(selectors or {"expected_names": [PARTIAL, MERGE]})) + + def test_strips_whole_correctness_and_warmup_sequences_and_separates_active_from_wall(self): + path = self.directory / "trace.sqlite" + trace(path, sequences() + [("fill", 950000, 951000, 0, 7)]) + result = self.sample(path) + self.assertEqual(result["active_samples_us"], [36, 16, 26]) + self.assertEqual(result["wall_samples_us"], [40, 20, 30]) + self.assertEqual(result["correctness_active_samples_us"], [905, 805]) + self.assertEqual(result["warmup_active_samples_us"], [705, 605]) + self.assertEqual(result["kernel_launches"], {PARTIAL: 7, MERGE: 7}) + self.assertEqual(result["auxiliary_kernel_counts"], {"fill": 1}) + self.assertEqual(result["trace_stream_id"], 7) + + def test_rejects_missing_extra_reordered_overlapping_or_interleaved_kernels(self): + rows = sequences() + cases = { + "missing": rows[:-1], + "extra": rows + [(MERGE, 90000000, 90001000, 0, 7)], + "reordered": [(MERGE, *rows[0][1:]), (PARTIAL, *rows[1][1:])] + rows[2:], + "overlap": [rows[0], (MERGE, 100, 1000, 0, 7)] + rows[2:], + "inside-pair": rows + [("other", 31000 + 40000000, 32000 + 40000000, 0, 7)], + "between-samples": rows + [("other", 49000000, 49001000, 0, 7)], + } + for name, launches in cases.items(): + with self.subTest(name=name): + path = self.directory / f"{name}.sqlite" + trace(path, launches) + with self.assertRaises(ValueError): + self.sample(path) + + def test_rejects_mismatched_or_missing_device_stream_and_context(self): + for column in ("deviceId", "streamId", "contextId"): + for value in (1, None): + with self.subTest(column=column, value=value): + path = self.directory / f"{column}-{value}.sqlite" + trace(path, sequences()) + with sqlite3.connect(path) as db: + if column == "contextId": + db.execute("ALTER TABLE CUPTI_ACTIVITY_KIND_KERNEL ADD COLUMN contextId INTEGER DEFAULT 0") + db.execute(f"UPDATE CUPTI_ACTIVITY_KIND_KERNEL SET {column}=? WHERE start=0", (value,)) + with self.assertRaises(ValueError): + self.sample(path) + + def test_resolves_full_demangled_names_and_rejects_ambiguous_selectors(self): + rows = sequences() + path = self.directory / "full.sqlite" + full = [f"authored_{row[0]}_specialization" for row in rows] + trace(path, [("truncated", *row[1:]) for row in rows], full_names=full) + result = self.sample(path, name_contains=[PARTIAL, MERGE]) + self.assertEqual(result["kernel_names"], full[:2]) + self.assertEqual(result["kernel_name_field"], "demangledName") + with self.assertRaises(ValueError): + self.sample(path, expected_names=["truncated", MERGE]) + + def test_rejects_trace_without_stream_metadata_or_missing_database(self): + missing = self.directory / "missing.sqlite" + with self.assertRaises(ValueError): + self.sample(missing) + self.assertFalse(missing.exists()) + path = self.directory / "no-stream.sqlite" + trace(path, sequences()) + with sqlite3.connect(path) as db: + db.execute("ALTER TABLE CUPTI_ACTIVITY_KIND_KERNEL DROP COLUMN streamId") + with self.assertRaises(ValueError): + self.sample(path) + + def test_checks_every_output_hash_and_decode_sequence_metadata(self): + report, values = self.validate(self.write_evidence(), time.time_ns() - 1000000000) + self.assertEqual(report["validation"]["elements"], 64) + self.assertEqual(values, self.output) + changes = {"workload": "single-query-mha", "replay_bitwise_equal": 1, + "validation.elements": 1, "validation.max_tolerance_ratio": 1.01, + "validation.max_abs_error": float("nan"), "validation.atol": 1.0, + "validation.oracle": "gpu", "provenance.tool": "cuda-events-single-kernel", + "provenance.kernel_sequence": [MERGE, PARTIAL], "provenance.split_count": 1, + "provenance.kernel_count": 1, "provenance.iters": 2, + "input_hashes": {"q": "0" * 16}, "shape": [True, 1, 65, 1, 64]} + for field, value in changes.items(): + with self.subTest(field=field): + changed = self.report() + keys = field.split(".") + owner = changed if len(keys) == 1 else changed[keys[0]] + owner[keys[-1]] = value + with self.assertRaises(ValueError): + self.validate(self.write_evidence(changed)) + + def test_rejects_stale_truncated_tampered_nonfinite_outputs(self): + paths = self.write_evidence() + os.utime(paths[0], ns=(1, 1)) + with self.assertRaises(ValueError): + self.validate(paths, time.time_ns()) + for values in (self.output[:-1], [1.0] + self.output[1:], [float("nan")] + self.output[1:]): + with self.subTest(values=values[:1]): + with self.assertRaises(ValueError): + self.validate(self.write_evidence(values=values)) + + def test_recomputes_full_cpu_oracle_even_when_false_runner_evidence_is_self_consistent(self): + wrong = [1.0] * 64 + report = self.report() + report["output_hash"] = fixture.bits_hash(wrong) + with self.assertRaisesRegex(ValueError, "output mismatch"): + self.validate(self.write_evidence(report, wrong)) + + def test_rejects_missing_invalid_and_inconsistent_event_sequence_samples(self): + for value in ([], [1, 2], [1, 2, -1], [1, True, 2], [1, 2, 3]): + with self.subTest(samples=value): + report = self.report() + report["samples_us"] = value + with self.assertRaises(ValueError): + self.validate(self.write_evidence(report)) + paths = self.write_evidence() + paths[0].write_text(paths[0].read_text().replace('"metric_value": 999.0', '"metric_value": 1e999')) + with self.assertRaises(ValueError): + self.validate(paths) + + def test_specialization_selection_and_report_mismatch_are_rejected(self): + cases = ((64, 8, "auto", "static"), (65, 3, "auto", "static"), + (1024, 8, "auto", "static"), (4097, 8, "auto", "static"), + (65, 1, "auto", "runtime"), (1024, 8, "runtime", "runtime"), + (65, 3, "static", "static")) + self.assertTrue(hasattr(compare, "effective_specialization"), "specialization selection is not implemented") + for length, heads, requested, expected in cases: + with self.subTest(length=length, heads=heads, requested=requested): + self.assertEqual(compare.effective_specialization(length, heads, requested), expected) + with self.assertRaises(ValueError): + compare.effective_specialization(65, 1, "static") + report = self.report() + report["provenance"]["shape_specialization"] = "static" + with self.assertRaisesRegex(ValueError, "shape_specialization"): + self.validate(self.write_evidence(report)) + + def test_cli_propagates_modes_and_rejects_unsupported_forced_static(self): + arguments = self.arguments() + with contextlib.redirect_stderr(io.StringIO()): + with self.assertRaises(SystemExit): + compare.parse_arguments(arguments + ["--specialization", "static"]) + args = compare.parse_arguments(arguments + ["--specialization", "runtime"]) + self.assertEqual(args.specialization, "runtime") + + def test_input_hashes_match_fixture_for_tails_large_logits_and_zero_query(self): + for length, heads, scale, zero in ((1, 1, 1, False), (65, 3, 16, False), (257, 8, 1, True)): + args = options() + args.seq_k, args.heads, args.input_scale, args.zero_query = length, heads, scale, zero + expected = {name: fixture.bits_hash(values) for name, values in + zip("qkv", fixture.make_inputs(1, length, heads, scale, zero))} + self.assertEqual(compare.expected_input_hashes(args), expected) + + def test_aggregate_preserves_distinct_raw_metrics_and_round_medians(self): + entries = [{"label": "rust", "active_samples_us": active, "wall_samples_us": wall} + for active, wall in (([1, 100, 1], [3, 102, 3]), ([9, 10, 11], [13, 14, 15]), + ([3, 2, 4], [4, 3, 5]))] + result = compare.aggregate_rounds(entries)["rust"] + self.assertEqual(result["active"]["round_medians_us"], [1, 10, 3]) + self.assertEqual(result["active"]["median_of_round_medians_us"], 3) + self.assertEqual(result["wall"]["median_of_round_medians_us"], 4) + self.assertEqual(result["active"]["samples_us"], [1, 100, 1, 9, 10, 11, 3, 2, 4]) + + def arguments(self): + binary = self.directory / "rust" + binary.write_text("synthetic executable") + binary.chmod(0o755) + return ["--rust-binary", str(binary), "--output-dir", str(self.directory / "evidence"), + "--seq-k", "65", "--heads", "1", "--warmup", "2", "--samples", "3", "--rounds", "3"] + + def fake_command(self, command, timeout, environment, log_path): + self.assertEqual(environment["CUDA_VISIBLE_DEVICES"], "GPU-external") + log_path.write_text("synthetic version\n") + if command[1:2] == ["profile"]: + result = Path(command[command.index("--json") + 1]) + prefix = Path(command[command.index("--output") + 1]) + output = Path(command[-1]) + report = self.report() + if "authored-cute-dsl" in prefix.name: + report["implementation"] = "cutlass-cute-dsl" + result.write_text(json.dumps(report)) + output.write_bytes(fixture.encode_f16(self.output)) + prefix.with_suffix(".nsys-rep").write_bytes(b"synthetic raw trace") + elif command[1:2] == ["export"]: + path = Path(command[command.index("--output") + 1]) + rows = sequences() + if "authored-cute-dsl" in path.name: + rows = [(f"authored_{row[0]}_specialized", *row[1:]) for row in rows] + else: + rows = [(f"{row[0]}__generic_hash", *row[1:]) for row in rows] + trace(path, rows[:-1] if getattr(self, "drop_last_launch", False) else rows) + return 0 + + def test_command_flow_alternates_complete_serial_runs_and_preserves_all_three_metrics(self): + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-external"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare.main(self.arguments()), 0) + directory, = (self.directory / "evidence").iterdir() + report = json.loads((directory / "report.json").read_text()) + self.assertEqual(report["status"], "passed") + self.assertEqual([entry["label"] for entry in report["runs"]], + ["rust", "authored-cute-dsl", "authored-cute-dsl", "rust", "rust", "authored-cute-dsl"]) + self.assertEqual(report["summary"]["rust"]["active"]["median_of_round_medians_us"], 26) + self.assertEqual(report["summary"]["rust"]["wall"]["median_of_round_medians_us"], 30) + self.assertEqual(report["summary"]["rust"]["active"]["samples_us"], [36, 16, 26] * 3) + self.assertEqual(report["runs"][0]["runner_report"]["metric_value"], 999) + self.assertEqual(report["environment"]["CUDA_VISIBLE_DEVICES"], "GPU-external") + self.assertEqual(report["parameters"]["specialization"], "auto") + for command in (entry["command"] for entry in report["commands"] if "profile" in entry["command"]): + self.assertEqual(command[command.index("--specialization") + 1], "auto") + for previous, following in zip(report["runs"], report["runs"][1:]): + self.assertLessEqual(previous["finished_ns"], following["started_ns"]) + + def test_failed_sequence_stops_before_another_runner_and_never_publishes_summary(self): + self.drop_last_launch = True + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-external"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(self.arguments()), 1) + directory, = (self.directory / "evidence").iterdir() + report = json.loads((directory / "report.json").read_text()) + self.assertEqual(report["status"], "failed") + self.assertEqual(len(report["runs"]), 1) + self.assertNotIn("summary", report) + + def test_build_manifest_rejects_wrong_binary_before_any_command(self): + manifest = self.directory / "build.json" + manifest.write_text(json.dumps({"binary_sha256": "0" * 64})) + with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(self.arguments() + ["--rust-build-manifest", str(manifest)]), 1) + directory, = (self.directory / "evidence").iterdir() + report = json.loads((directory / "report.json").read_text()) + self.assertEqual(report["commands"], []) + self.assertEqual(report["status"], "failed") + + def test_validation_sweep_covers_tails_heads_and_unequal_split_uniform_weights(self): + self.assertTrue((HERE / "validate.py").exists(), "dedicated decode validator is not implemented") + validation = load("decode_validation_test", HERE / "validate.py") + regular = {(sk, heads) for _, sk, heads, scale, zero in validation.CASES if scale == 1 and not zero} + self.assertEqual(regular, {(sk, heads) for sk in (1, 3, 31, 32, 33, 63, 64, 65, 127, 257, 1024, 4097) + for heads in (1, 3, 8)}) + self.assertTrue(any(sk > 64 and sk % 64 and zero for _, sk, _, _, zero in validation.CASES)) + self.assertTrue(any(sk > 64 and scale == 16 for _, sk, _, scale, _ in validation.CASES)) + + def test_validation_plan_pairs_every_portfolio_shape_in_both_modes(self): + validation = load("decode_validation_specialization_test", HERE / "validate.py") + self.assertTrue(hasattr(validation, "validation_plan"), "runtime/static sweep is not implemented") + plan = validation.validation_plan(validation.CASES, "auto") + self.assertEqual(len(plan), 52) + self.assertEqual([item[5] for item in plan[:40]], ["auto"] * 40) + self.assertEqual({(sk, heads, scale, mode) for _, sk, heads, scale, _, mode in plan[40:]}, + {(sk, heads, scale, mode) for sk, heads, scale in + ((64, 8, 1), (65, 3, 1), (1024, 8, 1), (4097, 8, 1), (65, 3, 16), (4097, 8, 16)) + for mode in ("runtime", "static")}) + + def test_validator_records_fresh_checked_outputs_and_stops_on_nonzero_exit(self): + validation = load("decode_validation_flow_test", HERE / "validate.py") + binary = self.directory / "rust" + binary.write_text("synthetic executable") + binary.chmod(0o755) + calls = [] + + def execute(command, timeout, environment, log_path): + calls.append(command) + log_path.write_text("synthetic run\n") + if len(calls) == 3: + return 9 + report = self.report() + report["provenance"].update(warmup=0, iters=0) + report.update(samples_us=[], metric_value=None) + if command[0] != str(binary): + report["implementation"] = "cutlass-cute-dsl" + Path(command[command.index("--json") + 1]).write_text(json.dumps(report)) + Path(command[command.index("--output") + 1]).write_bytes(fixture.encode_f16(self.output)) + return 0 + + parent = self.directory / "validation-evidence" + with mock.patch.object(validation.compare, "run_command", side_effect=execute), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + result = validation.main(["--binary", str(binary), "--reference", "--output-dir", str(parent)], + cases=(("first", 65, 1, 1, False), ("second", 65, 1, 1, False))) + self.assertEqual(result, 1) + directory, = parent.iterdir() + report = json.loads((directory / "report.json").read_text()) + self.assertEqual(report["status"], "failed") + self.assertEqual([entry["status"] for entry in report["checks"]], ["passed", "passed", "failed"]) + self.assertEqual(report["checks"][-1]["exit_code"], 9) + self.assertEqual(len(report["cross_checks"]), 1) + self.assertEqual(report["checks"][0]["independent_cpu_validation"]["elements"], 64) + + def test_validator_checks_runtime_static_outputs_directly_and_records_parity_failure(self): + validation = load("decode_validation_pair_flow", HERE / "validate.py") + binary = self.directory / "rust" + binary.write_text("synthetic executable") + binary.chmod(0o755) + cache = {} + + def execute(command, timeout, environment, log_path): + seq_k = int(command[command.index("--seq-k") + 1]) + heads = int(command[command.index("--heads") + 1]) + scale = int(command[command.index("--input-scale") + 1]) + mode = command[command.index("--specialization") + 1] + key = (seq_k, heads, scale) + if key not in cache: + inputs = fixture.make_inputs(1, seq_k, heads, scale) + expected = fixture.reference_cpu(*inputs, 1, seq_k, heads, False) + cache[key] = expected, {name: fixture.bits_hash(values) for name, values in zip("qkv", inputs)} + expected, hashes = cache[key] + values = fixture.decode_f16(fixture.encode_f16(expected)) + if fail_cross_check: + # Both runners remain within CPU tolerance, but opposite signed + # errors exceed the direct runtime/static comparison tolerance. + values[0] += 0.00019 * (1 if mode == "runtime" else -1) + elif mode == "static": + values[0] += 2 ** -13 + values = fixture.decode_f16(fixture.encode_f16(values)) + report = self.report() + report.update(shape=[1, 1, seq_k, heads, 64], input_scale=scale, input_hashes=hashes, + output_hash=fixture.bits_hash(values), validation=fixture.verify(values, expected), + samples_us=[], metric_value=None) + report["provenance"].update(warmup=0, iters=0, shape_specialization=mode, + split_count=(seq_k + 63) // 64) + if command[0] != str(binary): + report["implementation"] = "cutlass-cute-dsl" + log_path.write_text("synthetic specialization execution\n") + Path(command[command.index("--json") + 1]).write_text(json.dumps(report)) + Path(command[command.index("--output") + 1]).write_bytes(fixture.encode_f16(values)) + return 0 + + for fail_cross_check in (False, True): + with self.subTest(fail_cross_check=fail_cross_check): + parent = self.directory / f"pairs-{fail_cross_check}" + with mock.patch.object(validation.compare, "run_command", side_effect=execute), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + result = validation.main(["--binary", str(binary), "--reference", "--output-dir", str(parent)], + cases=()) + directory, = parent.iterdir() + report = json.loads((directory / "report.json").read_text()) + if fail_cross_check: + self.assertEqual(result, 1) + self.assertEqual(report["checks"][-1]["status"], "failed") + self.assertIn("output mismatch", report["error"]) + self.assertEqual(len(report["checks"]), 3) + else: + self.assertEqual(result, 0) + self.assertEqual(len(report["checks"]), 24) + self.assertEqual(len(report["specialization_cross_checks"]), 12) + self.assertTrue(all(not entry["bitwise_equal"] and entry["max_abs_error"] > 0 + for entry in report["specialization_cross_checks"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_decode.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_decode.py new file mode 100644 index 0000000000..955b8a8c53 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/test_decode.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU checks for the single-query example's shape and attention contract.""" + +from contextlib import redirect_stderr +from io import StringIO +import unittest + +import launch + + +class DecodeTests(unittest.TestCase): + def test_defaults_select_one_query_noncausal_mha(self): + args = launch.parse_args([]) + self.assertEqual((args.seq_q, args.seq_k, args.heads, args.causal), (1, 1024, 8, False)) + self.assertEqual((args.specialization, args.shape_specialization), ("auto", "static")) + + def test_supported_shapes_allow_static_and_runtime_modes(self): + for seq_k, heads in ((64, 8), (65, 3), (1024, 8), (4097, 8)): + for request, effective in (("auto", "static"), ("static", "static"), ("runtime", "runtime")): + with self.subTest(seq_k=seq_k, heads=heads, request=request): + args = launch.parse_args(["--seq-k", str(seq_k), "--heads", str(heads), + "--specialization", request]) + self.assertEqual(args.shape_specialization, effective) + + def test_auto_falls_back_and_explicit_static_rejects_unlisted_shapes(self): + for seq_k, heads in ((1, 8), (65, 8), (1024, 3), (4096, 8)): + argv = ["--seq-k", str(seq_k), "--heads", str(heads)] + with self.subTest(seq_k=seq_k, heads=heads): + self.assertEqual(launch.parse_args(argv).shape_specialization, "runtime") + with redirect_stderr(StringIO()), self.assertRaises(SystemExit): + launch.parse_args(argv + ["--specialization", "static"]) + + def test_rejects_causal_and_prefill_invocations(self): + for argv in (["--causal"], ["--seq-q", "64"], ["--seq-k", "0"], ["--heads", "0"]): + with self.subTest(argv=argv), redirect_stderr(StringIO()), self.assertRaises(SystemExit) as error: + launch.parse_args(argv) + self.assertNotEqual(error.exception.code, 0) + + def test_zero_query_averages_all_cache_rows_for_each_head(self): + q = [0.0] * 128 + k = [0.0] * 256 + v = [2.0] * 64 + [10.0] * 64 + [4.0] * 64 + [14.0] * 64 + got = launch.reference_cpu(q, k, v, 1, 2, 2, False) + self.assertEqual(got, [3.0] * 64 + [12.0] * 64) + + def test_ragged_cache_lengths_have_nonempty_last_split(self): + for length, splits in ((1, 1), (3, 1), (63, 1), (64, 1), (65, 2), (257, 5), (1024, 16)): + with self.subTest(length=length): + args = launch.parse_args(["--seq-k", str(length)]) + self.assertEqual(args.seq_k, length) + self.assertEqual(launch.split_count(length), splits) + rows = [list(range(s * 64, min((s + 1) * 64, length))) for s in range(splits)] + self.assertTrue(all(rows)) + self.assertEqual([row for split in rows for row in split], list(range(length))) + + def test_rejects_index_overflow_and_negative_measurement_arguments(self): + for argv in (["--seq-k", str(1 << 32)], ["--warmup", "-1"], ["--iters", "-1"], ["--device", "-1"]): + with self.subTest(argv=argv), redirect_stderr(StringIO()), self.assertRaises(SystemExit): + launch.parse_args(argv) + + def test_ragged_zero_query_averages_every_split(self): + # The final one-row split must contribute to the denominator and numerator. + length = 65 + q = [0.0] * 64 + k = [0.0] * (length * 64) + v = [float(row) for row in range(length) for _ in range(64)] + got = launch.reference_cpu(q, k, v, 1, length, 1, False) + self.assertEqual(got, [32.0] * 64) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/validate.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/validate.py new file mode 100644 index 0000000000..53d1d2f953 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/reference/validate.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validate full split-KV decode outputs, tails, replay, hashes and head isolation. + +Runs no timing iterations. --reference additionally checks the independently +authored cuda-oxide CuTeDSL decode example on identical inputs. Select the GPU +externally with CUDA_VISIBLE_DEVICES. Each run writes fresh evidence. +""" + +import argparse +import importlib.util +import json +import math +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import time +from types import SimpleNamespace + + +HERE = Path(__file__).resolve().parent +_spec = importlib.util.spec_from_file_location("split_kv_decode_comparison", HERE / "compare.py") +compare = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(compare) + +CASES = tuple( + (f"sk-{length}-heads-{heads}", length, heads, 1, False) + for length in (1, 3, 31, 32, 33, 63, 64, 65, 127, 257, 1024, 4097) + for heads in (1, 3, 8) +) + ( + ("unequal-splits-uniform-65", 65, 3, 1, True), + ("unequal-splits-uniform-257", 257, 8, 1, True), + ("large-logits-tail-65", 65, 3, 16, False), + ("large-logits-tail-4097", 4097, 8, 16, False), +) + + +def validation_plan(cases, specialization): + """Keep the original shape sweep, then force both paths on the AOT portfolio.""" + for _, seq_k, heads, _, _ in cases: + compare.effective_specialization(seq_k, heads, specialization) + plan = [(*case, specialization) for case in cases] + paired = [(seq_k, heads, 1) for seq_k, heads in compare.STATIC_PORTFOLIO] + paired.extend(((65, 3, 16), (4097, 8, 16))) + plan.extend((f"specialization-sk-{seq_k}-heads-{heads}-scale-{scale}-{mode}", + seq_k, heads, scale, False, mode) + for seq_k, heads, scale in paired for mode in ("runtime", "static")) + return plan + + +def main(argv=None, *, cases=CASES): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path, required=True) + parser.add_argument("--reference", action="store_true") + parser.add_argument("--reference-runner", type=Path, default=HERE / "launch.py") + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--specialization", choices=("auto", "runtime", "static"), default="auto", + help="mode for the base shape sweep; runtime/static portfolio checks always follow") + parser.add_argument("--timeout", type=float, default=300) + args = parser.parse_args(argv) + if not args.binary.is_file() or not os.access(args.binary, os.X_OK): + parser.error("--binary must be an existing executable") + if args.reference and not args.reference_runner.is_file(): + parser.error("--reference-runner must be an existing file") + if args.device < 0 or not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("device must be nonnegative; timeout must be finite and positive") + try: + plan = validation_plan(cases, args.specialization) + except ValueError as error: + parser.error(str(error)) + args.binary = args.binary.resolve() + args.reference_runner = args.reference_runner.resolve() + parent = args.output_dir.resolve() if args.output_dir else None + if parent: + parent.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="decode-validation-", dir=parent)) + environment = os.environ.copy() + summary = { + "status": "running", "started_ns": time.time_ns(), + "binary": compare.file_identity(args.binary), "checks": [], "cross_checks": [], + "specialization_cross_checks": [], "specialization": args.specialization, + "static_portfolio": [list(shape) for shape in compare.STATIC_PORTFOLIO], + "reference_origin": "independently authored cuda-oxide CuTeDSL decode example", + "environment": {key: environment.get(key) for key in ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER")}, + "verification": "every FP16 output against independent CPU FP64 stable-softmax oracle", + "launch_contract": {"correctness_sequences": 2, "warmup_sequences": 0, + "measured_sequences": 0, "kernel_sequence": compare.KERNEL_SEQUENCE}, + } + if args.reference: + summary["reference_runner"] = compare.file_identity(args.reference_runner) + report_path = directory / "report.json" + + def save(): + temporary = directory / "report.json.tmp" + temporary.write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n") + temporary.replace(report_path) + + print(f"Evidence: {directory}", flush=True) + entry, gpu_identity = None, None + oracle_cache, paired_outputs = {}, {} + save() + try: + for name, seq_k, heads, scale, zero_query, specialization in plan: + options = SimpleNamespace(seq_q=1, seq_k=seq_k, heads=heads, causal=False, + input_scale=scale, zero_query=zero_query, device=args.device, + warmup=0, samples=0, specialization=specialization) + shape_key = (seq_k, heads, scale, zero_query) + if shape_key not in oracle_cache: + oracle_cache[shape_key] = (compare.expected_input_hashes(options), compare.cpu_oracle(options)) + input_hashes, expected = oracle_cache[shape_key] + cli = ["--seq-q", "1", "--seq-k", str(seq_k), "--heads", str(heads), "--noncausal", + "--input-scale", str(scale), "--device", str(args.device), "--warmup", "0", "--iters", "0", + "--specialization", specialization] + if zero_query: + cli.append("--zero-query") + implementations = [("rust", [str(args.binary), "--rust-only"])] + if args.reference: + implementations.append(("authored-cute-dsl", [args.python, str(args.reference_runner)])) + outputs = {} + for label, prefix in implementations: + result_path = directory / f"{name}-{label}.json" + output_path = directory / f"{name}-{label}.f16" + command = [*prefix, *cli, "--json", str(result_path), "--output", str(output_path)] + entry = {"name": name, "implementation": label, "command": command, "status": "running", + "started_ns": time.time_ns(), "runner_report": str(result_path), + "output": str(output_path), "expected_input_hashes": input_hashes, + "specialization": specialization, + "shape_specialization": compare.effective_specialization(seq_k, heads, specialization)} + summary["checks"].append(entry) + save() + if result_path.exists() or output_path.exists(): + raise ValueError("refusing to reuse existing validation evidence") + entry["exit_code"] = compare.run_command(command, args.timeout, environment, + directory / f"{name}-{label}.log") + if entry["exit_code"]: + raise ValueError(f"{name}/{label} exited {entry['exit_code']}; see its log") + checked, output = compare.validate_report(result_path, output_path, options, entry["started_ns"], + input_hashes, expected) + compare.require_fields(checked, { + "implementation": "cute-rs" if label == "rust" else "cutlass-cute-dsl"}, "report") + identity = {key: checked["provenance"][key] for key in ("gpu", "compute_capability", "device")} + if gpu_identity is None: + gpu_identity = identity + summary["gpu"] = identity + elif identity != gpu_identity: + raise ValueError("GPU identity changed between validation runners/cases") + outputs[label] = output + entry.update(checked_report=checked, independent_cpu_validation=compare.fixture.verify(output, expected)) + if name.startswith("specialization-"): + key = (*shape_key, label) + if specialization == "runtime": + paired_outputs[key] = output + else: + runtime = paired_outputs.pop(key) + summary["specialization_cross_checks"].append({ + "shape": [1, 1, seq_k, heads, 64], "input_scale": scale, + "zero_query": zero_query, "implementation": label, "status": "passed", + "runtime_output_hash": compare.fixture.bits_hash(runtime), + "static_output_hash": compare.fixture.bits_hash(output), + "bitwise_equal": compare.fixture.encode_f16(runtime) == compare.fixture.encode_f16(output), + **compare.fixture.verify(output, runtime), + }) + entry.update(status="passed", finished_ns=time.time_ns()) + save() + if args.reference: + summary["cross_checks"].append({"name": name, "status": "passed", + **compare.fixture.verify(outputs["rust"], outputs["authored-cute-dsl"])}) + print(f"PASS {name}", flush=True) + save() + summary["status"] = "passed" + except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + summary.update(status="failed", error=str(error) or "interrupted") + if entry and entry["status"] == "running": + entry.update(status="failed", error=summary["error"], finished_ns=time.time_ns()) + print(f"FAIL: {summary['error']}", file=sys.stderr) + finally: + summary["finished_ns"] = time.time_ns() + save() + print(f"Validation report: {report_path}") + return 0 if summary["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/config.rs b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/config.rs new file mode 100644 index 0000000000..c3ad013501 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/config.rs @@ -0,0 +1,557 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Decode launch geometry, CLI options, and static-shape selection. + +use crate::fixture::{HEAD_DIM, Shape}; +use std::path::PathBuf; + +/// Number of key/value rows handled by one partial CTA. +pub const KV_CHUNK: usize = 64; +/// FP32 elements per split: a maximum, a sum, and 64 numerator values. +pub const SCRATCH_STRIDE: usize = 66; +/// Threads in each partial-attention CTA. +pub const PARTIAL_THREADS: u32 = 128; +/// Shared bytes reserved for staged K/V, scores, and reduction state. +pub const PARTIAL_SHARED_BYTES: u32 = 17_696; +/// One warp combines a head's partial results. +pub const MERGE_THREADS: u32 = 32; + +/// Policy for choosing a runtime or precompiled shape. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Specialization { + /// Use a precompiled shape when available, otherwise runtime dimensions. + Auto, + /// Read the KV length and head count from kernel arguments. + Runtime, + /// Require a precompiled kernel for the exact requested shape. + Static, +} + +impl Specialization { + /// Return the CLI spelling used by the comparison driver. + pub fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Runtime => "runtime", + Self::Static => "static", + } + } +} + +/// Validated options for a decode run or CuTeDSL comparison. +#[derive(Debug)] +pub struct Args { + /// Run the Rust kernel alone instead of launching the CuTeDSL comparison. + pub rust_only: bool, + /// Optional Python interpreter for the comparison process. + pub python: Option, + /// Compare every supported mode combination for the selected workload. + pub all_cases: bool, + /// Attention dimensions and mask after CLI validation. + pub shape: Shape, + /// Requested runtime/static kernel selection policy. + pub specialization: Specialization, + /// Power-of-two multiplier applied to Q and K, leaving V unchanged. + pub input_scale: u32, + /// Use an all-zero Q tensor to exercise uniform softmax. + pub zero_query: bool, + /// CUDA device ordinal after visible-device remapping. + pub device: usize, + /// Untimed invocations before collecting samples. + pub warmup: usize, + /// Number of timed samples; zero requests correctness checks only. + pub iters: usize, + /// Optional path for raw little-endian FP16 output in standalone mode. + pub output: Option, + /// Optional path for the standalone or comparison JSON report. + pub json: Option, +} + +/// Checked grid and scratch sizes for the two decode launches. +#[derive(Debug)] +pub struct LaunchPlan { + /// Number of 64-key splits per head, including a short final split. + pub splits: u32, + /// Total partial CTAs across all heads and splits. + pub partial_blocks: u32, + /// FP32 elements needed for all partial states. + pub scratch_elements: usize, +} + +impl LaunchPlan { + /// Validate the decode shape and compute sizes without overflowing kernel indices. + pub fn new(shape: Shape) -> Result { + if shape.seq_q != 1 || shape.causal { + return Err("the decode example requires --seq-q 1 and --noncausal".into()); + } + if shape.seq_k == 0 || shape.heads == 0 { + return Err("KV length and head count must be positive".into()); + } + // The prefill fixture's validator has tile restrictions that do not + // apply to decode. Keep the common allocation/index bounds here. + for length in [shape.seq_q, shape.seq_k] { + let elements = length + .checked_mul(shape.heads) + .and_then(|n| n.checked_mul(HEAD_DIM)); + if elements.is_none_or(|n| n > u32::MAX as usize) { + return Err("tensor element count must fit u32 indexing".into()); + } + } + let splits = shape.seq_k.div_ceil(KV_CHUNK); + let partial_blocks = shape + .heads + .checked_mul(splits) + .ok_or("partial grid size overflow")?; + let scratch_elements = partial_blocks + .checked_mul(SCRATCH_STRIDE) + .filter(|&n| n <= u32::MAX as usize) + .ok_or("scratch element count must fit u32 indexing")?; + scratch_elements + .checked_mul(size_of::()) + .ok_or("scratch byte count overflow")?; + Ok(Self { + splits: u32::try_from(splits).map_err(|_| "split count must fit u32")?, + partial_blocks: u32::try_from(partial_blocks) + .map_err(|_| "partial grid must fit u32")?, + scratch_elements, + }) + } +} + +impl Args { + /// Return the exact precompiled shape or explain why it is unavailable. + fn static_kernel_shape(&self) -> Result<(u32, u32), String> { + let dimensions = (self.shape.seq_k, self.shape.heads); + if matches!(dimensions, (64, 8) | (65, 3) | (1024, 8) | (4097, 8)) { + Ok((dimensions.0 as u32, dimensions.1 as u32)) + } else { + Err(format!( + "no static specialization for Sk={}, H={}; available (Sk,H): (64,8), (65,3), (1024,8), (4097,8)", + dimensions.0, dimensions.1, + )) + } + } + + /// Select static dimensions, or `(0, 0)` for the runtime kernel. + pub fn kernel_shape(&self) -> Result<(u32, u32), String> { + match self.specialization { + Specialization::Runtime => Ok((0, 0)), + Specialization::Auto => Ok(self.static_kernel_shape().unwrap_or((0, 0))), + Specialization::Static => self.static_kernel_shape(), + } + } + + /// Parse and validate CLI options; return `None` when help is requested. + pub fn parse_from(cli: impl IntoIterator) -> Result, String> { + let mut args = Self { + rust_only: false, + python: None, + all_cases: false, + shape: Shape { + seq_q: 1, + seq_k: 1024, + heads: 8, + causal: false, + }, + specialization: Specialization::Auto, + input_scale: 1, + zero_query: false, + device: 0, + warmup: 100, + iters: 31, + output: None, + json: None, + }; + let mut cli = cli.into_iter(); + let mut mask = None; + while let Some(key) = cli.next() { + match key.as_str() { + "--help" | "-h" => return Ok(None), + "--causal" | "--noncausal" => { + let causal = key == "--causal"; + if mask.is_some_and(|previous| previous != causal) { + return Err("--causal and --noncausal are mutually exclusive".into()); + } + mask = Some(causal); + args.shape.causal = causal; + } + "--rust-only" => args.rust_only = true, + "--all-cases" => args.all_cases = true, + "--python" => { + args.python = Some(cli.next().ok_or("missing value for --python")?.into()); + } + "--zero-query" => args.zero_query = true, + "--specialization" => { + let value = cli.next().ok_or("missing value for --specialization")?; + args.specialization = match value.as_str() { + "auto" => Specialization::Auto, + "runtime" => Specialization::Runtime, + "static" => Specialization::Static, + _ => return Err("--specialization must be auto, runtime, or static".into()), + }; + } + "--seq-q" | "--seq-k" | "--heads" | "--input-scale" | "--device" | "--warmup" + | "--iters" | "--output" | "--json" => { + let value = cli + .next() + .ok_or_else(|| format!("missing value for {key}"))?; + if key == "--output" { + args.output = Some(value.into()); + continue; + } + if key == "--json" { + args.json = Some(value.into()); + continue; + } + let number = value + .parse::() + .map_err(|_| format!("invalid nonnegative integer for {key}: {value}"))?; + match key.as_str() { + "--seq-q" => args.shape.seq_q = number, + "--seq-k" => args.shape.seq_k = number, + "--heads" => args.shape.heads = number, + "--input-scale" => { + args.input_scale = u32::try_from(number) + .map_err(|_| "--input-scale must be 1, 4, or 16")? + } + "--device" => args.device = number, + "--warmup" => args.warmup = number, + "--iters" => args.iters = number, + _ => unreachable!(), + } + } + _ => return Err(format!("unknown argument {key}")), + } + } + crate::comparison::validate_mode( + args.rust_only, + args.python.as_ref(), + args.output.as_ref(), + )?; + if args.all_cases && args.rust_only { + return Err("--all-cases cannot be combined with --rust-only".into()); + } + LaunchPlan::new(args.shape)?; + if ![1, 4, 16].contains(&args.input_scale) { + return Err("--input-scale must be 1, 4, or 16".into()); + } + if args.all_cases { + args.static_kernel_shape()?; + } else { + args.kernel_shape()?; + } + Ok(Some(args)) + } +} + +/// Reject a static kernel whose dimensions differ from the workload. +pub fn validate_kernel_shape( + shape: Shape, +) -> Result<(), String> { + if (SEQ_K, HEADS) == (0, 0) { + return Ok(()); + } + if SEQ_K == 0 || HEADS == 0 || shape.seq_k != SEQ_K as usize || shape.heads != HEADS as usize { + return Err(format!( + "kernel specialization Sk={SEQ_K}, H={HEADS} does not match workload Sk={}, H={}", + shape.seq_k, shape.heads, + )); + } + Ok(()) +} + +/// Tests for decode CLI validation, specialization selection, and launch bounds. +#[cfg(test)] +mod tests { + use super::*; + + /// Build a noncausal single-query shape for launch-plan tests. + fn shape(seq_k: usize, heads: usize) -> Shape { + Shape { + seq_q: 1, + seq_k, + heads, + causal: false, + } + } + + /// Parse borrowed CLI words without reading process arguments. + fn parse(cli: &[&str]) -> Result, String> { + Args::parse_from(cli.iter().map(|value| (*value).into())) + } + + /// Check that raw output is standalone-only and Python selection is comparison-only. + #[test] + fn comparison_cli_requires_standalone_mode_for_output() { + assert!(parse(&["--python", "/tmp/python"]).is_ok()); + assert!(parse(&["--output", "o.f16"]).is_err()); + assert!(parse(&["--rust-only", "--output", "o.f16"]).is_ok()); + assert!(parse(&["--rust-only", "--python", "/tmp/python"]).is_err()); + } + + /// Check that the comparison matrix accepts every supported static shape in each mode. + #[test] + fn all_cases_accepts_every_static_shape_and_selected_specialization() { + for (seq_k, heads) in [(64, 8), (65, 3), (1024, 8), (4097, 8)] { + for mode in ["auto", "runtime", "static"] { + let args = parse(&[ + "--all-cases", + "--seq-k", + &seq_k.to_string(), + "--heads", + &heads.to_string(), + "--specialization", + mode, + ]) + .unwrap() + .unwrap(); + assert!(args.all_cases); + assert!(!args.shape.causal); + } + } + } + + /// Ensure all-cases comparisons require a matching static kernel in every mode. + #[test] + fn all_cases_rejects_shapes_without_static_kernels_even_in_runtime_mode() { + for (seq_k, heads) in [(1, 8), (64, 3), (65, 8), (1024, 3), (4097, 3)] { + for mode in ["auto", "runtime", "static"] { + let error = parse(&[ + "--all-cases", + "--seq-k", + &seq_k.to_string(), + "--heads", + &heads.to_string(), + "--specialization", + mode, + ]) + .unwrap_err(); + assert!(error.contains("static specialization"), "{error}"); + } + } + } + + /// Reject options that cannot participate in the decode comparison matrix. + #[test] + fn all_cases_rejects_standalone_execution_output_and_causal_attention() { + for cli in [ + &["--all-cases", "--rust-only"][..], + &["--rust-only", "--all-cases"], + &["--all-cases", "--output", "o.f16"], + &["--all-cases", "--causal"], + ] { + assert!(parse(cli).is_err(), "{cli:?}"); + } + } + + /// Check CTA and scratch counts at cache lengths around split boundaries. + #[test] + fn launch_plan_covers_short_caches_and_split_tails() { + for (seq_k, splits, blocks, scratch_elements) in [ + (1, 1, 3, 198), + (63, 1, 3, 198), + (64, 1, 3, 198), + (65, 2, 6, 396), + (127, 2, 6, 396), + (128, 2, 6, 396), + (129, 3, 9, 594), + (1024, 16, 48, 3168), + (4097, 65, 195, 12870), + ] { + let plan = LaunchPlan::new(shape(seq_k, 3)).unwrap(); + assert_eq!( + (plan.splits, plan.partial_blocks, plan.scratch_elements), + (splits, blocks, scratch_elements), + "seq_k={seq_k}" + ); + } + } + + /// Reject empty dimensions, multiple queries, and causal decode. + #[test] + fn launch_plan_rejects_invalid_decode_workloads() { + for invalid in [ + shape(0, 1), + shape(1, 0), + Shape { + seq_q: 0, + ..shape(64, 1) + }, + Shape { + seq_q: 64, + ..shape(64, 1) + }, + Shape { + causal: true, + ..shape(1, 1) + }, + ] { + assert!(LaunchPlan::new(invalid).is_err(), "{invalid:?}"); + } + } + + /// Check both tensor and scratch indexing limits before allocation. + #[test] + fn launch_plan_rejects_input_and_scratch_index_overflow() { + for invalid in [ + shape(usize::MAX, 1), + shape(1, usize::MAX), + shape(1, 67_108_864), // Q/K/V exceed u32 elements. + shape(1, 65_075_263), // Tensors fit; the 66-element scratch slots do not. + ] { + assert!(LaunchPlan::new(invalid).is_err(), "{invalid:?}"); + } + let plan = LaunchPlan::new(shape(1, 65_075_262)).unwrap(); + assert_eq!(plan.partial_blocks, 65_075_262); + assert_eq!(plan.scratch_elements, 4_294_967_292); + } + + /// Check defaults and the CLI path for verifying a short final split. + #[test] + fn cli_selects_decode_defaults_and_accepts_tail_verification() { + let defaults = parse(&[]).unwrap().unwrap(); + assert!(!defaults.all_cases); + assert_eq!( + ( + defaults.shape.seq_q, + defaults.shape.seq_k, + defaults.shape.heads, + defaults.shape.causal + ), + (1, 1024, 8, false) + ); + let args = parse(&[ + "--seq-q", + "1", + "--seq-k", + "65", + "--heads", + "3", + "--noncausal", + "--input-scale", + "16", + "--zero-query", + "--device", + "2", + "--warmup", + "7", + "--iters", + "0", + "--rust-only", + "--output", + "o.f16", + "--json", + "r.json", + ]) + .unwrap() + .unwrap(); + assert_eq!( + (args.shape.seq_k, args.shape.heads, args.input_scale), + (65, 3, 16) + ); + assert!(args.zero_query); + assert_eq!((args.device, args.warmup, args.iters), (2, 7, 0)); + assert_eq!(args.output, Some(PathBuf::from("o.f16"))); + assert_eq!(args.json, Some(PathBuf::from("r.json"))); + assert!(parse(&["--help"]).unwrap().is_none()); + } + + /// Reject invalid dimensions, values, masks, and unknown options. + #[test] + fn cli_rejects_unsupported_shapes_and_malformed_arguments() { + for cli in [ + &["--causal"][..], + &["--seq-q", "64"], + &["--seq-k", "0"], + &["--heads", "0"], + &["--seq-k", "-1"], + &["--iters", "-1"], + &["--input-scale", "2"], + &["--seq-k"], + &["--unknown"], + &["--splits", "2"], + &["--device", "abc"], + &["--causal", "--noncausal"], + &["--noncausal", "--causal"], + ] { + assert!(parse(cli).is_err(), "{cli:?}"); + } + } + + /// Check that every supported shape selects the requested static or runtime kernel. + #[test] + fn specialization_selects_only_exact_aot_shapes() { + assert_eq!( + parse(&[]).unwrap().unwrap().kernel_shape().unwrap(), + (1024, 8) + ); + for (seq_k, heads) in [(64, 8), (65, 3), (1024, 8), (4097, 8)] { + for mode in ["auto", "static", "runtime"] { + let args = parse(&[ + "--seq-k", + &seq_k.to_string(), + "--heads", + &heads.to_string(), + "--specialization", + mode, + ]) + .unwrap() + .unwrap(); + let expected = if mode == "runtime" { + (0, 0) + } else { + (seq_k, heads) + }; + assert_eq!(args.kernel_shape().unwrap(), expected); + } + } + } + + /// Allow runtime fallback in auto mode while rejecting unavailable static shapes. + #[test] + fn specialization_falls_back_only_when_permitted() { + for (seq_k, heads) in [(1, 8), (64, 3), (65, 8), (1024, 3), (4097, 3)] { + for mode in ["auto", "runtime"] { + let args = parse(&[ + "--seq-k", + &seq_k.to_string(), + "--heads", + &heads.to_string(), + "--specialization", + mode, + ]) + .unwrap() + .unwrap(); + assert_eq!(args.kernel_shape().unwrap(), (0, 0)); + } + assert!( + parse(&[ + "--seq-k", + &seq_k.to_string(), + "--heads", + &heads.to_string(), + "--specialization", + "static" + ]) + .is_err() + ); + } + for cli in [&["--specialization"][..], &["--specialization", "invalid"]] { + assert!(parse(cli).is_err()); + } + } + + /// Check that static dimensions are validated before a kernel launch. + #[test] + fn specialization_rejects_kernel_shape_mismatches_before_launch() { + assert!(validate_kernel_shape::<0, 0>(shape(65, 3)).is_ok()); + assert!(validate_kernel_shape::<65, 3>(shape(65, 3)).is_ok()); + assert!(validate_kernel_shape::<65, 3>(shape(64, 3)).is_err()); + assert!(validate_kernel_shape::<65, 3>(shape(65, 8)).is_err()); + assert!(validate_kernel_shape::<0, 3>(shape(65, 3)).is_err()); + assert!(validate_kernel_shape::<65, 0>(shape(65, 3)).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/kernel.rs b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/kernel.rs new file mode 100644 index 0000000000..bf3d3dbb32 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/kernel.rs @@ -0,0 +1,294 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Single-token attention: parallel KV tiles followed by a stable split reduction. +//! +//! CuTe copies stage 64 keys and values per CTA. Four lanes cooperate on each +//! dot product; no query padding or matrix-multiply query rows are allocated. +//! Scores, probabilities, split numerators and reduction states remain FP32. +use core::marker::PhantomData; +use ctm_rs::{ + async_copy, barrier, cuda_module, float, kernel, launch_bounds, launch_contract, + memory::{DisjointSlice, DynamicSharedArray}, + thread, warp, +}; +use cute_rs::markers::{C, ComposedLayout, CopyG2S, Layout, LeadingDim, RowMajor, Swizzle}; +use cute_rs::{GmemMatrix, SmemTile, copy_g2s}; + +/// Swizzled 64-by-64 FP16 layout used for both K and V in shared memory. +type KvLayout = ComposedLayout, 0, RowMajor<64, 64>>; +/// Swizzled 16-by-64 region covered by one cooperative copy. +type CopyLayout = ComposedLayout, 0, RowMajor<16, 64>>; +/// Map 128 threads to 16 rows and eight 16-byte copy positions. +type CopyThreads = Layout<(C<16>, C<8>), (C<8>, C<1>)>; +/// Eight adjacent FP16 values copied by each thread. +type CopyValues = Layout<(C<1>, C<8>), (C<8>, C<1>)>; +/// Typed shared-memory view of one 64-row K or V tile. +type KvTile = SmemTile; + +/// Attention scale converted to base two for the exponential instruction. +const SCALE_LOG2: f32 = 0.125 * core::f32::consts::LOG2_E; +/// FP32 elements per split: a maximum, a sum, and 64 numerator values. +const STATE_SIZE: usize = 66; + +/// Zero/zero selects runtime dimensions. Every static shape is checked while +/// Rust monomorphizes the kernel, before any device code can be launched. +const fn assert_shape() { + assert!((SEQ_K == 0) == (HEADS == 0)); + assert!(SEQ_K as u64 * HEADS as u64 <= u32::MAX as u64 / 64); + assert!((SEQ_K as u64).div_ceil(64) * HEADS as u64 <= u32::MAX as u64 / 66); +} + +/// Expand fixed iterations with constant indices so fragments stay in registers. +macro_rules! static_for { + ($index:ident in [$($value:expr),*] $body:block) => { + $({ #[doc = "Compile-time index for this unrolled iteration."] const $index: usize = $value; $body })* + }; +} + +/// Select the larger value; comparisons involving NaN fall back to `b`. +#[inline(always)] +fn max(a: f32, b: f32) -> f32 { + if a > b { a } else { b } +} + +/// Reduce a maximum across all 32 participating lanes. +#[inline(always)] +fn warp_max(mut value: f32) -> f32 { + static_for!(MASK in [16, 8, 4, 2, 1] { + value = max(value, warp::shuffle_xor_f32(value, MASK as u32)); + }); + value +} + +/// Sum values across all 32 participating lanes. +#[inline(always)] +fn warp_sum(mut value: f32) -> f32 { + static_for!(MASK in [16, 8, 4, 2, 1] { + value += warp::shuffle_xor_f32(value, MASK as u32); + }); + value +} + +/// Map a logical K/V coordinate through the tile's shared-memory swizzle. +#[inline(always)] +fn shared_offset(row: usize, col: usize) -> usize { + let linear = row * 64 + col; + linear ^ ((linear & (7 << 6)) >> 3) +} + +/// Wrap storage for 4096 FP16 values as a typed K/V tile. +/// +/// # Safety +/// `base` must point to live shared storage with the tile's required alignment. +#[inline(always)] +unsafe fn tile(base: *mut f16) -> KvTile { + SmemTile { + base, + capacity: 4096, + layout: PhantomData, + } +} + +/// Issue cooperative copies for one KV split, zero-filling rows past the sequence. +/// +/// # Safety +/// All 128 threads must participate with valid source and destination views; wait before reading the tile. +#[inline(always)] +unsafe fn stage(src: &GmemMatrix>, split: usize, dst: &KvTile, tid: u32) { + static_for!(GROUP in [0, 1, 2, 3] { + let mut part = SmemTile:: { + base: unsafe { dst.base.add(GROUP * 1024) }, capacity: 1024, layout: PhantomData, + }; + // CuTe's copy predicate zero-fills rows beyond the actual KV length. + unsafe { + copy_g2s::, CopyThreads, CopyValues, RowMajor<16, 64>, CopyLayout, f16>( + src, (split * 4 + GROUP, 0), &mut part, tid, + ); + } + }); +} + +/// Device entry points for partial attention and the split reduction. +#[cuda_module] +pub mod kernels { + use super::*; + + /// Compute one head's 64-key split and write its FP32 numerator, maximum, and sum. + #[launch_bounds(128)] + #[launch_contract(domain=1, block=(128,1,1), dynamic_shared=17696, + dynamic_shared_alignment=1024, min_compute_capability=(12,0))] + #[kernel] + pub fn fmha_decode_partial( + q: &[f16], + k: &[f16], + v: &[f16], + mut partials: DisjointSlice, + seq_k: u32, + heads: u32, + ) { + const { assert_shape::() }; + let seq_k = if SEQ_K == 0 { seq_k } else { SEQ_K }; + let heads = if HEADS == 0 { heads } else { HEADS }; + let tid = thread::threadIdx_x() as usize; + let lane = tid % 32; + let warp_id = tid / 32; + let splits = (seq_k as usize).div_ceil(64); + let split = thread::blockIdx_x() as usize % splits; + let head = thread::blockIdx_x() as usize / splits; + unsafe { + let smem = DynamicSharedArray::::get_raw().cast::(); + let sk = tile(smem); + let sv = tile(smem.add(4096)); + let scores = smem.add(8192).cast::(); + let maxima = scores.add(64); + let sums = maxima.add(4); + let warp_outputs = sums.add(4); + let global_k = GmemMatrix { + base: k.as_ptr().add(head * 64), + rows: seq_k as usize, + cols: 64, + leading_dim: LeadingDim::<8> { + elements: heads as usize * 64, + }, + }; + let global_v = GmemMatrix { + base: v.as_ptr().add(head * 64), + rows: seq_k as usize, + cols: 64, + leading_dim: LeadingDim::<8> { + elements: heads as usize * 64, + }, + }; + stage(&global_k, split, &sk, tid as u32); + stage(&global_v, split, &sv, tid as u32); + async_copy::cp_async_commit_group(); + async_copy::cp_async_wait_group::<0>(); + barrier::barrier_cta_sync(1, 128); + + // 32 independent keys at a time, with four lanes per 64-D dot. + static_for!(PASS in [0, 1] { + let row = tid / 4 + PASS * 32; + let mut dot = 0.0f32; + static_for!(D in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] { + let column = tid % 4 * 16 + D; + dot += *q.as_ptr().add(head * 64 + column) as f32 + * *sk.base.add(shared_offset(row, column)) as f32; + }); + dot += warp::shuffle_xor_f32(dot, 1); + dot += warp::shuffle_xor_f32(dot, 2); + if tid.is_multiple_of(4) { + *scores.add(row) = if split * 64 + row < seq_k as usize { + dot + } else { f32::NEG_INFINITY }; + } + }); + barrier::barrier_cta_sync(1, 128); + let score = if tid < 64 { + *scores.add(tid) + } else { + f32::NEG_INFINITY + }; + let local_maximum = warp_max(score); + if lane == 0 { + *maxima.add(warp_id) = local_maximum; + } + barrier::barrier_cta_sync(1, 128); + let maximum = max( + max(*maxima, *maxima.add(1)), + max(*maxima.add(2), *maxima.add(3)), + ); + // Every launched split has at least one valid key, hence a finite maximum. + let weight = float::ex2_approx_f32((score - maximum) * SCALE_LOG2); + if tid < 64 { + *scores.add(tid) = weight; + } + let local_sum = warp_sum(weight); + if lane == 0 { + *sums.add(warp_id) = local_sum; + } + barrier::barrier_cta_sync(1, 128); + + // Each warp integrates 16 keys into two output dimensions per lane. + let mut numerator = [0.0f32; 2]; + static_for!(ROW in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] { + let row = warp_id * 16 + ROW; + let probability = *scores.add(row); + static_for!(D in [0, 1] { + numerator[D] += probability * *sv.base.add(shared_offset(row, lane * 2 + D)) as f32; + }); + }); + static_for!(D in [0, 1] { + *warp_outputs.add(warp_id * 64 + lane * 2 + D) = numerator[D]; + }); + barrier::barrier_cta_sync(1, 128); + if warp_id == 0 { + let dst = partials + .as_mut_ptr() + .add((head * splits + split) * STATE_SIZE); + static_for!(D in [0, 1] { + let column = lane * 2 + D; + let mut sum = 0.0; + static_for!(WARP in [0, 1, 2, 3] { + sum += *warp_outputs.add(WARP * 64 + column); + }); + *dst.add(2 + column) = sum; + }); + if lane == 0 { + *dst = maximum; + *dst.add(1) = *sums + *sums.add(1) + *sums.add(2) + *sums.add(3); + } + } + } + } + + /// Stably combine all splits for one head and write its normalized FP16 output. + #[launch_bounds(32)] + #[launch_contract(domain=1, block=(32,1,1), min_compute_capability=(12,0))] + #[kernel] + pub fn fmha_decode_merge( + partials: &[f32], + mut out: DisjointSlice, + heads: u32, + splits: u32, + ) { + const { assert_shape::() }; + let heads = if HEADS == 0 { heads } else { HEADS }; + let splits = if SEQ_K == 0 { + splits + } else { + SEQ_K.div_ceil(64) + }; + let head = thread::blockIdx_x() as usize; + let lane = thread::threadIdx_x() as usize; + if head < heads as usize { + unsafe { + let base = partials.as_ptr().add(head * splits as usize * STATE_SIZE); + let mut maximum = f32::NEG_INFINITY; + let mut split = 0usize; + while split < splits as usize { + maximum = max(maximum, *base.add(split * STATE_SIZE)); + split += 1; + } + let mut denominator = 0.0f32; + let mut numerator = [0.0f32; 2]; + split = 0; + while split < splits as usize { + let state = base.add(split * STATE_SIZE); + let rescale = float::ex2_approx_f32((*state - maximum) * SCALE_LOG2); + denominator += *state.add(1) * rescale; + static_for!(D in [0, 1] { + numerator[D] += *state.add(2 + lane * 2 + D) * rescale; + }); + split += 1; + } + let reciprocal = float::rcp_approx_ftz_f32(denominator); + static_for!(D in [0, 1] { + *out.as_mut_ptr().add(head * 64 + lane * 2 + D) = (numerator[D] * reciprocal) as f16; + }); + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/main.rs new file mode 100644 index 0000000000..021a259b7a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/main.rs @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM120a single-token attention example and comparison entry point. + +#![feature(f16)] + +/// Shared CuTeDSL comparison launcher used by both SM120a attention examples. +#[path = "../../fmha_prefill_sm120a_cute/src/comparison.rs"] +mod comparison; +/// Checked decode options and launch geometry. +mod config; +#[path = "../../fmha_prefill_sm120a_cute/src/fixture.rs"] +mod fixture; +mod kernel; +mod runner; + +/// Run decode verification, timing, or the requested CuTeDSL comparison. +fn main() -> Result<(), Box> { + runner::run() +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/runner.rs b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/runner.rs new file mode 100644 index 0000000000..18f105c57b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/src/runner.rs @@ -0,0 +1,241 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Checked host launch and complete partial-plus-merge timing for decode. + +use crate::{ + config::{ + Args, KV_CHUNK, LaunchPlan, MERGE_THREADS, PARTIAL_SHARED_BYTES, PARTIAL_THREADS, + SCRATCH_STRIDE, validate_kernel_shape, + }, + fixture::{self, HEAD_DIM, bits_hash, make_inputs, reference, verify}, + kernel, +}; +use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D, sys}; +use serde_json::json; +use std::error::Error; + +/// Host-side errors from CUDA, file I/O, validation, or the comparison process. +type Result = std::result::Result>; + +/// Parse options and dispatch the comparison or selected standalone kernel shape. +pub fn run() -> Result<()> { + let Some(args) = Args::parse_from(std::env::args().skip(1))? else { + println!( + "{} [--seq-q 1] [--seq-k 1024] [--heads 8] [--noncausal]\n \ + [--input-scale 1|4|16] [--zero-query] [--device 0]\n \ + [--specialization auto|runtime|static]\n \ + [--warmup 100] [--iters 31] [--json result.json] [--python PATH] [--rust-only | --all-cases]\nDefault: verify and compare CuTeDSL and cute-rs-preview; prints two performance lines.\n--all-cases prints runtime and static comparisons at the selected noncausal shape.\nIt overrides --specialization, requires a listed static shape, and cannot be combined with --rust-only.\nLive timings use CUDA events; archived Nsight measurements are a separate metric.\n--json saves the combined report; --rust-only saves standalone JSON and permits --output o.f16.\nPython selection: --python, CUDA_OXIDE_CUTE_PYTHON, example .venv, cute-rs-preview/.venv, then python3.\n\ + B=1, Sq=1, D=64, FP16, softmax scale=0.125, noncausal only.\n\ + One 128-thread CTA per head/64-key split; one 32-thread merge CTA per head.\n\ + KV length may be any positive value. Partial state and merge arithmetic are FP32.\n\ + Auto selects a static kernel for (Sk,H)=(64,8),(65,3),(1024,8),(4097,8);\n\ + other shapes use the runtime kernel. Static mode requires a listed shape.\n\ + Every output is checked against a CPU FP64 oracle before timing.\n\ + --input-scale scales Q/K only; --zero-query exercises uniform softmax.\n\ + --iters 0 verifies two partial-plus-merge invocations without timing.\n\ + CUDA events measure the full partial-plus-merge sequence with reused buffers.\n\ + Device ordinal is after CUDA_VISIBLE_DEVICES remapping.", + env!("CARGO_PKG_NAME"), + ); + return Ok(()); + }; + if !args.rust_only { + return crate::comparison::run( + "decode", + args.python.as_ref(), + args.all_cases, + args.shape, + args.input_scale, + args.zero_query, + args.device, + args.warmup, + args.iters, + "--specialization", + args.specialization.as_str(), + args.json.as_ref(), + ); + } + match args.kernel_shape()? { + (0, 0) => run_shape::<0, 0>(args), + (64, 8) => run_shape::<64, 8>(args), + (65, 3) => run_shape::<65, 3>(args), + (1024, 8) => run_shape::<1024, 8>(args), + (4097, 8) => run_shape::<4097, 8>(args), + _ => Err("unsupported kernel specialization dispatch".into()), + } +} + +/// Verify both decode launches, time their combined work, and write the report. +fn run_shape(args: Args) -> Result<()> { + let shape = args.shape; + validate_kernel_shape::(shape)?; + let plan = LaunchPlan::new(shape)?; + let (q, k, v) = make_inputs(shape, args.input_scale, args.zero_query); + let expected = reference(shape, &q, &k, &v); + let ctx = CudaContext::new(args.device)?; + let capability = ctx.compute_capability()?; + if capability.0 < 12 { + return Err("this Rust specialization requires SM120 or newer".into()); + } + let stream = ctx.default_stream(); + // All AOT specializations live in this executable's one CUTLASS cubin. + // The generic cuda_module loader merges PTX bundles and skips cubins, so + // load the named bundle directly and bind its matching typed launch API. + let image = cuda_host::load_embedded_module(&ctx, env!("CARGO_PKG_NAME"))?; + // SAFETY: this package's bundle contains the kernels and launch contracts + // declared by kernel::kernels, including every dispatched specialization. + let module = unsafe { kernel::kernels::from_module(image) }?; + let q_dev = DeviceBuffer::from_host(&stream, &q)?; + let k_dev = DeviceBuffer::from_host(&stream, &k)?; + let v_dev = DeviceBuffer::from_host(&stream, &v)?; + let output_poison = vec![f16::NAN; shape.q_elements()]; + let scratch_poison = vec![f32::NAN; plan.scratch_elements]; + let mut output = DeviceBuffer::from_host(&stream, &output_poison)?; + let mut scratch = DeviceBuffer::from_host(&stream, &scratch_poison)?; + let partial = module.prepare_fmha_decode_partial::(LaunchConfig1D::new( + plan.partial_blocks, + PARTIAL_THREADS, + PARTIAL_SHARED_BYTES, + ))?; + let merge = module.prepare_fmha_decode_merge::(LaunchConfig1D::new( + shape.heads as u32, + MERGE_THREADS, + 0, + ))?; + let partial_registers = partial.function().num_registers()?; + let merge_registers = merge.function().num_registers()?; + let invoke = |scratch: &mut DeviceBuffer, output: &mut DeviceBuffer| -> Result<()> { + // Q/K/V have the checked contiguous BSHD extents. Each partial CTA + // owns one disjoint scratch slot; each merge CTA owns one head's output. + // The same stream orders merge reads after all partial writes. Buffers + // remain live until synchronization, with both launch contracts met. + module.fmha_decode_partial::( + &stream, + &partial, + &q_dev, + &k_dev, + &v_dev, + scratch, + shape.seq_k as u32, + shape.heads as u32, + )?; + module.fmha_decode_merge::( + &stream, + &merge, + scratch, + output, + shape.heads as u32, + plan.splits, + )?; + Ok(()) + }; + + invoke(&mut scratch, &mut output)?; + let checked = output.to_host_vec(&stream)?; + let validation = verify(&checked, &expected)?; + let checked_bits: Vec = checked.iter().map(|x| x.to_bits()).collect(); + // Poison the producer's intermediate state as well as the final output: + // otherwise missing partial writes could reuse valid data during replay. + scratch.copy_from_host(&stream, &scratch_poison)?; + output.copy_from_host(&stream, &output_poison)?; + invoke(&mut scratch, &mut output)?; + let replay = output.to_host_vec(&stream)?; + verify(&replay, &expected)?; + if !replay + .iter() + .map(|x| x.to_bits()) + .eq(checked_bits.iter().copied()) + { + return Err("output changed across identical partial-plus-merge invocations".into()); + } + + let warmup = if args.iters == 0 { 0 } else { args.warmup }; + for _ in 0..warmup { + invoke(&mut scratch, &mut output)?; + } + stream.synchronize()?; + let events = (0..args.iters) + .map(|_| { + Ok(( + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + )) + }) + .collect::>>()?; + let mut samples = Vec::with_capacity(args.iters); + for (start, end) in &events { + start.record(&stream)?; + invoke(&mut scratch, &mut output)?; + end.record(&stream)?; + end.synchronize()?; + let elapsed = f64::from(start.elapsed_ms(end)?) * 1000.0; + if !elapsed.is_finite() || elapsed <= 0.0 { + return Err(format!("invalid event timing: {elapsed}").into()); + } + samples.push(elapsed); + } + let after = output.to_host_vec(&stream)?; + verify(&after, &expected)?; + if !after + .iter() + .map(|x| x.to_bits()) + .eq(checked_bits.iter().copied()) + { + return Err("output changed during timing".into()); + } + if let Some(path) = args.output { + let bytes: Vec = after + .iter() + .flat_map(|value| value.to_bits().to_le_bytes()) + .collect(); + std::fs::write(path, bytes)?; + } + let mut sorted = samples.clone(); + sorted.sort_by(f64::total_cmp); + let median = if sorted.is_empty() { + None + } else { + Some((sorted[(sorted.len() - 1) / 2] + sorted[sorted.len() / 2]) * 0.5) + }; + let report = json!({ + "implementation": "cute-rs", "workload": "fmha-decode", + "shape": [1, shape.seq_q, shape.seq_k, shape.heads, HEAD_DIM], + "causal": shape.causal, "input_scale": args.input_scale, "zero_query": args.zero_query, + "input_hashes": {"q": format!("{:016x}", bits_hash(&q)), + "k": format!("{:016x}", bits_hash(&k)), "v": format!("{:016x}", bits_hash(&v))}, + "output_hash": format!("{:016x}", bits_hash(&after)), "verification": "cpu-pass", + "validation": {"elements": validation.elements, "max_abs_error": validation.max_abs_error, + "max_tolerance_ratio": validation.max_tolerance_ratio, "atol": fixture::ATOL, "rtol": fixture::RTOL, + "oracle": "cpu-fp64-stable-softmax"}, + "replay_bitwise_equal": true, "metric_value": median, "metric_unit": "us", "samples_us": samples, + "provenance": { + "gpu": ctx.device_name()?, "compute_capability": [capability.0, capability.1], + "device": args.device, "tool": "cuda-events-kernel-sequence", "timing_mode": "direct", + "warmup": warmup, "iters": args.iters, "dtype": "float16", "accumulator_dtype": "float32", + "algorithm": "single-query-split-kv", "kv_chunk": KV_CHUNK, "split_count": plan.splits, + "shape_specialization": if SEQ_K == 0 { "runtime" } else { "static" }, + "specialization_request": args.specialization.as_str(), + "kernel_count": 2, "kernel_sequence": ["fmha_decode_partial", "fmha_decode_merge"], + "scratch": {"shape": [shape.heads, plan.splits as usize, SCRATCH_STRIDE], + "dtype": "float32", "bytes": plan.scratch_elements * size_of::(), + "layout": "maximum, denominator, 64 unnormalized numerator components"}, + "kernels": [ + {"name": "fmha_decode_partial", "grid": [plan.partial_blocks, 1, 1], + "threads": PARTIAL_THREADS, "shared_memory_bytes": PARTIAL_SHARED_BYTES, + "registers_per_thread": partial_registers}, + {"name": "fmha_decode_merge", "grid": [shape.heads, 1, 1], + "threads": MERGE_THREADS, "shared_memory_bytes": 0, + "registers_per_thread": merge_registers} + ] + } + }); + let rendered = serde_json::to_string_pretty(&report)?; + if let Some(path) = args.json { + std::fs::write(path, format!("{rendered}\n"))?; + } + println!("{rendered}"); + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/validate.py b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/validate.py new file mode 100644 index 0000000000..3020802028 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/validate.py @@ -0,0 +1,12 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the dedicated decode validation suite, including ragged KV caches.""" + +from pathlib import Path +import runpy + + +if __name__ == "__main__": + runpy.run_path(str(Path(__file__).parent / "reference/validate.py"), run_name="__main__") diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.lock new file mode 100644 index 0000000000..7495b04e33 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "fmha_prefill_sm120a_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", + "serde_json", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.toml new file mode 100644 index 0000000000..52dfb10e64 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "fmha_prefill_sm120a_cute" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "SM120a: FP16 prefill and FMHA with warp MMA and online softmax in cute-rs" + +[workspace] + +[dependencies] +cute-rs = { path = "../../cute-rs" } +ctm-rs = { path = "../../ctm-rs" } +# CTM reexports kernel macros, whose expansions still name ::cuda_device. +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } +serde_json = "1" diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/README.md b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/README.md new file mode 100644 index 0000000000..fefed1881e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/README.md @@ -0,0 +1,66 @@ +# FP16 prefill with cute-rs — SM120a (RTX 5090) + +This cuda-oxide example computes `softmax(Q @ K.T / sqrt(64)) @ V` with +FP16 inputs/output, FP32 accumulation, and 64-by-64 tiles. + +```text +Q/K/V tensor views → typed shared-memory copies → warp MMA fragments + → online softmax → weighted V accumulation → FP16 output +``` + +- Batch **1**, head dimension **64**, contiguous BSHD tensors and equal Q/K/V heads. +- Sequence lengths are positive multiples of 64. **Causal** (default) requires + equal query/KV lengths; `--noncausal` permits unequal lengths and `Sq=1`. +- Runtime shapes are the default. `--shape-mode static` requires **Sq=Sk=1024, H=8**. + +From `cuda-oxide/`, once per checkout: + +```bash +cargo oxide toolchain install cutlass +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Compare all four cases on RTX 5090: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 \ +CUTE_DSL_ARCH=sm_120a CUTE_DSL_COMPILER_OPT=--opt-level=2 \ + cargo oxide run fmha_prefill_sm120a_cute --arch sm_120a -- \ + --all-cases --warmup 100 --iters 31 +``` + +Both implementations check every output against a CPU FP64 oracle. The command +also verifies matching input hashes and output parity, then prints a table of +CUDA-event medians. `--all-cases` requires **Sq=Sk=1024, H=8** and overrides +the shape mode and mask. + +For one case, replace `--all-cases` with `--shape-mode static --causal`; choose +`runtime` or `--noncausal` for the other cases. Add `--json PATH` to save the +comparison, `--iters 0` to verify without timing, or `--rust-only` for a standalone +Rust report. `--all-cases` cannot be combined with `--rust-only`. +Pass `--help` after `--` for all options. + +Use `--python PATH` to select another interpreter and `--device N` to select +another visible GPU; device 0 is the default. + +RTX 5090 results at `a0d5776b2` (2026-10-18), **B=1, Sq=Sk=1024, H=8, D=64**, +CUTLASS 4.7 at optimization level 2 on both sides. Each value is the median of +**31 single-invocation samples after 100 warmups**, with reused buffers. +Both implementations passed correctness checks. Times are microseconds: + +| Shapes | Attention | CUDA events: Rust (cute-rs) | CUDA events: CuTeDSL | Nsight kernel: Rust (cute-rs) | Nsight kernel: CuTeDSL | +| --- | --- | ---: | ---: | ---: | ---: | +| Runtime | Causal | 18.048 | 23.520 | 16.320 | 16.287 | +| Runtime | Noncausal | 17.952 | 23.872 | 16.416 | 16.576 | +| Static | Causal | 17.984 | 22.752 | 15.904 | 15.936 | +| Static | Noncausal | 17.984 | 23.936 | 15.936 | 16.991 | + +CUDA events include GPU idle time during host submission. Nsight Systems +measures the kernel's execution duration, including device stalls and excluding +submission gaps. These are separate runs; subtracting columns does not isolate +host overhead. Clocks were unlocked and CuTeDSL ran first in each case. +This single run per method does not establish a stable speed ranking. + +- [Reference setup and validation commands](reference/README.md) +- [Dedicated decode example](../fmha_decode_sm120a_cute/README.md) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/compare.py new file mode 100644 index 0000000000..f5990ddb3f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/compare.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Verify and compare FMHA CuTeDSL and cute-rs using direct CUDA events. + +Shared by the prefill and decode executables. Workload flags use the existing +reference/launch.py interface. Nsight profiling remains in reference/compare.py. +""" + +import argparse +import importlib.util +import json +import math +from pathlib import Path +import statistics +import subprocess +import sys +import tempfile +import time + + +HERE = Path(__file__).resolve().parent + + +SETUP_HINT = ("from cuda-oxide/ run `python3 -m venv cute-rs-preview/.venv && " + "cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt`, " + "then rerun; or set CUDA_OXIDE_CUTE_PYTHON to an environment that has them") + + +def require_reference_environment(): + """Stop before any GPU work when the shared comparison environment is missing.""" + missing = [name for name in ("numpy", "torch", "cutlass", "cuda") + if importlib.util.find_spec(name) is None] + if missing: + sys.exit(f"missing Python packages: {', '.join(missing)}; {SETUP_HINT}") + + +def load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def load_reference(workload): + directory = HERE if workload == "prefill" else HERE.parent / "fmha_decode_sm120a_cute" + launch = load_module("fmha_direct_launch", directory / "reference/launch.py") + # The prefill evidence helpers import launch by name. Keep that import + # local so callers/tests cannot accidentally supply the decode fixture. + previous = sys.modules.get("launch") + try: + sys.modules["launch"] = launch + checks = load_module("fmha_direct_checks", directory / "reference/compare.py") + finally: + if previous is None: + sys.modules.pop("launch", None) + else: + sys.modules["launch"] = previous + return directory / "reference/launch.py", launch, checks + + +def validate_timing(report, count): + samples = report.get("samples_us") + if (report.get("metric_unit") != "us" or not isinstance(samples, list) + or len(samples) != count + or any(type(value) not in (int, float) or not math.isfinite(value) or value <= 0 + for value in samples)): + raise ValueError("expected one positive finite CUDA-event sample per iteration in us") + metric = report.get("metric_value") + if samples: + if type(metric) not in (int, float) or metric != statistics.median(samples): + raise ValueError("CUDA-event timing must equal the sample median") + elif metric is not None: + raise ValueError("verification-only runs must not report a timing") + + +def runner_arguments(args, workload): + common = ["--seq-q", str(args.seq_q), "--seq-k", str(args.seq_k), "--heads", str(args.heads), + "--causal" if args.causal else "--noncausal", "--input-scale", str(args.input_scale), + "--device", str(args.device), "--warmup", str(args.warmup), "--iters", str(args.iters)] + if workload == "prefill": + common += ["--shape-mode", args.shape_mode] + else: + common += ["--specialization", args.specialization] + if args.zero_query: + common.append("--zero-query") + return common + + +def compare(options, runner_cli): + reference_path, launch, checks = load_reference(options.workload) + args = launch.parse_args(runner_cli) + if args.output is not None: + raise ValueError("--output requires --rust-only; comparison checks both outputs internally") + binary = options.rust_binary.resolve(strict=True) + args.samples = args.iters + args.warmup = args.warmup if args.iters else 0 + common = runner_arguments(args, options.workload) + hashes = checks.expected_input_hashes(args) + reports, outputs = {}, {} + # Preserve CUDA_VISIBLE_DEVICES and the reference compiler environment. + # Each process creates its own CUDA context and times only its own work. + with tempfile.TemporaryDirectory(prefix=f"fmha-{options.workload}-compare-") as temporary: + directory = Path(temporary) + for key, label, prefix, implementation in ( + ("python_reference", "CuTeDSL", [sys.executable, str(reference_path)], "cutlass-cute-dsl"), + ("cuda_oxide", "cute-rs", [str(binary), "--rust-only"], "cute-rs"), + ): + result_path, output_path = directory / f"{key}.json", directory / f"{key}.f16" + command = [*prefix, *common, "--json", str(result_path), "--output", str(output_path)] + started = time.time_ns() + completed = subprocess.run(command, capture_output=True, text=True, timeout=options.timeout) + if completed.stderr: + print(completed.stderr, file=sys.stderr, end="") + if completed.returncode: + if completed.stdout: + print(completed.stdout, file=sys.stderr, end="") + raise RuntimeError( + f"{label} runner exited {completed.returncode}; select the reference environment " + "with --python or CUDA_OXIDE_CUTE_PYTHON; use --rust-only for standalone execution" + ) + report, output = checks.validate_report(result_path, output_path, args, started, hashes) + checks.require_fields(report, {"implementation": implementation, "workload": args.workload}, "report") + validate_timing(report, args.iters) + reports[key], outputs[key] = report, output + checks.require_fields(reports["cuda_oxide"]["provenance"], { + key: reports["python_reference"]["provenance"][key] + for key in ("gpu", "compute_capability", "device") + }, "matched device") + parity = launch.verify(outputs["cuda_oxide"], outputs["python_reference"]) + return { + "workload": args.workload, "verification": "both-pass", **reports, + "output_comparison": parity, + "timing": { + "mode": "direct", "statistic": "median", "unit": "us", + "scope": "partial-plus-merge sequence" if options.workload == "decode" else "single kernel", + "order": ["python_reference", "cuda_oxide"], + "note": "CUDA-event intervals include host submission gaps; these are not Nsight kernel-active times.", + }, + } + + +def compare_all_cases(options, runner_cli): + _, launch, _ = load_reference(options.workload) + args = launch.parse_args(runner_cli) + if options.workload == "prefill": + if (args.seq_q, args.seq_k, args.heads) != launch.STATIC_SHAPE: + raise ValueError("--all-cases requires --seq-q 1024 --seq-k 1024 --heads 8 for static cases") + masks = (True, False) + else: + if (args.seq_k, args.heads) not in launch.STATIC_SHAPES: + raise ValueError("--all-cases requires (seq-k, heads) in " + str(launch.STATIC_SHAPES)) + masks = (False,) + if args.output is not None: + raise ValueError("--output requires --rust-only; comparison checks both outputs internally") + cases = [] + for shape_mode in ("runtime", "static"): + if options.workload == "prefill": + args.shape_mode = shape_mode + else: + args.specialization = shape_mode + for causal in masks: + args.causal = causal + print(f"Comparing {options.workload}: {shape_mode} / {'causal' if causal else 'noncausal'}...", + file=sys.stderr, flush=True) + report = compare(options, runner_arguments(args, options.workload)) + cases.append({"shape_mode": shape_mode, "causal": causal, **report}) + return {"workload": args.workload, "verification": "all-pass", "cases": cases, + "timing": cases[0]["timing"]} + + +def print_summary(report): + cases = report.get("cases") + first = cases[0] if cases is not None else report + rust = first["cuda_oxide"] + provenance = rust["provenance"] + batch, seq_q, seq_k, heads, dim = rust["shape"] + print(f"{report['workload']}: B={batch}, Sq={seq_q}, Sk={seq_k}, H={heads}, D={dim}") + if provenance["iters"]: + print(f"Timing: direct CUDA-event median (us); warmup={provenance['warmup']}, " + f"iters={provenance['iters']}") + print("CUDA-event intervals include host submission gaps; not Nsight kernel-active times.") + else: + print("Verification only (timing disabled).") + if cases is not None: + print("| Shapes | Attention | cute-rs (us) | CuTeDSL (us) | Relative delta |") + print("| --- | --- | ---: | ---: | ---: |") + for case in cases: + rust_time = case["cuda_oxide"]["metric_value"] + dsl_time = case["python_reference"]["metric_value"] + values = ("verified | verified | n/a" if rust_time is None else + f"{rust_time:.3f} | {dsl_time:.3f} | {(rust_time / dsl_time - 1) * 100:+.1f}%") + print(f"| {case['shape_mode'].title()} | {'Causal' if case['causal'] else 'Noncausal'} | {values} |") + if provenance["iters"]: + print("Relative delta = (cute-rs / CuTeDSL - 1) * 100%; negative means a shorter cute-rs event interval.") + else: + print(f"Case: {provenance['shape_specialization']} / {'causal' if rust['causal'] else 'noncausal'}") + for key, label in (("python_reference", "CuTeDSL"), ("cuda_oxide", "cute-rs")): + metric = report[key]["metric_value"] + print(f"{label}: {metric:.3f} us" if metric is not None + else f"{label}: verified (timing disabled)") + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) + parser.add_argument("--workload", choices=("decode", "prefill"), required=True) + parser.add_argument("--rust-binary", type=Path, required=True) + parser.add_argument("--all-cases", action="store_true", + help="compare runtime/static shapes; prefill also sweeps causal/noncausal attention") + parser.add_argument("--json", type=Path, help="save both checked reports and output comparison") + parser.add_argument("--timeout", type=float, default=600, help="seconds per implementation") + options, runner_cli = parser.parse_known_args(argv) + if not math.isfinite(options.timeout) or options.timeout <= 0: + parser.error("timeout must be positive and finite") + require_reference_environment() + try: + report = compare_all_cases(options, runner_cli) if options.all_cases else compare(options, runner_cli) + if options.json: + options.json.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + print_summary(report) + except (OSError, ValueError, KeyError, TypeError, RuntimeError, subprocess.TimeoutExpired) as error: + print(f"FAIL: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/README.md b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/README.md new file mode 100644 index 0000000000..219b4fe61a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/README.md @@ -0,0 +1,129 @@ +# CuTe DSL FlashAttention v2 reference + +`flash_attention_v2.py` is an unchanged copy of the public CUTLASS 4.7 example: + +- Repository commit: `dcf215af68a2d08d305076c152a06f201728cd53`. +- [Original source](https://github.com/NVIDIA/cutlass/blob/dcf215af68a2d08d305076c152a06f201728cd53/examples/python/CuTeDSL/cute/ampere/kernel/attention/flash_attention_v2.py). +- SHA-256: `263df1fd783729286cd6fa5ec994ec86e7614837bd1ff204eb6687192a8bae2d`. +- License: BSD-3-Clause; the original copyright, conditions, and disclaimer + remain in the source header. + +This is the ordinary CuTe DSL implementation using `@cute.jit`, `cp.async`, +warp-level FP16/BF16 MMA, and fused online softmax. It requires no other +repository helper files. The installed DSL runtime has its own NVIDIA EULA; +runtime source is not vendored here. + +Use Python 3.10 or later supported by CUTLASS DSL 4.7, a CUDA-enabled PyTorch +build supporting the selected GPU, CUDA bindings, and `nvidia-cutlass-dsl==4.7.0`. +The shared [`requirements.txt`](../../../requirements.txt) selects the CUDA 13 +package extra. The upstream CUDA 12.9 alternative is `nvidia-cutlass-dsl==4.7.0` +without that extra. Match the CUDA runtime, toolkit, and driver to the package +variant. + +Prepare the shared Python environment once for both prefill and decode, +from `cuda-oxide/`: + +```bash +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Select a CUDA-enabled PyTorch build that supports SM120 and matches your CUDA +runtime. The normal `cargo oxide run` command finds the shared environment +automatically; no activation or Python environment-variable export is needed. +Python selection is `--python PATH` after `--`, `CUDA_OXIDE_CUTE_PYTHON`, the +example's `.venv/bin/python`, the shared `cute-rs-preview/.venv/bin/python`, then `python3`. +See the [example README](../README.md) for direct-event comparisons. +The main command with `--all-cases --warmup 100 --iters 31` prints direct +CUDA-event medians for all four runtime/static and causal/noncausal cases. +These intervals can include host launch/submission delays, which Nsight Systems +kernel durations exclude. Comparing those two methods does not establish a +regression, and runs with different warmup and sample counts use different +sampling protocols. + +`compare.py` retains the separate Nsight Systems profiling workflow and +invokes the Rust binary with `--rust-only`. With Nsight Systems installed, +reproduce its four-case measurement protocol from `cuda-oxide/` using the +shared Python environment: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build --arch sm_120a --device-codegen-crate fmha_prefill_sm120a_cute \ + --cargo-target-dir /tmp/cuda-oxide-fmha-prefill -- \ + --manifest-path cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.toml --release + +for shape in runtime static; do + for mask in causal noncausal; do + CUTE_DSL_ARCH=sm_120a cute-rs-preview/.venv/bin/python \ + cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/compare.py \ + --rust-binary /tmp/cuda-oxide-fmha-prefill/release/fmha_prefill_sm120a_cute \ + --seq-q 1024 --seq-k 1024 --heads 8 --shape-mode "$shape" --"$mask" \ + --warmup 100 --samples 31 --rounds 3 \ + --output-dir "/tmp/prefill-nsys-$shape-$mask" + done +done +``` + +The harness reports the median of round medians and retains traces and raw +samples in fresh evidence directories. Hardware state, clocks and software +versions can affect new results; reproducing the protocol does not guarantee +the archived numbers. + +From this directory, run the causal prefill workload on the default device: + +```bash +CUTE_DSL_ARCH=sm_120a ../../../.venv/bin/python flash_attention_v2.py \ + --dtype Float16 --batch_size 1 --num_head 8 \ + --seqlen_q 1024 --seqlen_k 1024 --head_dim 64 \ + --m_block_size 64 --n_block_size 64 --num_threads 128 \ + --softmax_scale 0.125 --is_causal +``` + +Omit `--is_causal` for the noncausal FMHA workload. Both commands retain the +upstream PyTorch output check. The default benchmark uses three warmups and +ten iterations; add `--use_cold_l2` only when comparing the corresponding +cache policy. + +The selected tile allocates three 64-by-64 FP16 shared-memory arrays, totaling +24 KiB. Each array is already a multiple of its 1024-byte alignment. This fits +the SM120 per-block limit of 99 KiB documented in the +[Blackwell tuning guide](https://docs.nvidia.com/cuda/blackwell-tuning-guide/). +The source's `can_implement` checks the larger SM80 capacity, so other tile +choices must separately respect the actual target's capacity. + +The reference accepts FP16/BF16 with a head dimension divisible by eight, +16-byte contiguous alignment, a warp-multiple thread count, and +`2 * m_block_size` divisible by the thread count. Head dimensions are padded +to multiples of 32 for storage. It implements equal Q/K/V head counts and +optional top-left causal masking; it produces no training LSE output and +does not implement GQA, paged KV, or split-KV decode. Its use of the older +warp MMA instruction makes this an SM120 candidate without `tcgen05` or TMEM; +compatibility and performance still require a GPU run. + +`launch.py` invokes this unchanged kernel with the Rust example's deterministic +inputs and an independent CPU FP64 oracle. It supports the same `--seq-q`, +`--seq-k`, `--heads`, `--causal`/`--noncausal`, `--input-scale`, `--zero-query`, +`--warmup`, `--iters`, `--output`, and `--json` options as the Rust host. +`benchmark.py` is the same checked runner. The original script's own CLI above +uses its upstream random input generator instead of the shared fixture. + +## Validation + +`validate.py` in the example directory checks lengths 1, 64, 128, 192 and +1024, both masks, unequal head-local inputs, rectangular noncausal shapes, +large logits and zero queries. With `--reference` it also checks that the +input hashes agree and compares every Rust and CuTeDSL output element, +retaining per-case logs and JSON reports. From the example directory, with +the shared environment's interpreter as `python`: + +```bash +CUTE_DSL_ARCH=sm_120a python validate.py \ + --binary /path/to/fmha_prefill_sm120a_cute --reference --output-dir /tmp/prefill-validation +CUTE_DSL_ARCH=sm_120a python validate.py \ + --binary /path/to/fmha_prefill_sm120a_cute --reference --shape-mode static \ + --output-dir /tmp/prefill-static-validation +python -m unittest discover -s reference -p test_fixture.py +``` + +`--rust-only` runs the Rust binary without the comparison; with `--warmup 0 +--iters 0` it performs two checked launches without timing. diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/benchmark.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/benchmark.py new file mode 100644 index 0000000000..0338b1e796 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/benchmark.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated CuTeDSL timing; shares all options and checks with launch.py.""" + +from launch import parse_args, run + + +if __name__ == "__main__": + run(parse_args()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/compare.py new file mode 100644 index 0000000000..6899dec882 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/compare.py @@ -0,0 +1,458 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compare checked attention runners using serial Nsight Systems kernel times. + +Select the GPU externally with CUDA_VISIBLE_DEVICES; --device is the resulting +visible ordinal. Defaults: 100 warmups, 31 samples, three alternating rounds. +Both runners must use launch.py's CLI/report protocol and issue two correctness +launches, then warmups, then samples of exactly one selected attention kernel. +CUDA event measurements are retained as runner evidence, never aggregated here. +The historical single-query prefill baseline uses --seq-q 1 --noncausal. +Dedicated split-KV decode uses fmha_decode_sm120a_cute/reference/compare.py, +which measures its partial and merge kernels as one ordered sequence. +""" + +import argparse +from collections import Counter +from datetime import datetime, timezone +import hashlib +import json +import math +import os +from pathlib import Path +import re +import shutil +import signal +import sqlite3 +import statistics +import struct +import subprocess +import sys +import tempfile +import time + +from launch import ATOL, RTOL, HEAD_DIM, STATIC_SHAPE, Q_SEED, K_SEED, V_SEED, decode_f16, mix32, verify + + +CORRECTNESS_LAUNCHES = 2 + + +def reject_nonfinite_json(value): + raise ValueError(f"nonfinite JSON value: {value}") + + +def hash_f16_bytes(data): + result = 0xCBF29CE484222325 + for byte in data: + result = ((result ^ byte) * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF + return f"{result:016x}" + + +def expected_input_hashes(args): + """Hash the shared deterministic fixture without allocating large tensors.""" + hashes = {} + for name, length, seed, scale in ( + ("q", args.seq_q, Q_SEED, args.input_scale), + ("k", args.seq_k, K_SEED, args.input_scale), + ("v", args.seq_k, V_SEED, 1), + ): + result = 0xCBF29CE484222325 + for index in range(length * args.heads * HEAD_DIM): + value = 0.0 if name == "q" and args.zero_query else ((mix32(index ^ seed) & 63) - 32) * scale / 128.0 + for byte in struct.pack(" 1: + raise ValueError("CPU verification tolerance exceeded") + provenance = result.get("provenance") + require_fields(provenance, {"device": args.device, "warmup": args.warmup, + "iters": args.samples, "dtype": "float16", + "causal_specialization": "static", + "shape_specialization": args.shape_mode, + "tool": "cuda-events-single-kernel", "timing_mode": "direct"}, "provenance") + capability = provenance.get("compute_capability") + if (not isinstance(provenance.get("gpu"), str) or not provenance["gpu"] + or not isinstance(capability, list) or len(capability) != 2 + or any(type(value) is not int or value < 0 for value in capability)): + raise ValueError("missing GPU identity/compute capability in runner report") + if len(data) != count * 2: + raise ValueError(f"output contains {len(data)} bytes; expected {count * 2}") + output = decode_f16(data) + if not all(map(math.isfinite, output)): + raise ValueError("nonfinite FP16 output") + if result.get("output_hash") != hash_f16_bytes(data): + raise ValueError("FP16 output does not match the checked report's output hash") + return result, output + + +def kernel_samples(path, warmup, samples, *, expected_name=None, name_contains=None): + """Resolve one full kernel identity, then remove checks/warmup by launch count. + + Nsight may truncate shortName. Prefer demangledName for identity; exact + overrides can name either field, but a colliding shortName is ambiguous. + Older schemas without demangledName retain exact shortName behavior. + """ + if bool(expected_name) == bool(name_contains): + raise ValueError("provide one exact kernel name or one name resolver") + if warmup < 0 or samples <= 0: + raise ValueError("warmup must be nonnegative and samples positive") + try: + with sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True) as db: + columns = {row[1] for row in db.execute("PRAGMA table_info(CUPTI_ACTIVITY_KIND_KERNEL)")} + device_column = "k.deviceId" if "deviceId" in columns else "NULL" + full_names = "demangledName" in columns + name_column = "d.value" if full_names else "s.value" + name_join = "LEFT JOIN StringIds AS d ON d.id = k.demangledName " if full_names else "" + rows = db.execute( + "SELECT k.start, k.end, s.value, " + device_column + ", " + name_column + " " + "FROM CUPTI_ACTIVITY_KIND_KERNEL AS k " + "LEFT JOIN StringIds AS s ON s.id = k.shortName " + + name_join + + "ORDER BY k.start, k.end" + ).fetchall() + gpu_metadata = [] + if db.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='TARGET_INFO_GPU'").fetchone(): + cursor = db.execute("SELECT * FROM TARGET_INFO_GPU") + fields = [column[0] for column in cursor.description] + gpu_metadata = [dict(zip(fields, [value.hex() if isinstance(value, bytes) else value + for value in row])) for row in cursor] + except sqlite3.Error as error: + raise ValueError(f"cannot read kernel trace {path}: {error}") from error + for start, end, short_name, _, full_name in rows: + if any(not isinstance(name, str) or not name for name in (short_name, full_name)): + raise ValueError("unresolved kernel shortName/demangledName in trace") + if type(start) is not int or type(end) is not int or end <= start: + raise ValueError(f"invalid kernel timestamps: {start!r}, {end!r}") + if name_contains: + matches = sorted({row[4] for row in rows if name_contains in row[4]}) + else: + matches = sorted({row[4] for row in rows if expected_name in (row[2], row[4])}) + if len(matches) != 1: + raise ValueError(f"expected one matching kernel name for {expected_name or name_contains!r}, found {matches!r}") + kernel_name = matches[0] + selected = [row for row in rows if row[4] == kernel_name] + expected = CORRECTNESS_LAUNCHES + warmup + samples + if len(selected) != expected: + raise ValueError(f"expected {expected} launches of exact kernel {kernel_name!r}, found {len(selected)}") + devices = {row[3] for row in selected} + if len(devices) != 1: + raise ValueError(f"selected kernel ran on multiple devices: {devices}") + durations = [(row[1] - row[0]) / 1000 for row in selected] + split = CORRECTNESS_LAUNCHES + warmup + return { + "kernel_name": kernel_name, "kernel_name_field": "demangledName" if full_names else "shortName", + "kernel_short_names": sorted({row[2] for row in selected}), "kernel_launches": len(selected), + "trace_device_id": selected[0][3], "trace_gpu_metadata": gpu_metadata, + "correctness_samples_us": durations[:CORRECTNESS_LAUNCHES], + "warmup_samples_us": durations[CORRECTNESS_LAUNCHES:split], "samples_us": durations[split:], + "selected_kernel_timestamps_ns": [{"start": row[0], "end": row[1]} for row in selected], + "auxiliary_kernel_counts": dict(sorted(Counter(row[4] for row in rows if row[4] != kernel_name).items())), + } + + +def aggregate_rounds(entries): + summary = {} + for entry in entries: + result = summary.setdefault(entry["label"], {"round_medians_us": [], "samples_us": []}) + result["round_medians_us"].append(entry["median_us"]) + result["samples_us"].extend(entry["samples_us"]) + for result in summary.values(): + result["median_of_round_medians_us"] = statistics.median(result["round_medians_us"]) + return summary + + +def run_command(command, timeout, environment, log_path): + """Finish each process before starting another; terminate its children on timeout.""" + with log_path.open("w") as log: + with subprocess.Popen(command, env=environment, stdout=log, stderr=subprocess.STDOUT, + start_new_session=True) as process: + try: + return process.wait(timeout=timeout) + except (subprocess.TimeoutExpired, KeyboardInterrupt): + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + raise + + +def file_identity(path): + path = path.resolve(strict=True) + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() if hasattr(hashlib, "file_digest") else None + if digest is None: + digest = hashlib.sha256(stream.read()).hexdigest() + return {"path": str(path), "sha256": digest, "size_bytes": path.stat().st_size} + + +def source_snapshot(reference_runner): + """Current source evidence is separate from proof of how a binary was built.""" + here = Path(__file__).resolve() + oxide = next((parent for parent in here.parents if (parent / "rust-toolchain.toml").is_file()), None) + paths = set() + for example in (here.parent.parent, reference_runner.parent.parent): + for child in (example / "src", example / "reference"): + if child.is_dir(): + paths.update(path for path in child.rglob("*") if path.suffix in (".py", ".rs", ".txt", ".md")) + paths.update(path for path in (example / "Cargo.toml", example / "Cargo.lock") if path.is_file()) + paths.add(reference_runner) + pins = {} + if oxide: + for relative in ("rust-toolchain.toml", "Cargo.lock", + "crates/cargo-oxide/src/cutlass_toolchain.rs", + "crates/cuda-oxide-mlir-export/src/profile.rs", + "crates/cuda-oxide-mlir-export/src/cute_attention.rs", + "crates/cuda-oxide-mlir-export/src/nvvm_attention.rs", + "cute-rs-preview/cute-rs/src/mma/mod.rs"): + path = oxide / relative + if path.is_file(): + paths.add(path) + pin_path = oxide / "crates/cargo-oxide/src/cutlass_toolchain.rs" + match = re.search(r'const CUTLASS_VERSION: &str = "([^"]+)"', pin_path.read_text()) if pin_path.is_file() else None + pins["rust_cutlass_compiler_source_pin"] = match[1] if match else None + pins["rust_toolchain_file"] = (oxide / "rust-toolchain.toml").read_text() + requirements = here.parent / "requirements.txt" + if requirements.is_file(): + pins["reference_requirements"] = requirements.read_text() + return {"scope": "current attention sources and selected compiler/API files; not binary build provenance", + "files": [file_identity(path) for path in sorted(paths) if path.is_file()], "pins": pins} + + +def parse_arguments(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rust-binary", type=Path, required=True) + parser.add_argument("--rust-build-manifest", type=Path, + help="build provenance JSON; must include matching top-level binary_sha256") + parser.add_argument("--python", default=sys.executable, help="reference Python interpreter") + parser.add_argument("--reference-runner", type=Path, default=Path(__file__).with_name("launch.py")) + parser.add_argument("--rust-kernel-name", + help="exact Nsight shortName or demangledName; otherwise resolve one full fmha_prefill name") + parser.add_argument("--reference-kernel-name", help="exact shortName or demangledName; otherwise resolve one full FlashAttentionForwardAmpere name") + parser.add_argument("--nsys", default=shutil.which("nsys") or "nsys") + parser.add_argument("--output-dir", type=Path, required=True, help="parent for a fresh evidence directory") + parser.add_argument("--seq-q", type=int, default=64) + parser.add_argument("--seq-k", type=int, default=64) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--shape-mode", choices=("runtime", "static"), default="runtime", + help="match both runners' shape and stride specialization; static requires Sq=Sk=1024, H=8") + masks = parser.add_mutually_exclusive_group() + masks.add_argument("--causal", dest="causal", action="store_true") + masks.add_argument("--noncausal", dest="causal", action="store_false") + parser.set_defaults(causal=True) + parser.add_argument("--input-scale", type=int, choices=(1, 4, 16), default=1) + parser.add_argument("--zero-query", action="store_true") + parser.add_argument("--device", type=int, default=0, help="ordinal after external CUDA_VISIBLE_DEVICES remapping") + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--samples", "--iters", type=int, default=31) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--timeout", type=float, default=600, help="seconds per tool invocation") + args = parser.parse_args(argv) + if args.shape_mode == "static" and (args.seq_q, args.seq_k, args.heads) != STATIC_SHAPE: + parser.error("static shape mode requires --seq-q 1024 --seq-k 1024 --heads 8") + if min(args.seq_q, args.seq_k, args.heads, args.samples, args.rounds) <= 0: + parser.error("lengths, heads, samples, and rounds must be positive") + if (args.seq_q != 1 and args.seq_q % 64) or args.seq_k % 64: + parser.error("query length must be 1 or a multiple of 64; KV length must be a multiple of 64") + if args.causal and args.seq_q != args.seq_k: + parser.error("causal cases require equal query and KV lengths") + if max(args.seq_q, args.seq_k) * args.heads * HEAD_DIM > 0xFFFFFFFF: + parser.error("tensor element count must fit u32 indexing") + if args.device < 0 or args.warmup < 0 or not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("device/warmup must be nonnegative and timeout finite and positive") + if any(name is not None and not name.strip() + for name in (args.rust_kernel_name, args.reference_kernel_name)): + parser.error("kernel names must be nonempty") + for field in ("rust_binary", "reference_runner", "rust_build_manifest"): + path = getattr(args, field) + if path is None: + continue + if not path.is_file() or field == "rust_binary" and not os.access(path, os.X_OK): + parser.error(f"{field} must be an existing {'executable ' if field == 'rust_binary' else ''}file: {path}") + setattr(args, field, path.resolve()) + return args + + +def main(argv=None): + args = parse_arguments(argv) + args.output_dir.mkdir(parents=True, exist_ok=True) + directory = Path(tempfile.mkdtemp(prefix="run-", dir=args.output_dir.resolve())) + environment = os.environ.copy() + report = { + "status": "running", "metric": "Nsight Systems kernel-active duration", "metric_unit": "us", + "duration_source": "CUPTI_ACTIVITY_KIND_KERNEL.end - start (nanoseconds / 1000)", + "started_ns": time.time_ns(), "started_utc": datetime.now(timezone.utc).isoformat(), + "parameters": {key: getattr(args, key) for key in ("seq_q", "seq_k", "heads", "causal", "shape_mode", + "input_scale", "zero_query", "device", "warmup", "samples", "rounds", "timeout")}, + "launch_contract": {"correctness": CORRECTNESS_LAUNCHES, "warmup": args.warmup, "measured": args.samples}, + "environment": {key: environment.get(key) for key in ("CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", + "CUTE_DSL_ARCH", "CUDA_HOME", "CUDA_PATH")}, + "aggregation": "median of each round's kernel samples, then median of round medians", + "runner_event_times": "retained in runner_report; excluded from kernel-active summaries", + "execution_order": [["rust", "cutlass"] if round_index % 2 == 0 else ["cutlass", "rust"] + for round_index in range(args.rounds)], + "provenance": {}, "commands": [], "runs": [], + } + + def save(): + temporary = directory / "report.json.tmp" + temporary.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n") + temporary.replace(directory / "report.json") + + def execute(command, log_path): + entry = {"command": command, "log": str(log_path), "status": "running", "started_ns": time.time_ns()} + report["commands"].append(entry) + save() + try: + entry["exit_code"] = run_command(command, args.timeout, environment, log_path) + if entry["exit_code"]: + raise ValueError(f"command exited {entry['exit_code']}; see {log_path}") + entry["status"] = "passed" + except (OSError, ValueError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + entry.update(status="failed", error=str(error) or "interrupted") + raise + finally: + entry["finished_ns"] = time.time_ns() + save() + + print(f"Evidence: {directory}", flush=True) + active = None + save() + try: + provenance = report["provenance"] + provenance["rust_binary"] = file_identity(args.rust_binary) + provenance["reference_runner"] = file_identity(args.reference_runner) + provenance["source_snapshot"] = source_snapshot(args.reference_runner) + provenance["rust_build"] = {"status": "unavailable; current source pins do not prove the binary's build toolchain"} + if args.rust_build_manifest: + manifest = json.loads(args.rust_build_manifest.read_text(), parse_constant=reject_nonfinite_json) + if not isinstance(manifest, dict) or manifest.get("binary_sha256") != provenance["rust_binary"]["sha256"]: + raise ValueError("Rust build manifest must contain binary_sha256 matching --rust-binary") + provenance["rust_build"] = {"file": file_identity(args.rust_build_manifest), "manifest": manifest} + for key, command in ( + ("nsys_version", [args.nsys, "--version"]), + ("python_version", [args.python, "--version"]), + ("git_head", ["git", "-C", str(Path(__file__).parent), "rev-parse", "HEAD"]), + ("git_status", ["git", "-C", str(Path(__file__).parent), "status", "--short"]), + ): + log = directory / f"{key}.log" + execute(command, log) + provenance[key] = log.read_text().strip() + report["expected_input_hashes"] = expected_input_hashes(args) + common = ["--seq-q", str(args.seq_q), "--seq-k", str(args.seq_k), "--heads", str(args.heads), + "--shape-mode", args.shape_mode, + "--causal" if args.causal else "--noncausal", "--input-scale", str(args.input_scale), + "--device", str(args.device), "--warmup", str(args.warmup), "--iters", str(args.samples)] + if args.zero_query: + common.append("--zero-query") + implementations = {"rust": [str(args.rust_binary), "--rust-only"], "cutlass": [args.python, str(args.reference_runner)]} + output_hashes, kernel_names, baseline = {}, {}, None + for round_index, order in enumerate(report["execution_order"], start=1): + outputs = {} + for label in order: + prefix = directory / f"round-{round_index:02d}-{label}" + result_path, output_path = prefix.with_suffix(".json"), prefix.with_suffix(".f16") + trace_path, database_path = prefix.with_suffix(".nsys-rep"), prefix.with_suffix(".sqlite") + profile_log, export_log = prefix.with_suffix(".profile.log"), prefix.with_suffix(".export.log") + profile = [args.nsys, "profile", "--trace=cuda,nvtx", "--sample=none", "--cpuctxsw=none", + "--output", str(prefix), *implementations[label], *common, + "--json", str(result_path), "--output", str(output_path)] + export = [args.nsys, "export", "--type=sqlite", "--output", str(database_path), str(trace_path)] + active = {"round": round_index, "label": label, "status": "running", "started_ns": time.time_ns(), + "correctness_report": str(result_path), "output": str(output_path), + "raw_trace": str(trace_path), "sqlite": str(database_path)} + report["runs"].append(active) + save() + if any(path.exists() for path in (result_path, output_path, trace_path, database_path)): + raise ValueError("refusing to reuse existing profiling evidence") + print(f"Round {round_index}/{args.rounds}: {label}", flush=True) + execute(profile, profile_log) + checked, output = validate_report(result_path, output_path, args, active["started_ns"], + report["expected_input_hashes"]) + active["runner_report"] = checked + outputs[label] = output + gpu = {key: checked["provenance"][key] for key in ("gpu", "compute_capability", "device")} + if baseline is None: + baseline = gpu + provenance["gpu"] = gpu + elif gpu != baseline: + raise ValueError(f"GPU identity changed between runners/rounds: {gpu!r} != {baseline!r}") + if label in output_hashes and output_hashes[label] != checked["output_hash"]: + raise ValueError(f"{label} output changed across rounds") + output_hashes[label] = checked["output_hash"] + if not trace_path.is_file(): + raise ValueError(f"missing raw Nsight Systems trace: {trace_path}") + execute(export, export_log) + name = args.rust_kernel_name if label == "rust" else args.reference_kernel_name + resolver = "fmha_prefill" if label == "rust" else "FlashAttentionForwardAmpere" + trace_data = kernel_samples(database_path, args.warmup, args.samples, expected_name=name, + name_contains=None if name else resolver) + if label in kernel_names and kernel_names[label] != trace_data["kernel_name"]: + raise ValueError(f"{label} kernel name changed across rounds") + kernel_names[label] = trace_data["kernel_name"] + if "trace_device_id" not in provenance: + provenance["trace_device_id"] = trace_data["trace_device_id"] + elif trace_data["trace_device_id"] != provenance["trace_device_id"]: + raise ValueError("Nsight kernel device changed between runners/rounds") + active.update(trace_data, median_us=statistics.median(trace_data["samples_us"]), + status="passed", finished_ns=time.time_ns()) + save() + report.setdefault("cross_checks", []).append({"round": round_index, "status": "passed", + **verify(outputs["rust"], outputs["cutlass"])}) + save() + report.update(status="passed", summary=aggregate_rounds(report["runs"])) + except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired, KeyboardInterrupt) as error: + report.update(status="failed", error=str(error) or "interrupted") + if active is not None and active["status"] == "running": + active.update(status="failed", error=report["error"], finished_ns=time.time_ns()) + print(f"FAIL: {report['error']}", file=sys.stderr) + finally: + report.update(finished_ns=time.time_ns(), finished_utc=datetime.now(timezone.utc).isoformat()) + save() + if report["status"] == "passed": + for label, result in report["summary"].items(): + print(f"{label}: {result['median_of_round_medians_us']:.3f} us kernel-active (median of round medians)") + print(f"Comparison report: {directory / 'report.json'}") + return 0 if report["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/flash_attention_v2.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/flash_attention_v2.py new file mode 100644 index 0000000000..95e3497b87 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/flash_attention_v2.py @@ -0,0 +1,1344 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +from types import SimpleNamespace +from typing import Type, Callable + +import cuda.bindings.driver as cuda +import cutlass.cute.testing as testing +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cute.runtime import from_dlpack +import cutlass.pipeline as pipeline +import cutlass.utils as utils + +""" +A flash attention v2 forward pass example for NVIDIA Ampere SM80 architecture using CUTE DSL. + +- Matrix Q is BxSqxNxH, B is batch dimension, Sq is query sequence length, N is number of heads, H is head dimension +- Matrix K is BxSkxNxH, B is batch dimension, Sk is key sequence length, N is number of heads, H is head dimension +- Matrix V is BxSkxNxH, B is batch dimension, Sk is key sequence length, N is number of heads, H is head dimension +- Matrix O is BxSqxNxH, B is batch dimension, Sq is query sequence length, N is number of heads, H is head dimension + +This kernel supports the following features: + - Utilizes CpAsync for efficient memory operations + - Utilizes Ampere's tensor core for matrix multiply-accumulate (MMA) operations + - Utilizes register pipeline to overlap shared memory-to-register transfers with computations. + - Leverages DSL to implement an integrated online softmax fusion pattern. + +This kernel works as follows: +1. Load Q and K matrices from global memory (GMEM) to shared memory (SMEM) using CpAsync operations. +2. Perform matrix multiply-accumulate (MMA) operations using tensor core instructions to compute intermediate result S. +3. Apply padding mask or causal mask to S during initial iterations. +4. Apply online softmax to S and rescale O using results from previous iteration. +5. Load V matrices and perform matrix multiply-accumulate (MMA) operations to compute final result O. +6. Normalize O after all iterations complete and store result back to global memory (GMEM). + +To run this example: + +.. code-block:: bash + + python examples/ampere/flash_attention_v2.py \ + --dtype Float16 --head_dim 128 --m_block_size 128 --n_block_size 128 \ + --num_threads 128 --batch_size 1 --seqlen_q 1280 --seqlen_k 1536 \ + --num_head 16 --softmax_scale 1.0 --is_causal + +The above command configures the model to use float16 for inputs and outputs. The problem dimensions +are set to: batch size of 1, query sequence length of 1280, key sequence length of 1536, head dimension +of 128, and 16 attention heads. The softmax scale is set to 1.0 and causal masking is enabled. The computation +uses tiles of size 128x128 for m and n dimensions, and utilizes 128 parallel threads. + +To collect the performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/ampere/flash_attention_v2.py \ + --dtype Float16 --head_dim 128 --m_block_size 128 --n_block_size 128 \ + --num_threads 128 --batch_size 1 --seqlen_q 1280 --seqlen_k 1536 \ + --num_head 16 --softmax_scale 1.0 --is_causal --skip_ref_check + +There are some constraints for this example: +* Only fp16 and bf16 data types are supported. +* The contiguous dimension of each tensor must be at least 16 bytes aligned. +* The log-sum-exp(for training) is not computed in the kernel. +* The values of `m_block_size`, `n_block_size`, and `head_dim` must be selected to stay within shared memory capacity limits. +* `m_block_size * 2` must be divisible by `num_threads`, otherwise the kernel will not be able to get the correct result. +""" + + +class FlashAttentionForwardAmpere: + def __init__( + self, + head_dim: int, + m_block_size: int = 128, + n_block_size: int = 128, + num_threads: int = 128, + is_causal: bool = False, + ): + """Initializes the configuration for a flash attention v2 kernel. + + All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension + should be a multiple of 8. + + :param head_dim: head dimension + :type head_dim: int + :param m_block_size: m block size + :type m_block_size: int + :param n_block_size: n block size + :type n_block_size: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + """ + self._head_dim = head_dim + self._m_block_size = m_block_size + self._n_block_size = n_block_size + # padding head_dim to a multiple of 32 as k_block_size + self._head_dim_padded = (head_dim + 31) // 32 * 32 + self._num_threads = num_threads + self._is_causal = is_causal + + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, num_threads=num_threads + ) + + @staticmethod + def can_implement( + dtype, head_dim, m_block_size, n_block_size, num_threads, is_causal + ) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param m_block_size: m block size + :type m_block_size: int + :param n_block_size: n block size + :type n_block_size: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :type is_causal: bool + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + # Check if data type is fp16 or bf16 + if dtype != cutlass.Float16 and dtype != cutlass.BFloat16: + return False + + # Check if head dimension is a multiple of 8 + if head_dim % 8 != 0: + return False + + # Check if number of threads is a multiple of 32 + if num_threads % 32 != 0: + return False + + # Check if block size setting is out of shared memory capacity + # Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size + smem_usage = (m_block_size * head_dim + n_block_size * head_dim * 2) * 2 + smem_capacity = utils.get_smem_capacity_in_bytes("sm_80") + if smem_usage > smem_capacity: + return False + + # Check if twice the block size is divisible by the number of threads + if (m_block_size * 2) % num_threads != 0: + return False + + return True + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + softmax_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Configures and launches the flash attention v2 kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(seqlen_q * num_head * head_dim, num_head * head_dim, head_dim, 1) + + Prepares the shared memory layout, tiled copy atoms, tiled mma and shared memory storage. + Then launches the kernel function with the prepared parameters. + + :param mQ: query tensor + :type mQ: cute.Tensor + :param mK: key tensor + :type mK: cute.Tensor + :param mV: value tensor + :type mV: cute.Tensor + :param mO: output tensor + :type mO: cute.Tensor + :param softmax_scale: softmax scale + :type softmax_scale: cutlass.Float32 + """ + # Get the data type and check if it is fp16 or bf16 + if cutlass.const_expr( + not ( + mQ.element_type == mK.element_type == mV.element_type == mO.element_type + ) + ): + raise TypeError("All tensors must have the same data type") + if cutlass.const_expr( + not ( + mQ.element_type == cutlass.Float16 + or mQ.element_type == cutlass.BFloat16 + ) + ): + raise TypeError("Only Float16 or BFloat16 is supported") + self._dtype: Type[cutlass.Numeric] = mQ.element_type + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: Q/K/V + # /////////////////////////////////////////////////////////////////////////////// + smem_k_block_size = 64 if self._head_dim_padded % 64 == 0 else 32 + swizzle_bits = 3 if smem_k_block_size == 64 else 2 + sQ_layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 3), + 0, + cute.make_layout((8, smem_k_block_size), stride=(smem_k_block_size, 1)), + ) + sQ_layout = cute.tile_to_shape( + sQ_layout_atom, + (self._m_block_size, self._head_dim_padded), + (0, 1), + ) + + sKV_layout_atom = sQ_layout_atom + sKV_layout = cute.tile_to_shape( + sKV_layout_atom, + (self._n_block_size, self._head_dim_padded), + (0, 1), + ) + + sO_layout = sQ_layout + + @cute.struct + class SharedStorage: + sQ: cute.struct.Align[ + cute.struct.MemRange[self._dtype, cute.cosize(sQ_layout)], 1024 + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self._dtype, cute.cosize(sKV_layout)], 1024 + ] + sV: cute.struct.Align[ + cute.struct.MemRange[self._dtype, cute.cosize(sKV_layout)], 1024 + ] + + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self._dtype.width + # atom_async_copy: async copy atom for QKV load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self._dtype, + num_bits_per_copy=universal_copy_bits, + ) + # atom_universal_copy: universal copy atom for O store + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self._dtype, + num_bits_per_copy=universal_copy_bits, + ) + # tQKV_layout: thread layout for QKV load + tQKV_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems + tQKV_layout = cute.make_layout( + (self._num_threads // tQKV_shape_dim_1, tQKV_shape_dim_1), + stride=(tQKV_shape_dim_1, 1), + ) + # tO_layout: thread layout for O store + tO_layout = tQKV_layout + + # Value layouts for copies + vQKV_layout = cute.make_layout((1, async_copy_elems)) + vO_layout = vQKV_layout + + # gmem_tiled_copy_QKV: tiled copy for QKV load + gmem_tiled_copy_QKV = cute.make_tiled_copy_tv( + atom_async_copy, tQKV_layout, vQKV_layout + ) + # gmem_tiled_copy_O: tiled copy for O store + gmem_tiled_copy_O = cute.make_tiled_copy_tv( + atom_universal_copy, tO_layout, vO_layout + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Tiled mma + # /////////////////////////////////////////////////////////////////////////////// + tiled_mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(self._dtype, cutlass.Float32, (16, 8, 16)), + (self._num_threads // 32, 1, 1), + permutation_mnk=(self._num_threads // 32 * 16, 16, 16), + ) + + # grid_dim: (m_block, batch_size, num_head) + grid_dim = ( + cute.ceil_div(mQ.shape[1], self._m_block_size), + cute.size(mQ.shape[0]), + cute.size(mQ.shape[2]), + ) + LOG2_E = 1.4426950408889634074 + softmax_scale_log2 = softmax_scale * LOG2_E + self.kernel( + mQ, + mK, + mV, + mO, + softmax_scale_log2, + sQ_layout, + sKV_layout, + sO_layout, + gmem_tiled_copy_QKV, + gmem_tiled_copy_O, + tiled_mma, + SharedStorage, + ).launch( + grid=grid_dim, + block=[self._num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + sQ_layout: cute.ComposedLayout, + sKV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + gmem_tiled_copy_QKV: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma: cute.TiledMma, + SharedStorage: cutlass.Constexpr, + ): + """Kernel function for flash attention v2. + + :param mQ: query tensor + :type mQ: cute.Tensor + :param mK: key tensor + :type mK: cute.Tensor + :param mV: value tensor + :type mV: cute.Tensor + :param mO: output tensor + :type mO: cute.Tensor + :param softmax_scale_log2: softmax scale log2 + :type softmax_scale_log2: cutlass.Float32 + :param sQ_layout: query layout + :type sQ_layout: cute.ComposedLayout + :param sKV_layout: key/value layout + :type sKV_layout: cute.ComposedLayout + :param sO_layout: output layout + :type sO_layout: cute.ComposedLayout + :param gmem_tiled_copy_QKV: tiled copy for QKV load + :type gmem_tiled_copy_QKV: cute.TiledCopy + :param gmem_tiled_copy_O: tiled copy for O store + :type gmem_tiled_copy_O: cute.TiledCopy + :param tiled_mma: tiled mma + :type tiled_mma: cute.TiledMma + :param SharedStorage: shared storage + :type SharedStorage: cutlass.Constexpr + """ + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + m_block, batch_size, num_head = cute.arch.block_idx() + + n_block_max = cute.ceil_div(mK.shape[1], self._n_block_size) + if self._is_causal: + n_block_max = min( + cute.ceil_div( + (m_block + 1) * self._m_block_size, + self._n_block_size, + ), + n_block_max, + ) + n_block = n_block_max - 1 + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + # (m_block_size, head_dim) + gQ = cute.local_tile( + mQ[batch_size, None, num_head, None], + (self._m_block_size, self._head_dim_padded), + (m_block, 0), + ) + # (n_block_size, head_dim, n_block) + gK = cute.local_tile( + mK[batch_size, None, num_head, None], + (self._n_block_size, self._head_dim_padded), + (None, 0), + ) + # (n_block_size, head_dim, n_block) + gV = cute.local_tile( + mV[batch_size, None, num_head, None], + (self._n_block_size, self._head_dim_padded), + (None, 0), + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + + storage = smem.allocate(SharedStorage) + sQ = storage.sQ.get_tensor(sQ_layout) + sK = storage.sK.get_tensor(sKV_layout) + sV = storage.sV.get_tensor(sKV_layout) + + # Transpose view of V to tensor with layout (head_dim, n_block_size) for tiled mma + sVt = cute.composition( + sV, + cute.make_layout( + (self._head_dim_padded, self._n_block_size), + stride=(self._n_block_size, 1), + ), + ) + + gmem_thr_copy_QKV = gmem_tiled_copy_QKV.get_slice(tidx) + # (CPY_Atom, CPY_M, CPY_K) + tQgQ = gmem_thr_copy_QKV.partition_S(gQ) + tQsQ = gmem_thr_copy_QKV.partition_D(sQ) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tKgK = gmem_thr_copy_QKV.partition_S(gK) + tKsK = gmem_thr_copy_QKV.partition_D(sK) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tVgV = gmem_thr_copy_QKV.partition_S(gV) + tVsV = gmem_thr_copy_QKV.partition_D(sV) + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma = tiled_mma.get_slice(tidx) + tSrQ = thr_mma.make_fragment_A(thr_mma.partition_A(sQ)) + tSrK = thr_mma.make_fragment_B(thr_mma.partition_B(sK)) + tOrVt = thr_mma.make_fragment_B(thr_mma.partition_B(sVt)) + acc_shape_O = thr_mma.partition_shape_C( + (self._m_block_size, self._head_dim_padded) + ) + acc_O = cute.make_rmem_tensor(acc_shape_O, cutlass.Float32) + acc_O.fill(0.0) + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom_Q = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + self._dtype, + ) + smem_copy_atom_K = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + self._dtype, + ) + smem_copy_atom_V = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + self._dtype, + ) + smem_tiled_copy_Q = cute.make_tiled_copy_A(smem_copy_atom_Q, tiled_mma) + smem_tiled_copy_K = cute.make_tiled_copy_B(smem_copy_atom_K, tiled_mma) + smem_tiled_copy_V = cute.make_tiled_copy_B(smem_copy_atom_V, tiled_mma) + + smem_thr_copy_Q = smem_tiled_copy_Q.get_slice(tidx) + smem_thr_copy_K = smem_tiled_copy_K.get_slice(tidx) + smem_thr_copy_V = smem_tiled_copy_V.get_slice(tidx) + + tSsQ = smem_thr_copy_Q.partition_S(sQ) + tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ) + tSsK = smem_thr_copy_K.partition_S(sK) + tSrK_copy_view = smem_thr_copy_K.retile(tSrK) + tOsVt = smem_thr_copy_V.partition_S(sVt) + tOrVt_copy_view = smem_thr_copy_V.retile(tOrVt) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + # Construct identity layout for Q and KV + mcQ = cute.make_identity_tensor(mQ.layout.shape) + mcKV = cute.make_identity_tensor(mK.layout.shape) + cQ = cute.local_tile( + mcQ[batch_size, None, num_head, None], + (self._m_block_size, self._head_dim_padded), + (m_block, 0), + ) + cKV = cute.local_tile( + mcKV[batch_size, None, num_head, None], + (self._n_block_size, self._head_dim_padded), + (n_block, 0), + ) + + # Repeat the partitioning with identity layouts + tQcQ = gmem_thr_copy_QKV.partition_S(cQ) + tKVcKV = gmem_thr_copy_QKV.partition_S(cKV) + # Allocate predicate tensors for m and n, here we only allocate the tile of k, and do special process for mn. + # This is to reduce register pressure and gets 2-3% performance gain compared with allocating the whole tile. + tQpQ = cute.make_rmem_tensor( + cute.make_layout( + ( + tQsQ.shape[0][1], + cute.size(tQsQ, mode=[1]), + cute.size(tQsQ, mode=[2]), + ), + stride=(cute.size(tQsQ, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + tKVpKV = cute.make_rmem_tensor( + cute.make_layout( + ( + tKsK.shape[0][1], + cute.size(tKsK, mode=[1]), + cute.size(tKsK, mode=[2]), + ), + stride=(cute.size(tKsK, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + # Set predicates for head_dim bounds, seqlen_q/k bounds is processed at the first tile. + for rest_v in cutlass.range_constexpr(tQpQ.shape[0]): + for rest_k in cutlass.range_constexpr(tQpQ.shape[2]): + tQpQ[rest_v, 0, rest_k] = cute.elem_less( + tQcQ[(0, rest_v), 0, rest_k][3], mQ.layout.shape[3] + ) + for rest_v in cutlass.range_constexpr(tKVpKV.shape[0]): + for rest_k in cutlass.range_constexpr(tKVpKV.shape[2]): + tKVpKV[rest_v, 0, rest_k] = cute.elem_less( + tKVcKV[(0, rest_v), 0, rest_k][3], mK.layout.shape[3] + ) + # /////////////////////////////////////////////////////////////////////////////// + # Prefetch Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Start async loads of the last mn-tile, where we take care of the mn residue + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + if cute.elem_less(tQcQ[0, m, 0][1], mQ.layout.shape[1]): + cute.copy( + gmem_tiled_copy_QKV, + tQgQ[None, m, None], + tQsQ[None, m, None], + pred=tQpQ[None, m, None], + ) + else: + # Clear the smem tiles to account for predicated off loads + tQsQ[None, m, None].fill(0) + for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])): + if cute.elem_less(tKVcKV[0, n, 0][1], mK.layout.shape[1]): + cute.copy( + gmem_tiled_copy_QKV, + tKgK[None, n, None, n_block], + tKsK[None, n, None], + pred=tKVpKV[None, n, None], + ) + else: + # Clear the smem tiles to account for predicated off loads + tKsK[None, n, None].fill(0) + + cute.arch.cp_async_commit_group() + # /////////////////////////////////////////////////////////////////////////////// + # Softmax intermediate result: row_max and row_sum + # /////////////////////////////////////////////////////////////////////////////// + # shape: (atom_v_m * rest_m) + row_max = cute.make_rmem_tensor( + (acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32 + ) + # shape: (atom_v_m * rest_m) + row_sum = cute.make_rmem_tensor( + (acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32 + ) + row_max.fill(-cutlass.Float32.inf) + row_sum.fill(0.0) + + # group parameters for compute_one_n_block + basic_params = SimpleNamespace( + m_block=m_block, + n_block=n_block, + mQ=mQ, + mK=mK, + batch_size=batch_size, + num_head=num_head, + ) + mma_params = SimpleNamespace( + thr_mma=thr_mma, + tiled_mma=tiled_mma, + tSrQ=tSrQ, + tSrK=tSrK, + tOrVt=tOrVt, + acc_O=acc_O, + ) + gmem_copy_params = SimpleNamespace( + gmem_tiled_copy_QKV=gmem_tiled_copy_QKV, + tKVcKV=tKVcKV, + tKgK=tKgK, + tKsK=tKsK, + tVgV=tVgV, + tVsV=tVsV, + tKVpKV=tKVpKV, + ) + smem_copy_params = SimpleNamespace( + smem_tiled_copy_Q=smem_tiled_copy_Q, + smem_tiled_copy_K=smem_tiled_copy_K, + smem_tiled_copy_V=smem_tiled_copy_V, + tSsQ=tSsQ, + tSrQ_copy_view=tSrQ_copy_view, + tSsK=tSsK, + tSrK_copy_view=tSrK_copy_view, + tOsVt=tOsVt, + tOrVt_copy_view=tOrVt_copy_view, + ) + softmax_params = SimpleNamespace( + row_max=row_max, + row_sum=row_sum, + softmax_scale_log2=softmax_scale_log2, + ) + + # Start processing of the first n-block. + # For performance reason, we separate out two kinds of iterations: + # those that need masking on S, and those that don't. + # We need masking on S for the very last block when K and V has length not multiple of n_block_size. + # We also need masking on S if it's causal, for the last ceil_div(m_block_size, n_block_size) blocks. + # We will have at least 1 "masking" iteration. + mask_steps = 1 + if cutlass.const_expr(self._is_causal): + mask_steps = cute.ceil_div(self._m_block_size, self._n_block_size) + + for n_tile in cutlass.range_constexpr(mask_steps): + n_block = n_block_max - n_tile - 1 + basic_params.n_block = n_block + if cutlass.const_expr(self._is_causal): + if n_block >= 0: + self.compute_one_n_block( + basic_params, + mma_params, + gmem_copy_params, + smem_copy_params, + softmax_params, + is_first_n_block=(n_tile == 0), + in_mask_steps=True, + ) + else: + self.compute_one_n_block( + basic_params, + mma_params, + gmem_copy_params, + smem_copy_params, + softmax_params, + is_first_n_block=True, + in_mask_steps=True, + ) + + # Start async loads of rest k-tiles in reverse order, no k-residue handling needed + for n_tile in range(mask_steps, n_block_max, 1): + n_block = n_block_max - n_tile - 1 + basic_params.n_block = n_block + self.compute_one_n_block( + basic_params, + mma_params, + gmem_copy_params, + smem_copy_params, + softmax_params, + is_first_n_block=False, + in_mask_steps=False, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + # normalize acc_O by row_sum and calculate the lse + self.normalize_softmax(acc_O, row_sum) + # store acc_O + rO = cute.make_fragment_like(acc_O, self._dtype) + rO.store(acc_O.load().to(self._dtype)) + # reuse sQ's data iterator + sO = cute.make_tensor(sQ.iterator, sO_layout) + + # smem copy atom for O + smem_copy_atom_O = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self._dtype + ) + # tiled copy atom for O + smem_tiled_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma) + smem_thr_copy_O = smem_tiled_copy_O.get_slice(tidx) + taccOrO = smem_thr_copy_O.retile(rO) + taccOsO = smem_thr_copy_O.partition_D(sO) + # copy acc O from rmem to smem with the smem copy atom + cute.copy( + smem_copy_atom_O, + taccOrO, + taccOsO, + ) + gO = cute.local_tile( + mO[batch_size, None, num_head, None], + (self._m_block_size, self._head_dim_padded), + (m_block, 0), + ) + + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO) + tOgO = gmem_thr_copy_O.partition_D(gO) + tOrO = cute.make_fragment_like(tOgO, self._dtype) + # sync before all smem stores are done. + self.cta_sync_barrier.arrive_and_wait() + # load acc O from smem to rmem for wider vectorization + cute.copy( + gmem_tiled_copy_O, + tOsO, + tOrO, + ) + mcO = cute.make_identity_tensor(mO.layout.shape) + cO = cute.local_tile( + mcO[batch_size, None, num_head, None], + (self._m_block_size, self._head_dim_padded), + (m_block, 0), + ) + tOcO = gmem_thr_copy_O.partition_D(cO) + tOpO = cute.make_rmem_tensor( + cute.make_layout( + (tOgO.shape[0][1], tOgO.shape[1], tOgO.shape[2]), + stride=(tOgO.shape[2], 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tOpO.shape[0]): + for rest_n in cutlass.range_constexpr(cute.size(tOpO.shape[2])): + tOpO[rest_v, 0, rest_n] = cute.elem_less( + tOcO[(0, rest_v), 0, rest_n][3], mO.layout.shape[3] + ) + # copy acc O from rmem to gmem + for rest_m in cutlass.range_constexpr(cute.size(tOpO.shape[1])): + if cute.elem_less(tOcO[0, rest_m, 0][1], mO.layout.shape[1]): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None], + pred=tOpO[None, rest_m, None], + ) + + @cute.jit + def compute_one_n_block( + self, + basic_params: SimpleNamespace, + mma_params: SimpleNamespace, + gmem_copy_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax_params: SimpleNamespace, + is_first_n_block: cutlass.Constexpr, + in_mask_steps: cutlass.Constexpr, + ): + """Compute one n_block of S/O. + + This function provides different variants for processing the first n block versus subsequent blocks, + as well as variants for handling masked and unmasked steps. + + :param basic_params: basic parameters + :type basic_params: SimpleNamespace + :param mma_params: mma parameters + :type mma_params: SimpleNamespace + :param gmem_copy_params: gmem copy parameters + :type gmem_copy_params: SimpleNamespace + :param smem_copy_params: smem copy parameters + :type smem_copy_params: SimpleNamespace + :param softmax_params: softmax parameters + :type softmax_params: SimpleNamespace + :param is_first_n_block: is first n block + :type is_first_n_block: cutlass.Constexpr + """ + acc_shape_S = mma_params.thr_mma.partition_shape_C( + (self._m_block_size, self._n_block_size) + ) + acc_S = cute.make_rmem_tensor(acc_shape_S, cutlass.Float32) + acc_S.fill(0.0) + + # wait for smem tile QK before mma calculation for S + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + # load smem tile V for O, special process for the first tile to avoid loading nan. + # The `if` here is a constexpr, won't be generated in the IR. + if is_first_n_block: + for n in cutlass.range_constexpr(cute.size(gmem_copy_params.tVsV.shape[1])): + if cute.elem_less( + gmem_copy_params.tKVcKV[0, n, 0][1], + basic_params.mK.layout.shape[1], + ): + cute.copy( + gmem_copy_params.gmem_tiled_copy_QKV, + gmem_copy_params.tVgV[None, n, None, basic_params.n_block], + gmem_copy_params.tVsV[None, n, None], + pred=gmem_copy_params.tKVpKV[None, n, None], + ) + else: + gmem_copy_params.tVsV[None, n, None].fill(0.0) + else: + cute.copy( + gmem_copy_params.gmem_tiled_copy_QKV, + gmem_copy_params.tVgV[None, None, None, basic_params.n_block], + gmem_copy_params.tVsV, + pred=gmem_copy_params.tKVpKV, + ) + + cute.arch.cp_async_commit_group() + # /////////////////////////////////////////////////////////////////////////////// + # S gemm calculation + # /////////////////////////////////////////////////////////////////////////////// + # load first QK k-block from smem to rmem for mma + cute.copy( + smem_copy_params.smem_tiled_copy_Q, + smem_copy_params.tSsQ[None, None, 0], + smem_copy_params.tSrQ_copy_view[None, None, 0], + ) + cute.copy( + smem_copy_params.smem_tiled_copy_K, + smem_copy_params.tSsK[None, None, 0], + smem_copy_params.tSrK_copy_view[None, None, 0], + ) + # mma for S + for k in cutlass.range_constexpr(cute.size(smem_copy_params.tSsQ.shape[2])): + # load next QK k-block from smem to rmem for mma + k_next = (k + 1) % cute.size(smem_copy_params.tSsQ.shape[2]) + cute.copy( + smem_copy_params.smem_tiled_copy_Q, + smem_copy_params.tSsQ[None, None, k_next], + smem_copy_params.tSrQ_copy_view[None, None, k_next], + ) + cute.copy( + smem_copy_params.smem_tiled_copy_K, + smem_copy_params.tSsK[None, None, k_next], + smem_copy_params.tSrK_copy_view[None, None, k_next], + ) + cute.gemm( + mma_params.tiled_mma, + acc_S, + mma_params.tSrQ[None, None, k], + mma_params.tSrK[None, None, k], + acc_S, + ) + + # wait for smem tile V for O + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + + if basic_params.n_block > 0: + cute.copy( + gmem_copy_params.gmem_tiled_copy_QKV, + gmem_copy_params.tKgK[None, None, None, basic_params.n_block - 1], + gmem_copy_params.tKsK, + pred=gmem_copy_params.tKVpKV, + ) + cute.arch.cp_async_commit_group() + # /////////////////////////////////////////////////////////////////////////////// + # online softmax + # /////////////////////////////////////////////////////////////////////////////// + self.softmax_rescale_O( + basic_params, + mma_params, + softmax_params, + acc_S, + is_first_n_block, + in_mask_steps, + ) + + rP = cute.make_fragment_like(acc_S, self._dtype) + rP.store(acc_S.load().to(self._dtype)) + # /////////////////////////////////////////////////////////////////////////////// + # O gemm calculation + # /////////////////////////////////////////////////////////////////////////////// + # Convert layout of acc_S to gemm O accept layout. + # Due to the mma instruction shape is 16x8x16, we need to convert from (4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2)) + rP_layout_divided = cute.logical_divide(rP.layout, (None, None, 2)) + rP_mma_view = cute.make_layout( + ( + (rP_layout_divided.shape[0], rP_layout_divided.shape[2][0]), + rP_layout_divided.shape[1], + rP_layout_divided.shape[2][1], + ), + stride=( + (rP_layout_divided.stride[0], rP_layout_divided.stride[2][0]), + rP_layout_divided.stride[1], + rP_layout_divided.stride[2][1], + ), + ) + tOrS = cute.make_tensor(rP.iterator, rP_mma_view) + + # load first V k-block from smem to rmem for mma + cute.copy( + smem_copy_params.smem_tiled_copy_V, + smem_copy_params.tOsVt[None, None, 0], + smem_copy_params.tOrVt_copy_view[None, None, 0], + ) + # mma for O + for k in cutlass.range_constexpr(cute.size(tOrS.shape[2])): + # load next V k-block from smem to rmem for mma + k_next = (k + 1) % cute.size(tOrS.shape[2]) + cute.copy( + smem_copy_params.smem_tiled_copy_V, + smem_copy_params.tOsVt[None, None, k_next], + smem_copy_params.tOrVt_copy_view[None, None, k_next], + ) + cute.gemm( + mma_params.tiled_mma, + mma_params.acc_O, + tOrS[None, None, k], + mma_params.tOrVt[None, None, k], + mma_params.acc_O, + ) + + @cute.jit + def softmax_rescale_O( + self, + basic_params: SimpleNamespace, + mma_params: SimpleNamespace, + softmax_params: SimpleNamespace, + acc_S: cute.Tensor, + is_first_n_block: cutlass.Constexpr, + in_mask_steps: cutlass.Constexpr, + ): + """Apply online softmax and rescale acc_O. + + This function provides different variants for processing the first n block versus subsequent blocks, + as well as variants for handling masked and unmasked steps. + + :param basic_params: basic parameters + :type basic_params: SimpleNamespace + :param mma_params: mma parameters + :type mma_params: SimpleNamespace + :param softmax_params: softmax parameters + :type softmax_params: SimpleNamespace + :param acc_S: acc_S tensor + :type acc_S: cute.Tensor + :param is_first_n_block: is first n_block + :type is_first_n_block: cutlass.Constexpr + :param in_mask_steps: in mask steps + :type in_mask_steps: cutlass.Constexpr + """ + # Change acc_S to M,N layout view. + acc_S_mn = self._make_acc_tensor_mn_view(acc_S) + acc_O_mn = self._make_acc_tensor_mn_view(mma_params.acc_O) + row_max_prev = None + # if it is not the first tile, load the row r of previous row_max and compare with row_max_cur_row. + if cutlass.const_expr(not is_first_n_block): + row_max_prev = cute.make_fragment_like( + softmax_params.row_max, cutlass.Float32 + ) + cute.basic_copy(softmax_params.row_max, row_max_prev) + # if it is the first tile, create a mask for residual of S to -inf for softmax. + tScS_mn = None + if cutlass.const_expr(in_mask_steps): + mcS = cute.make_identity_tensor( + ( + basic_params.mQ.shape[0], + basic_params.mQ.shape[1], + basic_params.mQ.shape[2], + basic_params.mK.shape[1], + ) + ) + cS = cute.local_tile( + mcS[basic_params.batch_size, None, basic_params.num_head, None], + (self._m_block_size, self._n_block_size), + (basic_params.m_block, basic_params.n_block), + ) + tScS = mma_params.thr_mma.partition_C(cS) + tScS_mn = self._make_acc_tensor_mn_view(tScS) + + # Each iteration processes one row of acc_S + for r in cutlass.range_constexpr(cute.size(softmax_params.row_max)): + # mask residual of S with -inf + if cutlass.const_expr(in_mask_steps): + if cutlass.const_expr(not self._is_causal): + # traverse column index. + for c in cutlass.range_constexpr(cute.size(tScS_mn.shape[1])): + if cute.elem_less( + basic_params.mK.shape[1], tScS_mn[0, c][3] + 1 + ): + acc_S_mn[r, c] = -cutlass.Float32.inf + else: + # get the column index limit based on current row. Only consider the row index, so the column index sets to 0. + col_idx_limit = cutlass.min( + tScS_mn[r, 0][1] + 1, basic_params.mK.shape[1] + ) + # traverse column index. + for c in cutlass.range_constexpr(cute.size(tScS_mn.shape[1])): + # only consider the column index, so the row index sets to 0. + if cute.elem_less(col_idx_limit, tScS_mn[0, c][3] + 1): + acc_S_mn[r, c] = -cutlass.Float32.inf + + # (n_block_size) + acc_S_row = acc_S_mn[r, None].load() + # row_max_cur_row => f32 + row_max_cur_row = acc_S_row.reduce( + cute.ReductionOp.MAX, -cutlass.Float32.inf, 0 + ) + # quad reduction for row_max + row_max_cur_row = self._threadquad_reduce_max(row_max_cur_row) + row_max_prev_row = None + # if it is not the first tile, load the row r of previous row_max and compare with row_max_cur_row. + if cutlass.const_expr(not is_first_n_block): + row_max_prev_row = row_max_prev[r] + row_max_cur_row = cute.arch.fmax(row_max_prev_row, row_max_cur_row) + if cutlass.const_expr(self._is_causal): + row_max_cur_row = ( + 0.0 if row_max_cur_row == -cutlass.Float32.inf else row_max_cur_row + ) + + # compute exp(x - max) using exp2(x * log_2(e) - max * log_2(e)) + acc_S_row_exp = cute.math.exp2( + acc_S_row * softmax_params.softmax_scale_log2 + - row_max_cur_row * softmax_params.softmax_scale_log2, + fastmath=True, + ) + # acc_S_row_sum => f32 + acc_S_row_sum = acc_S_row_exp.reduce( + cute.ReductionOp.ADD, cutlass.Float32.zero, 0 + ) + # if it is not the first tile, load the row r of previous row_max and minus row_max_cur_row to update row_sum. + if cutlass.const_expr(not is_first_n_block): + prev_minus_cur_exp = cute.math.exp2( + row_max_prev_row * softmax_params.softmax_scale_log2 + - row_max_cur_row * softmax_params.softmax_scale_log2, + fastmath=True, + ) + acc_S_row_sum = ( + acc_S_row_sum + softmax_params.row_sum[r] * prev_minus_cur_exp + ) + acc_O_mn[r, None] = acc_O_mn[r, None].load() * prev_minus_cur_exp + # update row_max, row_sum and acc_S + softmax_params.row_max[r] = row_max_cur_row + softmax_params.row_sum[r] = acc_S_row_sum + acc_S_mn[r, None] = acc_S_row_exp + + @cute.jit + def normalize_softmax( + self, + acc_O: cute.Tensor, + row_sum: cute.Tensor, + ): + """Normalize acc_O by row_sum. + + :param acc_O: input tensor + :type acc_O: cute.Tensor + :param row_sum: row_sum tensor + :type row_sum: cute.Tensor + """ + # do quad reduction for row_sum. + acc_O_mn = self._make_acc_tensor_mn_view(acc_O) + for r in cutlass.range_constexpr(cute.size(row_sum)): + row_sum[r] = self._threadquad_reduce_sum(row_sum[r]) + # if row_sum is zero or nan, set acc_O_mn_row to 1.0 + acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r] + + scale = ( + 1.0 if acc_O_mn_row_is_zero_or_nan else cute.arch.rcp_approx(row_sum[r]) + ) + + acc_O_mn[r, None] = acc_O_mn[r, None].load() * scale + + def _make_acc_tensor_mn_view(self, acc: cute.Tensor) -> cute.Tensor: + """make acc tensor as mn layout view + + :param acc: input tensor + :type acc: cute.Tensor + :return: acc tensor mn layout view + :rtype: cute.Tensor + """ + acc_layout_col_major = cute.make_layout(acc.layout.shape) + acc_layout_mn = cute.make_layout( + ( + ( + acc_layout_col_major.shape[0][1], + acc_layout_col_major.shape[1], + ), # MMA_M + ( + acc_layout_col_major.shape[0][0], + acc_layout_col_major.shape[2], + ), # MMA_N + ), + stride=( + ( + acc_layout_col_major.stride[0][1], + acc_layout_col_major.stride[1], + ), # MMA_M + ( + acc_layout_col_major.stride[0][0], + acc_layout_col_major.stride[2], + ), # MMA_N + ), + ) + acc_layout_mn = cute.composition(acc.layout, acc_layout_mn) + return cute.make_tensor(acc.iterator, acc_layout_mn) + + def _threadquad_reduce(self, val: cutlass.Float32, op: Callable) -> cutlass.Float32: + """thread quad reduction + + :param val: register value + :type val: cutlass.Float32 + :param op: binary operator + :type op: Callable + :return: reduced value + :rtype: cutlass.Float32 + """ + val = op( + val, + cute.arch.shuffle_sync_bfly(val, offset=2, mask=-1, mask_and_clamp=31), + ) + val = op( + val, + cute.arch.shuffle_sync_bfly(val, offset=1, mask=-1, mask_and_clamp=31), + ) + return val + + def _threadquad_reduce_max(self, val: cutlass.Float32) -> cutlass.Float32: + """thread quad reduction max + + :param val: register value + :type val: cutlass.Float32 + :return: max value + :rtype: cutlass.Float32 + """ + return self._threadquad_reduce(val, lambda x, y: cute.arch.fmax(x, y)) + + def _threadquad_reduce_sum(self, val: cutlass.Float32) -> cutlass.Float32: + """thread quad reduction sum + + :param val: register value + :type val: cutlass.Float32 + :return: sum value + :rtype: cutlass.Float32 + """ + return self._threadquad_reduce(val, lambda x, y: x + y) + + +def run( + dtype: Type[cutlass.Numeric], + batch_size: int, + seqlen_q: int, + seqlen_k: int, + num_head: int, + head_dim: int, + softmax_scale: float = 1.0, + m_block_size: int = 128, + n_block_size: int = 128, + num_threads: int = 128, + is_causal: bool = False, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + import torch + import cutlass.torch as cutlass_torch + + # Skip unsupported testcase + if not FlashAttentionForwardAmpere.can_implement( + dtype, + head_dim, + m_block_size, + n_block_size, + num_threads, + is_causal, + ): + raise TypeError( + f"Unsupported testcase {dtype}, {head_dim}, {m_block_size}, {n_block_size}, {num_threads}, {is_causal}" + ) + + print("Running Ampere SM80 FlashAttentionForward test with:") + print(f" dtype: {dtype}") + print(f" batch_size: {batch_size}") + print(f" seqlen_q: {seqlen_q}") + print(f" seqlen_k: {seqlen_k}") + print(f" num_head: {num_head}") + print(f" head_dim: {head_dim}") + print(f" softmax_scale: {softmax_scale}") + print(f" m_block_size: {m_block_size}") + print(f" n_block_size: {n_block_size}") + print(f" num_threads: {num_threads}") + print(f" is_causal: {is_causal}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + # Create tensor Q/K/V/O + def create_tensor( + batch_size: int, + seqlen: int, + num_head: int, + head_dim: int, + dtype: Type[cutlass.Numeric], + ) -> cute.Tensor: + # (batch_size, seqlen, num_head, head_dim) + shape = (batch_size, seqlen, num_head, head_dim) + torch_tensor = ( + torch.empty(*shape, dtype=torch.int32) + .random_(-2, 2) + .to(dtype=cutlass_torch.dtype(dtype)) + .cuda() + ) + # assume input is 16B aligned. + cute_tensor = ( + from_dlpack(torch_tensor, assumed_align=16) + .mark_layout_dynamic(leading_dim=3) + .mark_compact_shape_dynamic( + mode=3, + stride_order=torch_tensor.dim_order(), + divisibility=(128 // dtype.width), + ) + ) + return cute_tensor, torch_tensor + + q, q_torch = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype) + k, k_torch = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype) + v, v_torch = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype) + o, o_torch = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype) + + fa2_fwd = FlashAttentionForwardAmpere( + head_dim, + m_block_size, + n_block_size, + num_threads, + is_causal, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + # Pass compile options if needed + compile_options = "" + # compile the fa2 forward pass + compiled_fa2_fwd = cute.compile( + fa2_fwd, q, k, v, o, softmax_scale, current_stream, options=compile_options + ) + + if not skip_ref_check: + compiled_fa2_fwd(q, k, v, o, softmax_scale, current_stream) + torch.cuda.synchronize() + q_ref = q_torch.permute(0, 2, 1, 3) + k_ref = k_torch.permute(0, 2, 1, 3) + v_ref = v_torch.permute(0, 2, 1, 3) + torch.backends.cuda.enable_flash_sdp(enabled=True) + ref_o = torch.nn.functional.scaled_dot_product_attention( + q_ref, k_ref, v_ref, scale=softmax_scale, is_causal=is_causal + ).permute(0, 2, 1, 3) + torch.testing.assert_close(o_torch.cpu(), ref_o.cpu(), atol=1e-02, rtol=1e-04) + print("Results verified successfully!") + + def generate_tensors(): + q_workspace, _ = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype) + k_workspace, _ = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype) + v_workspace, _ = create_tensor(batch_size, seqlen_k, num_head, head_dim, dtype) + o_workspace, _ = create_tensor(batch_size, seqlen_q, num_head, head_dim, dtype) + return testing.JitArguments( + q_workspace, + k_workspace, + v_workspace, + o_workspace, + softmax_scale, + current_stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_torch.numel() * q_torch.element_size() + + k_torch.numel() * k_torch.element_size() + + v_torch.numel() * v_torch.element_size() + + o_torch.numel() * o_torch.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + avg_time_us = testing.benchmark( + compiled_fa2_fwd, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="example of flash attention v2 with CuTe on GPU" + ) + parser.add_argument("--dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--seqlen_q", type=int, default=8192) + parser.add_argument("--seqlen_k", type=int, default=8192) + parser.add_argument("--num_head", type=int, default=16) + parser.add_argument("--head_dim", type=int, default=128) + parser.add_argument("--softmax_scale", type=float, default=0.5) + parser.add_argument("--m_block_size", type=int, default=128) + parser.add_argument("--n_block_size", type=int, default=64) + parser.add_argument("--num_threads", type=int, default=128) + parser.add_argument("--is_causal", action="store_true", help="Enable causal mask") + parser.add_argument("--warmup_iterations", type=int, default=3) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference check" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + args = parser.parse_args() + run( + args.dtype, + args.batch_size, + args.seqlen_q, + args.seqlen_k, + args.num_head, + args.head_dim, + args.softmax_scale, + args.m_block_size, + args.n_block_size, + args.num_threads, + args.is_causal, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/launch.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/launch.py new file mode 100644 index 0000000000..5c73c27bba --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/launch.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the unchanged public CuTe kernel with the Rust example's exact inputs. + +The fixture and scalar oracle require only the standard library. GPU imports +occur only in run(); the large CPU oracle uses PyTorch FP64 on the CPU. +""" + +import argparse +import json +import math +from pathlib import Path +import statistics +import struct + +HEAD_DIM = 64 +STATIC_SHAPE = (1024, 1024, 8) # seq_q, seq_k, heads; batch=1 and head_dim=64 +SOFTMAX_SCALE = 0.125 +ATOL = 2e-4 +RTOL = 2e-3 +Q_SEED, K_SEED, V_SEED = 0xA511E9B3, 0x63D83595, 0xC2B2AE35 + + +def mix32(value): + value ^= value >> 16 + value = (value * 0x7FEB352D) & 0xFFFFFFFF + value ^= value >> 15 + value = (value * 0x846CA68B) & 0xFFFFFFFF + return value ^ (value >> 16) + + +def make_inputs(seq_q, seq_k, heads, input_scale=1, zero_query=False): + """Return exact FP16-representable values in flattened BSHD order.""" + def values(length, seed, scale): + return [((mix32(i ^ seed) & 63) - 32) * scale / 128.0 for i in range(length)] + + q = ([0.0] * (seq_q * heads * HEAD_DIM) if zero_query else + values(seq_q * heads * HEAD_DIM, Q_SEED, input_scale)) + return (q, values(seq_k * heads * HEAD_DIM, K_SEED, input_scale), + values(seq_k * heads * HEAD_DIM, V_SEED, 1)) + + +def encode_f16(values): + return struct.pack(f"<{len(values)}e", *values) + + +def decode_f16(data): + if len(data) % 2: + raise ValueError("FP16 input must contain complete two-byte elements") + return list(struct.unpack(f"<{len(data) // 2}e", data)) + + +def bits_hash(values): + result = 0xCBF29CE484222325 + for byte in encode_f16(values): + result = ((result ^ byte) * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF + return f"{result:016x}" + + +def reference_cpu(q, k, v, seq_q, seq_k, heads, causal, *, use_torch=False): + """Stable mathematical attention with FP64 probabilities and accumulation.""" + if len(q) != seq_q * heads * HEAD_DIM or len(k) != seq_k * heads * HEAD_DIM or len(v) != len(k): + raise ValueError("fixture dimensions do not match tensor lengths") + if use_torch: + import torch + + q_cpu = torch.tensor(q, dtype=torch.float64, device="cpu").reshape(seq_q, heads, HEAD_DIM).permute(1, 0, 2) + k_cpu = torch.tensor(k, dtype=torch.float64, device="cpu").reshape(seq_k, heads, HEAD_DIM).permute(1, 0, 2) + v_cpu = torch.tensor(v, dtype=torch.float64, device="cpu").reshape(seq_k, heads, HEAD_DIM).permute(1, 0, 2) + scores = (q_cpu @ k_cpu.transpose(1, 2)) * SOFTMAX_SCALE + if causal: + allowed = torch.arange(seq_k)[None, :] <= torch.arange(seq_q)[:, None] + scores.masked_fill_(~allowed, float("-inf")) + probabilities = (scores - scores.amax(dim=-1, keepdim=True)).exp() + probabilities /= probabilities.sum(dim=-1, keepdim=True) + return (probabilities @ v_cpu).permute(1, 0, 2).contiguous().reshape(-1).tolist() + + output = [0.0] * len(q) + for head in range(heads): + for row in range(seq_q): + q_base = (row * heads + head) * HEAD_DIM + end = min(row + 1, seq_k) if causal else seq_k + scores = [] + for column in range(end): + k_base = (column * heads + head) * HEAD_DIM + scores.append(sum(q[q_base + d] * k[k_base + d] for d in range(HEAD_DIM)) * SOFTMAX_SCALE) + maximum = max(scores) + weights = [math.exp(score - maximum) for score in scores] + denominator = sum(weights) + for d in range(HEAD_DIM): + output[q_base + d] = sum( + weight * v[(column * heads + head) * HEAD_DIM + d] + for column, weight in enumerate(weights) + ) / denominator + return output + + +def verify(got, expected): + if len(got) != len(expected) or not got: + raise ValueError("output and nonempty oracle lengths must match") + maximum, ratio = 0.0, 0.0 + for index, (actual, want) in enumerate(zip(got, expected)): + if not math.isfinite(actual) or not math.isfinite(want): + raise ValueError(f"nonfinite output/oracle at {index}: {actual}, {want}") + error = abs(actual - want) + bound = ATOL + RTOL * abs(want) + if error > bound: + raise ValueError(f"output mismatch at {index}: actual={actual}, expected={want}, error={error}, bound={bound}") + maximum = max(maximum, error) + ratio = max(ratio, error / bound) + return {"elements": len(got), "max_abs_error": maximum, "max_tolerance_ratio": ratio, + "atol": ATOL, "rtol": RTOL, "oracle": "cpu-fp64-stable-softmax"} + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--seq-q", type=int, default=1024) + parser.add_argument("--seq-k", type=int, default=1024) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--shape-mode", choices=("runtime", "static"), default="runtime", + help="static specializes shapes and strides for Sq=Sk=1024, H=8 only") + masks = parser.add_mutually_exclusive_group() + masks.add_argument("--causal", action="store_true", dest="causal") + masks.add_argument("--noncausal", action="store_false", dest="causal") + parser.set_defaults(causal=True) + parser.add_argument("--input-scale", type=int, choices=(1, 4, 16), default=1, + help="exact power-of-two scaling of Q/K only") + parser.add_argument("--zero-query", action="store_true") + parser.add_argument("--device", type=int, default=0, help="ordinal after CUDA_VISIBLE_DEVICES remapping") + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--iters", type=int, default=31, help="zero selects verification without timing") + parser.add_argument("--output", type=Path, help="write checked FP16 BSHD output") + parser.add_argument("--json", type=Path) + args = parser.parse_args(argv) + if args.shape_mode == "static" and (args.seq_q, args.seq_k, args.heads) != STATIC_SHAPE: + parser.error("static shape mode requires --seq-q 1024 --seq-k 1024 --heads 8") + if args.seq_q <= 0 or args.seq_k <= 0 or args.heads <= 0: + parser.error("sequence lengths and heads must be positive") + if (args.seq_q != 1 and args.seq_q % 64) or args.seq_k % 64: + parser.error("query length must be 1 or a multiple of 64; KV length must be a multiple of 64") + if args.causal and args.seq_q != args.seq_k: + parser.error("initial causal cases require equal sequence lengths") + if max(args.seq_q, args.seq_k) * args.heads * HEAD_DIM > 0xFFFFFFFF: + parser.error("tensor element count must fit u32 indexing") + if args.warmup < 0 or args.iters < 0 or args.device < 0: + parser.error("device, warmup and iters must be nonnegative") + args.workload = "prefill-fmha" + return args + + +def make_tensor_view(tensor, from_dlpack, shape_mode): + """Keep static shapes and strides, or apply the original runtime annotations.""" + view = from_dlpack(tensor, assumed_align=16) + if shape_mode == "runtime": + view = view.mark_layout_dynamic(leading_dim=3).mark_compact_shape_dynamic( + mode=3, stride_order=tensor.dim_order(), divisibility=8, + ) + return view + + +def run(args): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + import torch + + from flash_attention_v2 import FlashAttentionForwardAmpere + + q_values, k_values, v_values = make_inputs( + args.seq_q, args.seq_k, args.heads, args.input_scale, args.zero_query, + ) + expected = reference_cpu(q_values, k_values, v_values, args.seq_q, args.seq_k, + args.heads, args.causal, use_torch=True) + torch.cuda.set_device(args.device) + device = torch.device("cuda", args.device) + tensors = [torch.tensor(values, dtype=torch.float16, device=device).reshape(1, length, args.heads, HEAD_DIM) + for values, length in ((q_values, args.seq_q), (k_values, args.seq_k), (v_values, args.seq_k))] + output = torch.full((1, args.seq_q, args.heads, HEAD_DIM), float("nan"), dtype=torch.float16, device=device) + tensors.append(output) + views = [make_tensor_view(tensor, from_dlpack, args.shape_mode) for tensor in tensors] + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + reference = FlashAttentionForwardAmpere(HEAD_DIM, 64, 64, 128, args.causal) + compiled = cute.compile(reference, *views, SOFTMAX_SCALE, stream) + + def invoke(): + compiled(*views, SOFTMAX_SCALE, stream) + + def read_output(): + return output.detach().cpu().reshape(-1).tolist() + + invoke() + checked = read_output() + verification = verify(checked, expected) + output.fill_(float("nan")) + invoke() + replay = read_output() + verify(replay, expected) + if encode_f16(replay) != encode_f16(checked): + raise ValueError("output changed across identical launches") + for _ in range(args.warmup if args.iters else 0): + invoke() + torch.cuda.synchronize() + events = [(torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(args.iters)] + samples = [] + for start, end in events: + start.record() + invoke() + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) * 1000.0 + if not math.isfinite(elapsed) or elapsed <= 0: + raise ValueError(f"invalid event timing: {elapsed}") + samples.append(elapsed) + after = read_output() + verify(after, expected) + if encode_f16(after) != encode_f16(checked): + raise ValueError("output changed during benchmark") + if args.output: + args.output.write_bytes(encode_f16(after)) + result = { + "implementation": "cutlass-cute-dsl", "workload": args.workload, + "shape": [1, args.seq_q, args.seq_k, args.heads, HEAD_DIM], + "causal": args.causal, "input_scale": args.input_scale, "zero_query": args.zero_query, + "input_hashes": {name: bits_hash(values) for name, values in zip("qkv", (q_values, k_values, v_values))}, + "output_hash": bits_hash(after), "verification": "cpu-pass", "validation": verification, + "replay_bitwise_equal": True, "metric_value": statistics.median(samples) if samples else None, + "metric_unit": "us", "samples_us": samples, + "provenance": {"gpu": torch.cuda.get_device_name(args.device), + "compute_capability": list(torch.cuda.get_device_capability(args.device)), + "device": args.device, "tool": "cuda-events-single-kernel", "timing_mode": "direct", + "causal_specialization": "static", + "shape_specialization": args.shape_mode, + "warmup": args.warmup if args.iters else 0, "iters": args.iters, + "grid": [(args.seq_q + 63) // 64, 1, args.heads], + "threads": 128, "shared_memory_bytes": 24576, "dtype": "float16", + "cutlass_version": cutlass.__version__, "torch_version": torch.__version__, + "source_commit": "dcf215af68a2d08d305076c152a06f201728cd53"}, + } + text = json.dumps(result, indent=2, allow_nan=False) + "\n" + if args.json: + args.json.write_text(text) + print(text, end="") + return result + + +if __name__ == "__main__": + run(parse_args()) diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/requirements.txt b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/requirements.txt new file mode 100644 index 0000000000..b2cb40f71c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/requirements.txt @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# This comparison uses the shared environment pinned in cute-rs-preview/requirements.txt. +-r ../../../requirements.txt diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_compare.py new file mode 100644 index 0000000000..e6e79b49e0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_compare.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only acceptance tests for matched attention profiling evidence.""" + +import contextlib +import io +import json +import os +from pathlib import Path +import sqlite3 +import tempfile +import time +from types import SimpleNamespace +import unittest +from unittest import mock + +from launch import ATOL, RTOL, bits_hash, encode_f16, make_inputs + +try: + import compare +except ImportError: + compare = None + + +REFERENCE_NAME = "kernel_FlashAttentionForwardAmpere_1_64_64" + + +def trace(path, launches, *, full_names=None): + """Insert out of order; durations are nanoseconds, including auxiliary work.""" + with sqlite3.connect(path) as db: + db.execute("CREATE TABLE StringIds (id INTEGER PRIMARY KEY, value TEXT)") + db.execute("CREATE TABLE CUPTI_ACTIVITY_KIND_KERNEL " + "(start INTEGER, end INTEGER, shortName INTEGER, deviceId INTEGER)") + names = {name: i for i, name in enumerate(dict.fromkeys(name for name, _ in launches))} + db.executemany("INSERT INTO StringIds VALUES (?, ?)", [(i, name) for name, i in names.items()]) + db.executemany("INSERT INTO CUPTI_ACTIVITY_KIND_KERNEL VALUES (?, ?, ?, 0)", + [(i * 1000000, i * 1000000 + duration, names[name]) + for i, (name, duration) in reversed(list(enumerate(launches)))]) + if full_names is not None: + db.execute("ALTER TABLE CUPTI_ACTIVITY_KIND_KERNEL ADD COLUMN demangledName INTEGER") + full_ids = {name: len(names) + i for i, name in enumerate(dict.fromkeys(full_names))} + db.executemany("INSERT INTO StringIds VALUES (?, ?)", [(i, name) for name, i in full_ids.items()]) + db.executemany("UPDATE CUPTI_ACTIVITY_KIND_KERNEL SET demangledName=? WHERE start=?", + [(full_ids[name], i * 1000000) for i, name in enumerate(full_names)]) + + +def options(): + return SimpleNamespace(seq_q=64, seq_k=64, heads=1, causal=True, input_scale=1, + zero_query=False, warmup=2, samples=3, device=0, shape_mode="runtime") + + +class ComparisonTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(compare, "comparison runner has not been implemented") + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + self.args = options() + self.hashes = compare.expected_input_hashes(self.args) + self.output = [0.0] * (64 * 64) + + def report(self): + return { + "implementation": "cute-rs", "workload": "prefill-fmha", + "shape": [1, self.args.seq_q, self.args.seq_k, self.args.heads, 64], + "causal": True, "input_scale": 1, + "zero_query": False, "input_hashes": self.hashes, + "output_hash": bits_hash(self.output), "verification": "cpu-pass", + "validation": {"elements": len(self.output), "max_abs_error": 0.0, + "max_tolerance_ratio": 0.0, "atol": ATOL, "rtol": RTOL, + "oracle": "cpu-fp64-stable-softmax"}, + "replay_bitwise_equal": True, "metric_value": 999.0, "metric_unit": "us", + "samples_us": [999.0] * 3, + "provenance": {"gpu": "test SM120 device", "compute_capability": [12, 0], + "device": 0, "warmup": 2, "iters": 3, "dtype": "float16", + "causal_specialization": "static", + "shape_specialization": self.args.shape_mode, + "tool": "cuda-events-single-kernel", "timing_mode": "direct"}, + } + + def write_evidence(self, report=None, values=None): + result, output = self.directory / "runner.json", self.directory / "output.f16" + result.write_text(json.dumps(self.report() if report is None else report)) + output.write_bytes(encode_f16(self.output if values is None else values)) + return result, output + + def validate(self, paths, started=0): + return compare.validate_report(*paths, self.args, started, self.hashes) + + def test_checks_every_output_and_hash_with_full_cpu_oracle_evidence(self): + started = time.time_ns() + report, values = self.validate(self.write_evidence(), started) + self.assertEqual(report["validation"]["elements"], 4096) + self.assertEqual(values, self.output) + + def test_streamed_input_hashes_match_shared_fixture_for_prefill_and_decode(self): + for query, scale, zero in ((64, 1, False), (1, 16, False), (64, 4, True)): + with self.subTest(query=query, scale=scale, zero=zero): + args = options() + args.seq_q, args.input_scale, args.zero_query = query, scale, zero + inputs = make_inputs(args.seq_q, args.seq_k, args.heads, scale, zero) + self.assertEqual(compare.expected_input_hashes(args), + {name: bits_hash(values) for name, values in zip("qkv", inputs)}) + + def test_rejects_wrong_shape_inputs_or_incomplete_cpu_verification(self): + mutations = [ + ("verification", "failed"), ("replay_bitwise_equal", 1), + ("shape", [True, 64, 64, 1, 64]), ("input_hashes", {"q": "0" * 16}), + ("validation.elements", 1), ("validation.max_tolerance_ratio", 1.01), + ("validation.max_abs_error", float("nan")), ("validation.atol", 1.0), + ("validation.oracle", "gpu-library"), ("provenance.warmup", 1), + ("provenance.iters", 31), ("provenance.device", 1), + ] + for field, value in mutations: + with self.subTest(field=field): + report = self.report() + owner = report + parts = field.split(".") + for part in parts[:-1]: + owner = owner[part] + owner[parts[-1]] = value + with self.assertRaises(ValueError): + self.validate(self.write_evidence(report)) + + def test_rejects_stale_corrupt_truncated_or_nonfinite_output(self): + paths = self.write_evidence() + os.utime(paths[0], ns=(1, 1)) + with self.assertRaisesRegex(ValueError, "stale"): + self.validate(paths, time.time_ns()) + for values in (self.output[:-1], [1.0] + self.output[1:], + [float("nan")] + self.output[1:]): + with self.subTest(length=len(values), first=values[0]): + with self.assertRaises(ValueError): + self.validate(self.write_evidence(values=values)) + + def test_exact_kernel_samples_strip_two_checks_and_all_warmups(self): + path = self.directory / "trace.sqlite" + trace(path, [("torch_fill", 777000), ("fmha_prefill", 900000), + ("torch_fill", 555000), ("fmha_prefill", 800000)] + + [("fmha_prefill", n) for n in [700000, 600000, 31000, 11000, 21000]]) + result = compare.kernel_samples(path, 2, 3, expected_name="fmha_prefill") + self.assertEqual(result["samples_us"], [31.0, 11.0, 21.0]) + self.assertEqual(result["correctness_samples_us"], [900.0, 800.0]) + self.assertEqual(result["warmup_samples_us"], [700.0, 600.0]) + self.assertEqual(result["auxiliary_kernel_counts"], {"torch_fill": 2}) + + def test_resolves_one_reference_name_then_requires_all_133_launches(self): + path = self.directory / "reference.sqlite" + trace(path, [(REFERENCE_NAME, 900000)] * 102 + + [(REFERENCE_NAME, n * 1000) for n in range(1, 32)]) + result = compare.kernel_samples(path, 100, 31, + name_contains="FlashAttentionForwardAmpere") + self.assertEqual(result["kernel_name"], REFERENCE_NAME) + self.assertEqual(result["samples_us"], list(range(1, 32))) + + def test_resolves_truncated_short_name_using_full_demangled_identity(self): + path = self.directory / "truncated.sqlite" + short = "kernel_cutlass_kernel_flash_at_1..." + full = "kernel_cutlass_kernel_flash_attention_v2FlashAttentionForwardAmpere_object_at__tensor_0" + launches = [(short, 900000)] * 102 + [(short, n * 1000) for n in range(1, 32)] + trace(path, launches + [("vectorized_elementwise_kernel", 1000)] * 2, + full_names=[full] * 133 + ["void at::native::vectorized_elementwise_kernel"] * 2) + for selector in ({"name_contains": "FlashAttentionForwardAmpere"}, + {"expected_name": short}, {"expected_name": full}): + with self.subTest(selector=selector): + result = compare.kernel_samples(path, 100, 31, **selector) + self.assertEqual(result["kernel_name"], full) + self.assertEqual(result["kernel_name_field"], "demangledName") + self.assertEqual(result["kernel_short_names"], [short]) + self.assertEqual(result["samples_us"], list(range(1, 32))) + self.assertEqual(sum(result["auxiliary_kernel_counts"].values()), 2) + + def test_rejects_ambiguous_truncated_name_and_keeps_distinct_full_kernels(self): + path = self.directory / "colliding.sqlite" + short = "kernel_cutlass_kernel_flash_at_1..." + first, second = REFERENCE_NAME + "_first", REFERENCE_NAME + "_second" + trace(path, [(short, 1000)] * 7 + [(short, 2000)] * 7, + full_names=[first] * 7 + [second] * 7) + for selector in ({"name_contains": "FlashAttentionForwardAmpere"}, {"expected_name": short}): + with self.subTest(selector=selector): + with self.assertRaisesRegex(ValueError, "one.*kernel name"): + compare.kernel_samples(path, 2, 3, **selector) + result = compare.kernel_samples(path, 2, 3, expected_name=first) + self.assertEqual(result["kernel_name"], first) + self.assertEqual(result["samples_us"], [1.0] * 3) + self.assertEqual(result["auxiliary_kernel_counts"], {second: 7}) + + def test_rejects_wrong_counts_ambiguous_names_and_invalid_durations(self): + cases = { + "missing": [("fmha_prefill", 1000)] * 6, + "extra": [("fmha_prefill", 1000)] * 8, + "partial-name": [("fmha_prefill_other", 1000)] * 7, + "invalid-duration": [("fmha_prefill", 1000)] * 6 + [("fmha_prefill", -1)], + } + for name, launches in cases.items(): + with self.subTest(name=name): + path = self.directory / f"{name}.sqlite" + trace(path, launches) + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3, expected_name="fmha_prefill") + path = self.directory / "ambiguous.sqlite" + trace(path, [(REFERENCE_NAME, 1000)] * 7 + [(REFERENCE_NAME + "_other", 1000)]) + with self.assertRaisesRegex(ValueError, "one.*kernel name"): + compare.kernel_samples(path, 2, 3, name_contains="FlashAttentionForwardAmpere") + + def test_rejects_missing_trace_unresolved_names_or_multiple_devices(self): + path = self.directory / "missing.sqlite" + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3, expected_name="fmha_prefill") + self.assertFalse(path.exists()) + trace(path, [("fmha_prefill", 1000)] * 7) + with sqlite3.connect(path) as db: + db.execute("UPDATE CUPTI_ACTIVITY_KIND_KERNEL SET deviceId=1 WHERE start=0") + with self.assertRaisesRegex(ValueError, "device"): + compare.kernel_samples(path, 2, 3, expected_name="fmha_prefill") + with sqlite3.connect(path) as db: + db.execute("DELETE FROM StringIds") + with self.assertRaises(ValueError): + compare.kernel_samples(path, 2, 3, expected_name="fmha_prefill") + + def test_aggregate_uses_round_medians_and_preserves_every_raw_sample(self): + entries = [{"label": "rust", "samples_us": values, "median_us": median} + for values, median in [([1, 100, 1], 1), ([9, 10, 11], 10), ([3, 2, 4], 3)]] + result = compare.aggregate_rounds(entries)["rust"] + self.assertEqual(result["round_medians_us"], [1, 10, 3]) + self.assertEqual(result["median_of_round_medians_us"], 3) + self.assertEqual(result["samples_us"], [1, 100, 1, 9, 10, 11, 3, 2, 4]) + + def arguments(self): + binary = self.directory / "rust" + binary.write_text("synthetic executable") + binary.chmod(0o755) + return ["--rust-binary", str(binary), "--output-dir", str(self.directory / "evidence"), + "--seq-q", "64", "--seq-k", "64", "--heads", "1", + "--warmup", "2", "--samples", "3", "--rounds", "3"] + + def fake_command(self, command, timeout, environment, log_path): + self.assertEqual(environment["CUDA_VISIBLE_DEVICES"], "GPU-externally-selected") + log_path.write_text("synthetic tool version\n") + if command[1:2] == ["profile"]: + result = Path(command[command.index("--json") + 1]) + prefix = Path(command[command.index("--output") + 1]) + output = Path(command[-1]) + report = self.report() + report["causal"] = "--causal" in command + if "--shape-mode" in command: + report["provenance"]["shape_specialization"] = command[command.index("--shape-mode") + 1] + label = "cutlass" if "cutlass" in prefix.name else "rust" + if label == "cutlass": + report["implementation"] = "cutlass-cute-dsl" + if label == getattr(self, "invalid_specialization_label", None): + if self.invalid_specialization is None: + del report["provenance"]["causal_specialization"] + else: + report["provenance"]["causal_specialization"] = self.invalid_specialization + if label == getattr(self, "invalid_shape_label", None): + if self.invalid_shape_specialization is None: + del report["provenance"]["shape_specialization"] + else: + report["provenance"]["shape_specialization"] = self.invalid_shape_specialization + result.write_text(json.dumps(report)) + output.write_bytes(encode_f16(self.output)) + prefix.with_suffix(".nsys-rep").write_bytes(b"synthetic trace") + elif command[1:2] == ["export"]: + path = Path(command[command.index("--output") + 1]) + name = (REFERENCE_NAME if "cutlass" in path.name + else getattr(self, "rust_kernel_name", "fmha_prefill")) + durations = [900000, 800000, 700000, 600000, 31000, 11000, 21000] + if getattr(self, "drop_last_launch", False): + durations.pop() + trace(path, [(name, n) for n in durations]) + return 0 + + def test_runner_keeps_external_device_and_separates_active_and_event_metrics(self): + arguments = self.arguments() + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(compare.main(arguments), 0) + directory, = (self.directory / "evidence").iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "passed") + self.assertEqual([entry["label"] for entry in result["runs"]], + ["rust", "cutlass", "cutlass", "rust", "rust", "cutlass"]) + self.assertEqual(result["summary"]["rust"]["median_of_round_medians_us"], 21) + self.assertEqual(result["summary"]["rust"]["samples_us"], [31, 11, 21] * 3) + self.assertEqual(result["runs"][0]["runner_report"]["metric_value"], 999) + self.assertIn("source_snapshot", result["provenance"]) + self.assertIn("nsys_version", result["provenance"]) + self.assertEqual(result["provenance"]["gpu"]["gpu"], "test SM120 device") + self.assertEqual(result["parameters"]["shape_mode"], "runtime") + for command in result["commands"]: + if command["command"][1:2] == ["profile"]: + argv = command["command"] + self.assertEqual(argv[argv.index("--shape-mode") + 1], "runtime") + + def test_static_mode_is_forwarded_to_both_runners_and_recorded(self): + self.args.seq_q, self.args.seq_k, self.args.heads = 1024, 1024, 8 + self.args.shape_mode = "static" + self.hashes = compare.expected_input_hashes(self.args) + self.output = [0.0] * (1024 * 8 * 64) + arguments = self.arguments() + ["--seq-q", "1024", "--seq-k", "1024", "--heads", "8", + "--shape-mode", "static", "--rounds", "1"] + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()) as errors: + self.assertEqual(compare.main(arguments), 0, errors.getvalue()) + directory, = (self.directory / "evidence").iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["parameters"]["shape_mode"], "static") + self.assertEqual([run["label"] for run in result["runs"]], ["rust", "cutlass"]) + for run in result["runs"]: + self.assertEqual(run["runner_report"]["provenance"]["shape_specialization"], "static") + profiles = [entry["command"] for entry in result["commands"] if entry["command"][1:2] == ["profile"]] + self.assertEqual(len(profiles), 2) + for command in profiles: + self.assertEqual(command[command.index("--shape-mode") + 1], "static") + + def test_static_comparison_rejects_unlisted_shapes_before_running_tools(self): + for seq_q, seq_k, heads in ((1, 1024, 8), (512, 1024, 8), (1024, 512, 8), (1024, 1024, 4)): + arguments = self.arguments() + ["--seq-q", str(seq_q), "--seq-k", str(seq_k), + "--heads", str(heads), "--noncausal"] + with self.subTest(seq_q=seq_q, seq_k=seq_k, heads=heads): + self.assertEqual(compare.parse_arguments(arguments).shape_mode, "runtime") + with mock.patch.object(compare, "run_command") as command, \ + contextlib.redirect_stderr(io.StringIO()) as errors, self.assertRaises(SystemExit): + compare.main(arguments + ["--shape-mode", "static"]) + command.assert_not_called() + self.assertIn("static shape mode requires", errors.getvalue()) + + def test_requested_static_mode_rejects_runtime_report(self): + report = self.report() + self.args.shape_mode = "static" + with self.assertRaisesRegex(ValueError, "shape_specialization"): + self.validate(self.write_evidence(report)) + + def test_runner_rejects_missing_wrong_or_invalid_shape_mode_on_either_side(self): + for label in ("rust", "cutlass"): + for mode in (None, "static", "auto"): + with self.subTest(label=label, mode=mode): + self.invalid_shape_label = label + self.invalid_shape_specialization = mode + evidence = self.directory / f"invalid-shape-{label}-{mode}" + arguments = self.arguments() + ["--output-dir", str(evidence), "--rounds", "1"] + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(arguments), 1) + directory, = evidence.iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "failed") + self.assertIn("shape_specialization", result["error"]) + self.assertEqual(len(result["runs"]), 1 if label == "rust" else 2) + self.assertEqual(result["runs"][-1]["label"], label) + self.assertNotIn("summary", result) + + def test_runner_resolves_generic_names_for_both_causal_specializations(self): + for causal, name in ((True, "fmha_prefill__generic_causal_hash"), + (False, "fmha_prefill__generic_noncausal_hash")): + with self.subTest(causal=causal): + self.rust_kernel_name = name + evidence = self.directory / f"generic-{causal}" + arguments = self.arguments() + ["--output-dir", str(evidence), "--rounds", "1", + "--causal" if causal else "--noncausal"] + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()) as errors: + self.assertEqual(compare.main(arguments), 0, errors.getvalue()) + directory, = evidence.iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["runs"][0]["kernel_name"], name) + self.assertEqual(result["runs"][0]["samples_us"], [31, 11, 21]) + self.assertEqual(result["parameters"]["causal"], causal) + for run in result["runs"]: + self.assertEqual(run["runner_report"]["causal"], causal) + self.assertEqual(run["runner_report"]["provenance"]["causal_specialization"], "static") + + def test_runner_rejects_missing_or_runtime_causal_specialization_on_either_side(self): + for label in ("rust", "cutlass"): + for specialization in (None, "runtime"): + with self.subTest(label=label, specialization=specialization): + self.invalid_specialization_label = label + self.invalid_specialization = specialization + evidence = self.directory / f"invalid-{label}-{specialization}" + arguments = self.arguments() + ["--output-dir", str(evidence), "--rounds", "1"] + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(arguments), 1) + directory, = evidence.iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "failed") + self.assertIn("causal_specialization", result["error"]) + self.assertEqual(len(result["runs"]), 1 if label == "rust" else 2) + self.assertEqual(result["runs"][-1]["label"], label) + self.assertEqual(result["runs"][-1]["status"], "failed") + self.assertNotIn("summary", result) + + def test_failed_trace_stops_serial_run_and_never_publishes_a_summary(self): + self.drop_last_launch = True + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-externally-selected"}), \ + mock.patch.object(compare, "run_command", side_effect=self.fake_command), \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(self.arguments()), 1) + directory, = (self.directory / "evidence").iterdir() + result = json.loads((directory / "report.json").read_text()) + self.assertEqual(result["status"], "failed") + self.assertEqual(len(result["runs"]), 1) + self.assertEqual(result["runs"][0]["status"], "failed") + self.assertNotIn("summary", result) + + def test_build_manifest_must_match_binary_before_any_tool_runs(self): + manifest = self.directory / "build.json" + manifest.write_text(json.dumps({"binary_sha256": "0" * 64})) + arguments = self.arguments() + ["--rust-build-manifest", str(manifest)] + with mock.patch.object(compare, "run_command") as command, \ + contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(compare.main(arguments), 1) + command.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_fixture.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_fixture.py new file mode 100644 index 0000000000..5e15ea57f4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/reference/test_fixture.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only checks for the shared fixture, oracle, and exhaustive verifier.""" + +import contextlib +import io +import math +import importlib.util +import unittest +from unittest import mock + +import launch + + +class FixtureTests(unittest.TestCase): + def test_shape_mode_defaults_to_runtime_and_static_accepts_both_masks(self): + self.assertEqual(launch.parse_args([]).shape_mode, "runtime") + for mask in ("--causal", "--noncausal"): + with self.subTest(mask=mask): + args = launch.parse_args(["--shape-mode", "static", mask]) + self.assertEqual((args.seq_q, args.seq_k, args.heads), (1024, 1024, 8)) + self.assertEqual(args.shape_mode, "static") + + def test_static_mode_rejects_every_unsupported_dimension(self): + for dimensions in ((1, 1024, 8), (512, 1024, 8), (1024, 512, 8), (1024, 1024, 4)): + seq_q, seq_k, heads = dimensions + argv = ["--seq-q", str(seq_q), "--seq-k", str(seq_k), "--heads", str(heads), "--noncausal"] + with self.subTest(dimensions=dimensions): + self.assertEqual(launch.parse_args(argv).shape_mode, "runtime") + with contextlib.redirect_stderr(io.StringIO()) as errors, self.assertRaises(SystemExit): + launch.parse_args(argv + ["--shape-mode", "static"]) + self.assertIn("static shape mode requires", errors.getvalue()) + + def test_static_views_keep_shapes_and_strides_while_runtime_marks_both_dynamic(self): + for mode in ("runtime", "static"): + with self.subTest(mode=mode): + tensor, view = mock.Mock(), mock.Mock() + tensor.dim_order.return_value = (0, 1, 2, 3) + view.mark_layout_dynamic.return_value = view + view.mark_compact_shape_dynamic.return_value = view + from_dlpack = mock.Mock(return_value=view) + self.assertIs(launch.make_tensor_view(tensor, from_dlpack, mode), view) + from_dlpack.assert_called_once_with(tensor, assumed_align=16) + if mode == "runtime": + view.mark_layout_dynamic.assert_called_once_with(leading_dim=3) + view.mark_compact_shape_dynamic.assert_called_once_with( + mode=3, stride_order=(0, 1, 2, 3), divisibility=8, + ) + else: + view.mark_layout_dynamic.assert_not_called() + view.mark_compact_shape_dynamic.assert_not_called() + tensor.dim_order.assert_not_called() + + def test_dyadic_inputs_are_exact_fp16_and_distinguish_heads_and_operands(self): + q, k, v = launch.make_inputs(2, 2, 2) + for values in (q, k, v): + self.assertEqual(len(values), 256) + self.assertTrue(all(-0.25 <= x <= 0.2421875 for x in values)) + self.assertEqual(launch.decode_f16(launch.encode_f16(values)), values) + self.assertGreater(len(set(values)), 40) + self.assertNotEqual(q, k) + self.assertNotEqual(k, v) + self.assertNotEqual(q[:64], q[64:128]) + + def test_input_byte_hashes_match_rust_fixture_abi(self): + q, k, v = launch.make_inputs(2, 2, 2) + self.assertEqual([launch.bits_hash(x) for x in (q, k, v)], + ["ab1887afc727f6b1", "c758962db6dc1ae0", "e543f46a99a93138"]) + scaled_q, scaled_k, scaled_v = launch.make_inputs(2, 2, 2, 16) + self.assertEqual(scaled_q, [x * 16 for x in q]) + self.assertEqual(scaled_k, [x * 16 for x in k]) + self.assertEqual(scaled_v, v) + self.assertEqual(launch.make_inputs(2, 2, 2, 16, True)[0], [0.0] * 256) + + @unittest.skipUnless(importlib.util.find_spec("torch"), "optional CPU Torch oracle") + def test_batched_cpu_oracle_matches_scalar_for_scaled_logits_and_both_masks(self): + for scale in (1, 4, 16): + for causal in (False, True): + q, k, v = launch.make_inputs(2, 2, 2, scale) + scalar = launch.reference_cpu(q, k, v, 2, 2, 2, causal) + batched = launch.reference_cpu(q, k, v, 2, 2, 2, causal, use_torch=True) + self.assertLess(max(abs(x - y) for x, y in zip(scalar, batched)), 1e-12) + + def test_oracle_masks_future_keys_and_keeps_bshd_heads_separate(self): + zeros = [0.0] * 256 + values = [2.0] * 64 + [10.0] * 64 + [4.0] * 64 + [14.0] * 64 + full = launch.reference_cpu(zeros, zeros, values, 2, 2, 2, False) + causal = launch.reference_cpu(zeros, zeros, values, 2, 2, 2, True) + self.assertEqual(full, ([3.0] * 64 + [12.0] * 64) * 2) + self.assertEqual(causal, [2.0] * 64 + [10.0] * 64 + [3.0] * 64 + [12.0] * 64) + + def test_oracle_subtracts_row_max_before_exponentiating(self): + result = launch.reference_cpu( + [100.0] * 64, [100.0] * 64 + [-100.0] * 64, + [3.0] * 64 + [-7.0] * 64, 1, 2, 1, False, + ) + self.assertEqual(result, [3.0] * 64) + + def test_verifier_rejects_nonfinite_wrong_length_and_last_element_error(self): + expected = [0.0, 0.01, -0.02, 0.03] + self.assertEqual(launch.verify(expected, expected)["elements"], 4) + for replacement in (math.nan, math.inf, -math.inf, -0.03): + with self.subTest(replacement=replacement), self.assertRaises(ValueError): + launch.verify(expected[:-1] + [replacement], expected) + with self.assertRaises(ValueError): + launch.verify(expected[:-1], expected) + with self.assertRaises(ValueError): + launch.verify(expected, [0.0, 0.01, -0.02, math.nan]) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/comparison.rs b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/comparison.rs new file mode 100644 index 0000000000..ff614f9af7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/comparison.rs @@ -0,0 +1,216 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Shared process launcher for the SM120a prefill and decode comparisons. + +use crate::fixture::Shape; +use std::{ + error::Error, + path::{Path, PathBuf}, + process::Command, +}; + +/// Host-side errors from comparison setup, paths, and process execution. +type Result = std::result::Result>; + +/// Choose explicit, environment, local-venv, shared-venv, then PATH Python. +fn resolve_python( + directory: &Path, + explicit: Option<&PathBuf>, + environment: Option, +) -> PathBuf { + let local_python = directory.join(".venv/bin/python"); + explicit.cloned().or(environment).unwrap_or_else(|| { + if local_python.is_file() { + local_python + } else { + directory + .parent() + .and_then(Path::parent) + .map(|cute_rs_preview| cute_rs_preview.join(".venv/bin/python")) + .filter(|shared_python| shared_python.is_file()) + .unwrap_or_else(|| PathBuf::from("python3")) + } + }) +} + +/// Reject Python selection in standalone mode and raw output in comparison mode. +pub fn validate_mode( + rust_only: bool, + python: Option<&PathBuf>, + output: Option<&PathBuf>, +) -> std::result::Result<(), String> { + if rust_only && python.is_some() { + return Err("--python applies only to the CuTeDSL comparison".into()); + } + if !rust_only && output.is_some() { + return Err("--output requires --rust-only".into()); + } + Ok(()) +} + +/// Run the Python comparison with the same shape, inputs, and timing options. +#[allow(clippy::too_many_arguments)] +pub fn run( + workload: &str, + python: Option<&PathBuf>, + all_cases: bool, + shape: Shape, + input_scale: u32, + zero_query: bool, + device: usize, + warmup: usize, + iters: usize, + mode_flag: &str, + mode: &str, + json: Option<&PathBuf>, +) -> Result<()> { + let directory = Path::new(env!("CARGO_MANIFEST_DIR")); + let prefill = directory + .parent() + .ok_or("example directory has no parent")? + .join("fmha_prefill_sm120a_cute"); + let script = prefill.join("compare.py"); + let reference = directory.join("reference"); + let kernel = if workload == "decode" { + "attention_decode.py" + } else { + "flash_attention_v2.py" + }; + if !script.is_file() + || !reference.join("launch.py").is_file() + || !reference.join(kernel).is_file() + || !prefill.join("reference/launch.py").is_file() + { + return Err(format!("CuTeDSL comparison requires local compare.py and reference/ assets in {}; use --rust-only for standalone execution", directory.display()).into()); + } + let python = resolve_python( + directory, + python, + std::env::var_os("CUDA_OXIDE_CUTE_PYTHON").map(PathBuf::from), + ); + // Dispatch before CUDA context creation; Python invokes this binary with --rust-only. + let mut command = Command::new(&python); + command + .arg(script) + .args(["--workload", workload]) + .arg("--rust-binary") + .arg(std::env::current_exe()?) + .args([ + "--seq-q", + &shape.seq_q.to_string(), + "--seq-k", + &shape.seq_k.to_string(), + "--heads", + &shape.heads.to_string(), + "--input-scale", + &input_scale.to_string(), + "--device", + &device.to_string(), + "--warmup", + &warmup.to_string(), + "--iters", + &iters.to_string(), + mode_flag, + mode, + ]) + .arg(if shape.causal { + "--causal" + } else { + "--noncausal" + }); + if zero_query { + command.arg("--zero-query"); + } + if all_cases { + command.arg("--all-cases"); + } + if let Some(path) = json { + command.arg("--json").arg(path); + } + let status = command.status().map_err(|error| format!( + "could not start CuTeDSL comparison with {}: {error}; select its Python environment with --python or CUDA_OXIDE_CUTE_PYTHON", python.display()))?; + if !status.success() { + return Err(format!( + "CuTeDSL comparison failed with {} ({status})", + python.display() + ) + .into()); + } + Ok(()) +} + +/// Tests for interpreter overrides and virtualenv fallback order. +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Temporary test directory removed when the fixture is dropped. + struct TempDirectory(#[doc = "Root directory containing fake interpreter paths."] PathBuf); + + impl TempDirectory { + /// Create a unique directory for one interpreter-selection test. + fn new() -> Self { + /// Process-local counter that keeps temporary test directory names distinct. + static NEXT: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "fmha-python-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + std::fs::create_dir(&path).unwrap(); + Self(path) + } + + /// Create a placeholder interpreter file at the requested relative path. + fn interpreter(&self, relative: &str) -> PathBuf { + let path = self.0.join(relative); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, "").unwrap(); + path + } + } + + impl Drop for TempDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + /// Check that explicit and environment Python choices override discovered virtualenvs. + #[test] + fn python_overrides_take_precedence_over_virtualenvs() { + let temp = TempDirectory::new(); + let directory = temp.0.join("examples/fmha"); + temp.interpreter("examples/fmha/.venv/bin/python"); + temp.interpreter(".venv/bin/python"); + let explicit = PathBuf::from("explicit-python"); + let environment = PathBuf::from("environment-python"); + assert_eq!( + resolve_python(&directory, Some(&explicit), Some(environment.clone())), + explicit + ); + assert_eq!( + resolve_python(&directory, None, Some(environment.clone())), + environment + ); + } + + /// Check the local-venv, shared-venv, and PATH fallback order. + #[test] + fn python_falls_back_from_local_to_shared_virtualenv_then_path() { + let temp = TempDirectory::new(); + let directory = temp.0.join("examples/fmha"); + assert_eq!( + resolve_python(&directory, None, None), + PathBuf::from("python3") + ); + let shared = temp.interpreter(".venv/bin/python"); + assert_eq!(resolve_python(&directory, None, None), shared); + let local = temp.interpreter("examples/fmha/.venv/bin/python"); + assert_eq!(resolve_python(&directory, None, None), local); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/fixture.rs b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/fixture.rs new file mode 100644 index 0000000000..5e28b09293 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/fixture.rs @@ -0,0 +1,338 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Deterministic BSHD inputs and an independent FP64 attention oracle. + +/// FP16 values per attention head. +pub const HEAD_DIM: usize = 64; +/// Dot-product scale, equal to the inverse square root of the head dimension. +pub const SOFTMAX_SCALE: f64 = 0.125; +/// Absolute error allowance when comparing FP16 output with the FP64 oracle. +pub const ATOL: f64 = 2.0e-4; +/// Relative error allowance, multiplied by the oracle value's magnitude. +pub const RTOL: f64 = 2.0e-3; +/// Fixed seed used to generate query values independently of K and V. +const Q_SEED: u32 = 0xa511_e9b3; +/// Fixed seed used to generate key values independently of Q and V. +const K_SEED: u32 = 0x63d8_3595; +/// Fixed seed used to generate value inputs independently of Q and K. +const V_SEED: u32 = 0xc2b2_ae35; + +/// Single-batch attention dimensions and mask selection in BSHD order. +#[derive(Clone, Copy, Debug)] +pub struct Shape { + /// Number of query rows. + pub seq_q: usize, + /// Number of key/value rows. + pub seq_k: usize, + /// Number of independent attention heads. + pub heads: usize, + /// Whether each query may attend only to its own and earlier positions. + pub causal: bool, +} + +impl Shape { + #[allow(dead_code)] + /// Check prefill tile constraints, causal dimensions, and 32-bit index bounds. + pub fn validate(self) -> Result<(), String> { + if self.seq_q == 0 || self.seq_k == 0 || self.heads == 0 { + return Err("sequence lengths and head count must be positive".into()); + } + if (self.seq_q != 1 && !self.seq_q.is_multiple_of(64)) || !self.seq_k.is_multiple_of(64) { + return Err( + "query length must be 1 or a multiple of 64; KV length must be a multiple of 64" + .into(), + ); + } + if self.causal && self.seq_q != self.seq_k { + return Err("the initial causal specialization requires equal sequence lengths".into()); + } + for length in [self.seq_q, self.seq_k] { + let elements = length + .checked_mul(self.heads) + .and_then(|n| n.checked_mul(HEAD_DIM)); + if elements.is_none_or(|n| n > u32::MAX as usize) { + return Err("tensor element count must fit u32 indexing".into()); + } + } + Ok(()) + } + + /// Number of FP16 query or output elements; call after shape validation. + pub fn q_elements(self) -> usize { + self.seq_q * self.heads * HEAD_DIM + } + + /// Number of FP16 elements in either K or V; call after shape validation. + pub fn kv_elements(self) -> usize { + self.seq_k * self.heads * HEAD_DIM + } +} + +/// Deterministically mix an index and seed using wrapping 32-bit arithmetic. +fn mix32(mut value: u32) -> u32 { + value ^= value >> 16; + value = value.wrapping_mul(0x7feb_352d); + value ^= value >> 15; + value = value.wrapping_mul(0x846c_a68b); + value ^ (value >> 16) +} + +/// Generate reproducible FP16 values with the requested power-of-two scale. +fn input(length: usize, seed: u32, scale: u32) -> Vec { + (0..length) + .map(|i| { + let integer = (mix32(i as u32 ^ seed) & 63) as i32 - 32; + (integer as f32 * scale as f32 / 128.0) as f16 + }) + .collect() +} + +/// Scale Q/K by an exact power of two to exercise larger logits; V is unchanged. +pub fn make_inputs( + shape: Shape, + input_scale: u32, + zero_query: bool, +) -> (Vec, Vec, Vec) { + let q = if zero_query { + vec![0.0_f16; shape.q_elements()] + } else { + input(shape.q_elements(), Q_SEED, input_scale) + }; + ( + q, + input(shape.kv_elements(), K_SEED, input_scale), + input(shape.kv_elements(), V_SEED, 1), + ) +} + +/// FNV-1a over canonical little-endian FP16 bytes, matching reference/launch.py. +pub fn bits_hash(values: &[f16]) -> u64 { + let mut hash = 0xcbf2_9ce4_8422_2325u64; + for value in values { + for byte in value.to_bits().to_le_bytes() { + hash = (hash ^ u64::from(byte)).wrapping_mul(0x0000_0100_0000_01b3); + } + } + hash +} + +/// Mathematical attention, with all arithmetic and probabilities in FP64. +/// This intentionally does not reproduce the kernel's FP16 probability rounding. +pub fn reference(shape: Shape, q: &[f16], k: &[f16], v: &[f16]) -> Vec { + assert_eq!(q.len(), shape.q_elements()); + assert_eq!(k.len(), shape.kv_elements()); + assert_eq!(v.len(), shape.kv_elements()); + let q: Vec = q.iter().map(|&x| x as f64).collect(); + let k: Vec = k.iter().map(|&x| x as f64).collect(); + let v: Vec = v.iter().map(|&x| x as f64).collect(); + let mut output = vec![0.0; shape.q_elements()]; + let mut scores = vec![0.0; shape.seq_k]; + for head in 0..shape.heads { + for row in 0..shape.seq_q { + let q_base = (row * shape.heads + head) * HEAD_DIM; + let end = if shape.causal { + (row + 1).min(shape.seq_k) + } else { + shape.seq_k + }; + let mut maximum = f64::NEG_INFINITY; + for (column, score) in scores[..end].iter_mut().enumerate() { + let k_base = (column * shape.heads + head) * HEAD_DIM; + let mut dot = 0.0; + for d in 0..HEAD_DIM { + dot += q[q_base + d] * k[k_base + d]; + } + *score = dot * SOFTMAX_SCALE; + maximum = maximum.max(*score); + } + let mut denominator = 0.0; + for score in &mut scores[..end] { + *score = (*score - maximum).exp(); + denominator += *score; + } + let mut accumulator = [0.0f64; HEAD_DIM]; + for (column, &weight) in scores[..end].iter().enumerate() { + let v_base = (column * shape.heads + head) * HEAD_DIM; + for d in 0..HEAD_DIM { + accumulator[d] += weight * v[v_base + d]; + } + } + for d in 0..HEAD_DIM { + output[q_base + d] = accumulator[d] / denominator; + } + } + } + output +} + +/// Error statistics from checking every output element against the oracle. +pub struct Verification { + /// Number of output elements checked. + pub elements: usize, + /// Largest absolute difference from the FP64 oracle. + pub max_abs_error: f64, + /// Largest error divided by its allowed absolute-plus-relative bound. + pub max_tolerance_ratio: f64, +} + +/// Require every element to be finite and satisfy atol + rtol * abs(expected). +pub fn verify(got: &[f16], expected: &[f64]) -> Result { + if got.len() != expected.len() || got.is_empty() { + return Err(format!( + "output length {} differs from nonempty oracle length {}", + got.len(), + expected.len() + )); + } + let mut result = Verification { + elements: got.len(), + max_abs_error: 0.0, + max_tolerance_ratio: 0.0, + }; + for (index, (&actual, &want)) in got.iter().zip(expected).enumerate() { + let actual = actual as f64; + if !actual.is_finite() || !want.is_finite() { + return Err(format!( + "nonfinite output/oracle at {index}: actual={actual}, expected={want}" + )); + } + let error = (actual - want).abs(); + let bound = ATOL + RTOL * want.abs(); + if error > bound { + return Err(format!( + "output mismatch at {index}: actual={actual}, expected={want}, error={error}, bound={bound}" + )); + } + result.max_abs_error = result.max_abs_error.max(error); + result.max_tolerance_ratio = result.max_tolerance_ratio.max(error / bound); + } + Ok(result) +} + +/// Tests for deterministic inputs, the attention oracle, and full-output validation. +#[cfg(test)] +mod tests { + use super::*; + + /// Check fixture hashes and exact power-of-two scaling against the Python convention. + #[test] + fn fixture_bytes_match_python_and_scaling_preserves_values() { + let shape = Shape { + seq_q: 2, + seq_k: 2, + heads: 2, + causal: false, + }; + let (q, k, v) = make_inputs(shape, 1, false); + assert_eq!( + [bits_hash(&q), bits_hash(&k), bits_hash(&v)], + [ + 0xab18_87af_c727_f6b1, + 0xc758_962d_b6dc_1ae0, + 0xe543_f46a_99a9_3138 + ] + ); + let (q16, k16, v16) = make_inputs(shape, 16, false); + assert!( + q16.iter() + .zip(&q) + .all(|(&a, &b)| a as f32 == b as f32 * 16.0) + ); + assert!( + k16.iter() + .zip(&k) + .all(|(&a, &b)| a as f32 == b as f32 * 16.0) + ); + assert_eq!(v16, v); + assert!(make_inputs(shape, 16, true).0.iter().all(|&x| x == 0.0_f16)); + } + + /// Check causal masking and independent head indexing using uniform logits. + #[test] + fn oracle_masks_future_keys_and_keeps_heads_separate() { + let mut shape = Shape { + seq_q: 2, + seq_k: 2, + heads: 2, + causal: false, + }; + let zeros = vec![0.0_f16; 256]; + let values: Vec = [2.0, 10.0, 4.0, 14.0] + .into_iter() + .flat_map(|x| [x as f16; 64]) + .collect(); + let expected: Vec = [3.0, 12.0, 3.0, 12.0] + .into_iter() + .flat_map(|x| [x; 64]) + .collect(); + assert_eq!(reference(shape, &zeros, &zeros, &values), expected); + shape.causal = true; + let expected: Vec = [2.0, 10.0, 3.0, 12.0] + .into_iter() + .flat_map(|x| [x; 64]) + .collect(); + assert_eq!(reference(shape, &zeros, &zeros, &values), expected); + } + + /// Check stable softmax with logits too large for unshifted exponentials. + #[test] + fn oracle_handles_logits_that_overflow_unshifted_exponentials() { + let shape = Shape { + seq_q: 1, + seq_k: 2, + heads: 1, + causal: false, + }; + let q = vec![100.0_f16; 64]; + let k: Vec = [100.0, -100.0] + .into_iter() + .flat_map(|x| [x as f16; 64]) + .collect(); + let v: Vec = [3.0, -7.0] + .into_iter() + .flat_map(|x| [x as f16; 64]) + .collect(); + assert_eq!(reference(shape, &q, &k, &v), vec![3.0; 64]); + } + + /// Ensure every element, including the tail, is checked for finite and accurate output. + #[test] + fn verifier_rejects_nonfinite_and_unsampled_tail_corruption() { + let expected = [0.0, 0.01, -0.02, 0.03]; + let mut got: Vec = expected.iter().map(|&x| x as f16).collect(); + assert_eq!(verify(&got, &expected).unwrap().elements, 4); + assert!(verify(&got[..3], &expected).is_err()); + for bad in [f16::NAN, f16::INFINITY, f16::NEG_INFINITY, -0.03_f16] { + got[3] = bad; + assert!(verify(&got, &expected).is_err()); + } + } + + /// Check positive dimensions, tile boundaries, and equal-length causal sequences. + #[test] + fn shape_rejects_zero_and_asymmetric_causal_sequences() { + let good = Shape { + seq_q: 64, + seq_k: 128, + heads: 8, + causal: false, + }; + assert!(good.validate().is_ok()); + assert!(Shape { seq_q: 1, ..good }.validate().is_ok()); + assert!(Shape { heads: 0, ..good }.validate().is_err()); + assert!(Shape { seq_q: 0, ..good }.validate().is_err()); + assert!(Shape { seq_k: 0, ..good }.validate().is_err()); + assert!( + Shape { + causal: true, + ..good + } + .validate() + .is_err() + ); + assert!(Shape { seq_q: 65, ..good }.validate().is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/kernel.rs b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/kernel.rs new file mode 100644 index 0000000000..938f9a58ca --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/kernel.rs @@ -0,0 +1,458 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Ampere-style FA2: typed CuTe copies, register MMA, and online softmax. +//! One CTA owns 64 queries; four warps each own 16 rows of a 64-column tile. +use core::marker::PhantomData; +use ctm_rs::{ + async_copy, barrier, convert, cuda_module, float, kernel, launch_bounds, launch_contract, + memory::{DisjointSlice, DynamicSharedArray}, + thread, warp, +}; +use cute_rs::markers::{C, ComposedLayout, CopyG2S, Layout, LeadingDim, RowMajor, Swizzle}; +use cute_rs::{FragmentA, FragmentBPair, FragmentC, GmemMatrix, SmemTile, copy_g2s}; + +/// Swizzled 64-by-64 FP16 layout shared by Q, K, V, and the output staging tile. +type SharedLayout = ComposedLayout, 0, RowMajor<64, 64>>; +/// Swizzled 16-by-64 subview used by one cooperative copy. +type CopyLayout = ComposedLayout, 0, RowMajor<16, 64>>; +/// Map 128 threads to 16 rows and eight 16-byte copy positions. +type CopyThreads = Layout<(C<16>, C<8>), (C<8>, C<1>)>; +/// Eight adjacent FP16 values copied by each thread. +type CopyValues = Layout<(C<1>, C<8>), (C<8>, C<1>)>; +/// Typed view of one 64-by-64 FP16 shared tile. +type SharedTile = SmemTile; + +// Constant indices keep fragments in registers and allow the compiler to +// specialize each MMA, without repeating the algorithm for each fragment. +/// Expand fixed iterations with constant indices so MMA fragments stay in registers. +macro_rules! static_for { + ($index:ident in [$($value:expr),*] $body:block) => { + $({ #[doc = "Compile-time index for this unrolled iteration."] const $index: usize = $value; $body })* + }; +} + +/// The eight 16x8 accumulators making one warp's 16x64 tile. +struct WarpTile(#[doc = "Eight adjacent 16-by-8 FP32 accumulator fragments."] [FragmentC; 8]); + +impl WarpTile { + /// Create an all-zero warp accumulator tile. + #[inline(always)] + fn zero() -> Self { + Self([FragmentC::ZERO; 8]) + } + + /// Accumulate one K16 step across all eight column fragments, one atom call each. + /// + /// # Safety + /// All warp lanes must participate with matching, valid operand fragments. + #[inline(always)] + unsafe fn gemm(&mut self, a: FragmentA, b: [FragmentBPair; 4]) { + static_for!(PAIR in [0, 1, 2, 3] { + self.0[2 * PAIR] = + unsafe { self.0[2 * PAIR].mma_atom_call(a, b[PAIR].get::<0>()) }; + self.0[2 * PAIR + 1] = + unsafe { self.0[2 * PAIR + 1].mma_atom_call(a, b[PAIR].get::<1>()) }; + }); + } + + /// Retile two adjacent C fragments into one A fragment, entirely in registers. + #[inline(always)] + fn probability(&self) -> FragmentA { + const { assert!(K < 4) }; + let left = self.0[2 * K].0; + let right = self.0[2 * K + 1].0; + FragmentA([ + pack(left[0], left[1]), + pack(left[2], left[3]), + pack(right[0], right[1]), + pack(right[2], right[3]), + ]) + } + + /// Narrow the whole probability tile before loading the value fragments. + #[inline(always)] + fn probabilities(&self) -> [FragmentA; 4] { + let mut fragments = [FragmentA([0; 4]); 4]; + static_for!(K in [0, 1, 2, 3] { + fragments[K] = self.probability::(); + }); + fragments + } + + /// Scale one of this lane's two rows across all accumulator fragments. + #[inline(always)] + fn rescale(&mut self, scale: f32) { + static_for!(N in [0, 1, 2, 3, 4, 5, 6, 7] { + self.0[N].0[ROW * 2] *= scale; + self.0[N].0[ROW * 2 + 1] *= scale; + }); + } + + /// Reduce this lane's 16 scores as a tree, as in the CuTe row reduction. + #[inline(always)] + fn maximum(&self) -> f32 { + let mut pairs = [0.0; 8]; + static_for!(N in [0, 1, 2, 3, 4, 5, 6, 7] { + pairs[N] = self.0[N].0[ROW * 2].max(self.0[N].0[ROW * 2 + 1]); + }); + static_for!(N in [0, 1, 2, 3] { + pairs[N] = pairs[2 * N].max(pairs[2 * N + 1]); + }); + static_for!(N in [0, 1] { + pairs[N] = pairs[2 * N].max(pairs[2 * N + 1]); + }); + pairs[0].max(pairs[1]) + } + + /// Mask and exponentiate scores, updating the running normalization and output scale. + #[inline(always)] + fn softmax( + &mut self, + output: &mut Self, + state: &mut Softmax, + query_row: u32, + lane: u32, + ) { + static_for!(ROW in [0, 1] { + let row = query_row + ROW as u32 * 8; + static_for!(N in [0, 1, 2, 3, 4, 5, 6, 7] { + static_for!(COL in [0, 1] { + let column = N as u32 * 8 + (lane % 4) * 2 + COL as u32; + let value = if CAUSAL && FIRST && column > row { f32::NEG_INFINITY } + else { self.0[N].0[ROW * 2 + COL] }; + self.0[N].0[ROW * 2 + COL] = value; + }); + }); + let mut maximum = self.maximum::(); + maximum = max(maximum, warp::shuffle_xor_f32(maximum, 2)); + maximum = max(maximum, warp::shuffle_xor_f32(maximum, 1)); + if !FIRST { maximum = max(maximum, state.maximum[ROW]); } + // Causal traversal may begin with a fully masked row in this tile. + if CAUSAL && maximum == f32::NEG_INFINITY { maximum = 0.0; } + /// Attention scale converted to base two for the exponential instruction. + const SCALE_LOG2: f32 = 0.125 * core::f32::consts::LOG2_E; + let mut sum = 0.0; + static_for!(N in [0, 1, 2, 3, 4, 5, 6, 7] { + static_for!(COL in [0, 1] { + let value = float::ex2_approx_ftz_f32( + self.0[N].0[ROW * 2 + COL] * SCALE_LOG2 - maximum * SCALE_LOG2, + ); + self.0[N].0[ROW * 2 + COL] = value; + sum += value; + }); + }); + if FIRST { + state.sum[ROW] = sum; + } else { + let old_scale = float::ex2_approx_ftz_f32( + state.maximum[ROW] * SCALE_LOG2 - maximum * SCALE_LOG2, + ); + state.sum[ROW] = sum + state.sum[ROW] * old_scale; + output.rescale::(old_scale); + } + state.maximum[ROW] = maximum; + }); + } + + /// Normalize the output and pack it into the shared tile. + /// + /// # Safety + /// All warp lanes must participate; `dst` must provide disjoint slots for each warp. + #[inline(always)] + unsafe fn store(&mut self, state: &Softmax, dst: &SharedTile, warp_id: usize, lane: usize) { + static_for!(ROW in [0, 1] { + let mut sum = state.sum[ROW]; + sum += warp::shuffle_xor_f32(sum, 2); + sum += warp::shuffle_xor_f32(sum, 1); + let scale = if sum == 0.0 || sum.is_nan() { 1.0 } + else { float::rcp_approx_ftz_f32(sum) }; + self.rescale::(scale); + static_for!(N in [0, 1, 2, 3, 4, 5, 6, 7] { + let row = warp_id * 16 + lane / 4 + ROW * 8; + let col = N * 8 + lane % 4 * 2; + unsafe { + *dst.base.add(shared_offset(row as u32, col as u32)).cast::() = + pack(self.0[N].0[ROW * 2], self.0[N].0[ROW * 2 + 1]); + } + }); + }); + } +} + +/// Running softmax statistics for the two rows represented by each lane. +struct Softmax { + /// Running maximum for each row, shared by its four cooperating lanes. + maximum: [f32; 2], + /// Lane-local sums to be combined across each row's four lanes at output time. + sum: [f32; 2], +} + +/// Return the floating-point maximum used by row reductions. +#[inline(always)] +fn max(a: f32, b: f32) -> f32 { + a.max(b) +} + +/// Round two FP32 values to the packed FP16 representation used by MMA and stores. +#[inline(always)] +fn pack(a: f32, b: f32) -> u32 { + convert::cvt_f16x2_f32(a, b) +} + +/// Physical coordinate of the typed `Swizzle<3,3,3>` shared layout. +#[inline(always)] +fn shared_offset(row: u32, column: u32) -> usize { + // Within this 64-column tile, the row and column occupy disjoint bits. + // Keep the row stride separate so adjacent fragment rows share an offset. + (row * 64 + (column ^ ((row & 7) << 3))) as usize +} + +/// Wrap storage for 4096 FP16 values as a typed shared tile. +/// +/// # Safety +/// `base` must point to live shared storage with the tile's required alignment. +#[inline(always)] +unsafe fn tile(base: *mut f16) -> SharedTile { + SmemTile { + base, + capacity: 4096, + layout: PhantomData, + } +} + +/// Issue four cooperative copies to stage one 64-row tile. +/// +/// # Safety +/// All 128 threads must participate with valid views; wait before reading the destination. +#[inline(always)] +unsafe fn stage(src: &GmemMatrix>, block: usize, dst: &SharedTile, tid: u32) { + static_for!(GROUP in [0, 1, 2, 3] { + let mut part = SmemTile:: { + base: unsafe { dst.base.add(GROUP * 1024) }, capacity: 1024, layout: PhantomData, + }; + unsafe { + copy_g2s::, CopyThreads, CopyValues, RowMajor<16, 64>, CopyLayout, f16>( + src, (block * 4 + GROUP, 0), &mut part, tid, + ); + } + }); +} + +/// Stage a complete KV tile. The checked host shape is a multiple of 64, +/// and reverse traversal keeps `block` within that shape. +#[inline(always)] +unsafe fn stage_full( + src: &GmemMatrix>, + block: usize, + dst: &SharedTile, + tid: u32, +) { + let source_tile = GmemMatrix { + base: unsafe { src.base.add(block * 64 * src.leading_dim.elements) }, + rows: 64, + cols: 64, + leading_dim: LeadingDim::<8> { + elements: src.leading_dim.elements, + }, + }; + unsafe { stage(&source_tile, 0, dst, tid) }; +} + +/// Load four pairs of K-major B fragments for one K16 reduction step. +/// +/// # Safety +/// The shared tile must be ready, and every warp lane must participate with its own lane ID. +#[inline(always)] +unsafe fn load_k(src: &SharedTile, lane: u32) -> [FragmentBPair; 4] { + let mut fragments = [FragmentBPair([0; 4]); 4]; + static_for!(N in [0, 1, 2, 3] { + fragments[N] = unsafe { cute_rs::load_matrix_b_pair_k_major(src, (N, K), lane) }; + }); + fragments +} + +/// Load four pairs of V fragments for one K16 reduction step. +/// +/// # Safety +/// The shared tile must be ready, and every warp lane must participate with its own lane ID. +#[inline(always)] +unsafe fn load_v(src: &SharedTile, lane: u32) -> [FragmentBPair; 4] { + let mut fragments = [FragmentBPair([0; 4]); 4]; + static_for!(N in [0, 1, 2, 3] { + fragments[N] = unsafe { cute_rs::load_matrix_b_pair(src, (K, N), lane) }; + }); + fragments +} + +/// Shared-memory pipeline for one query tile. Only the first visited KV tile +/// can contain future tokens; all earlier tiles are wholly visible. +struct TilePipeline { + /// Query tile, reused as output staging after the mainloop. + q: SharedTile, + /// Shared key tile, refilled while the current scores are consumed. + k: SharedTile, + /// Shared value tile used by the probability-times-value MMA. + v: SharedTile, + /// Global K matrix for the CTA's attention head. + keys: GmemMatrix>, + /// Global V matrix for the CTA's attention head. + values: GmemMatrix>, + /// CTA-local thread ID used for cooperative copies and warp ownership. + tid: u32, +} + +impl TilePipeline { + /// Consume one KV tile, update softmax and output, and prefetch the preceding key tile. + /// + /// # Safety + /// All 128 threads must call in lockstep with live, nonoverlapping shared tiles and valid global views. + #[inline(always)] + unsafe fn step( + &self, + key_block: usize, + output: &mut WarpTile, + state: &mut Softmax, + ) { + let tid = self.tid; + let lane = tid % 32; + let warp_id = (tid / 32) as usize; + // SAFETY: the caller supplies full KV tiles and the shared allocations + // remain live across all steps; waits and CTA barriers order reuse. + unsafe { + async_copy::cp_async_wait_group::<0>(); + barrier::barrier_cta_sync(1, 128); + stage_full(&self.values, key_block, &self.v, tid); + async_copy::cp_async_commit_group(); + let mut scores = WarpTile::zero(); + let mut a = cute_rs::load_matrix_a(&self.q, (warp_id, 0), lane); + let mut b = load_k::<0>(&self.k, lane); + static_for!(K in [0, 1, 2] { + let next_a = cute_rs::load_matrix_a(&self.q, (warp_id, K + 1), lane); + let next_b = load_k::<{ K + 1 }>(&self.k, lane); + scores.gemm(a, b); + a = next_a; + b = next_b; + }); + scores.gemm(a, b); + async_copy::cp_async_wait_group::<0>(); + barrier::barrier_cta_sync(1, 128); + if key_block > 0 { + stage_full(&self.keys, key_block - 1, &self.k, tid); + async_copy::cp_async_commit_group(); + } + scores.softmax::(output, state, warp_id as u32 * 16 + lane / 4, lane); + let probabilities = scores.probabilities(); + let mut values = load_v::<0>(&self.v, lane); + static_for!(K in [0, 1, 2] { + let next = load_v::<{ K + 1 }>(&self.v, lane); + output.gemm(probabilities[K], values); + values = next; + }); + output.gemm(probabilities[3], values); + } + } +} + +/// Device entry point and generated host launcher for prefill attention. +#[cuda_module] +pub mod kernels { + use super::*; + + /// Compute one head's 64-query tile with reverse KV traversal and online softmax. + #[launch_bounds(128)] + #[launch_contract(domain=1, block=(128,1,1), dynamic_shared=24576, + dynamic_shared_alignment=1024, min_compute_capability=(12,0))] + #[kernel] + pub fn fmha_prefill( + q: &[f16], + k: &[f16], + v: &[f16], + mut out: DisjointSlice, + seq_q: u32, + seq_k: u32, + heads: u32, + ) { + // Zero selects runtime geometry. Static instantiations fix all three + // dimensions; the host checks them before dispatching this entry point. + const { assert!((SQ == 0 && SK == 0 && HEADS == 0) || (SQ == 1024 && SK == 1024 && HEADS == 8)) }; + let seq_q = if SQ == 0 { seq_q } else { SQ }; + let seq_k = if SK == 0 { seq_k } else { SK }; + let heads = if HEADS == 0 { heads } else { HEADS }; + let tid = thread::threadIdx_x(); + let lane = tid % 32; + let warp_id = (tid / 32) as usize; + let query_blocks = seq_q.div_ceil(64); + let query_block = (thread::blockIdx_x() % query_blocks) as usize; + let head = (thread::blockIdx_x() / query_blocks) as usize; + let stride = heads as usize * 64; + unsafe { + let shared = DynamicSharedArray::::get_raw().cast::(); + let sq = tile(shared); + let sk = tile(shared.add(4096)); + let sv = tile(shared.add(8192)); + let q_global = GmemMatrix { + base: q.as_ptr().add(head * 64), + rows: seq_q as usize, + cols: 64, + leading_dim: LeadingDim::<8> { elements: stride }, + }; + let k_global = GmemMatrix { + base: k.as_ptr().add(head * 64), + rows: seq_k as usize, + cols: 64, + leading_dim: LeadingDim::<8> { elements: stride }, + }; + let v_global = GmemMatrix { + base: v.as_ptr().add(head * 64), + rows: seq_k as usize, + cols: 64, + leading_dim: LeadingDim::<8> { elements: stride }, + }; + // Causal shapes have equal sequence lengths. The first visited + // tile is the diagonal, so its mask uses tile-local coordinates. + let mut remaining = if CAUSAL { + query_block + 1 + } else { + seq_k as usize / 64 + }; + stage(&q_global, query_block, &sq, tid); + stage_full(&k_global, remaining - 1, &sk, tid); + async_copy::cp_async_commit_group(); + let mut output = WarpTile::zero(); + let mut state = Softmax { + maximum: [f32::NEG_INFINITY; 2], + sum: [0.0; 2], + }; + let pipeline = TilePipeline { + q: sq, + k: sk, + v: sv, + keys: k_global, + values: v_global, + tid, + }; + // The first tile initializes softmax and handles the causal diagonal. + // Keep the steady-state loop free of masks and first-tile rescaling. + remaining -= 1; + pipeline.step::(remaining, &mut output, &mut state); + while remaining > 0 { + remaining -= 1; + pipeline.step::(remaining, &mut output, &mut state); + } + output.store(&state, &pipeline.q, warp_id, lane as usize); + barrier::barrier_cta_sync(1, 128); + static_for!(GROUP in [0, 1, 2, 3] { + let row = tid as usize / 8 + GROUP * 16; + let col = tid as usize % 8 * 8; + if query_block * 64 + row < seq_q as usize { + let packed = *pipeline.q.base.add(shared_offset(row as u32, col as u32)).cast::(); + *out.as_mut_ptr().add((query_block * 64 + row) * stride + head * 64 + col) + .cast::() = packed; + } + }); + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/main.rs new file mode 100644 index 0000000000..125b6c30d0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/main.rs @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM120a prefill attention example and comparison entry point. + +#![feature(f16)] + +/// Shared CuTeDSL comparison process and interpreter selection. +mod comparison; +mod fixture; +mod kernel; +/// Host setup, correctness checks, and CUDA-event timing. +mod runner; + +/// Run prefill verification, timing, or the requested CuTeDSL comparison. +fn main() -> Result<(), Box> { + runner::run() +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/runner.rs b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/runner.rs new file mode 100644 index 0000000000..0d67ce582e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/src/runner.rs @@ -0,0 +1,400 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CLI dispatch, correctness checks, and timing for SM120a prefill attention. + +use crate::{fixture, kernel}; + +use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D, sys}; +use fixture::{HEAD_DIM, Shape, bits_hash, make_inputs, reference, verify}; +use serde_json::json; +use std::{error::Error, path::PathBuf}; + +/// Host-side errors from CUDA, file I/O, validation, or comparison execution. +type Result = std::result::Result>; +/// Threads in each 64-query CTA. +const THREADS: u32 = 128; +/// Shared bytes reserved for the three FP16 Q/K/V tiles. +const SHARED_BYTES: u32 = 24_576; + +/// Whether sequence dimensions are runtime arguments or a precompiled shape. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ShapeMode { + /// Pass query length, KV length, and head count as kernel arguments. + Runtime, + /// Use the precompiled 1024-query, 1024-key, eight-head shape. + Static, +} + +/// Validated prefill workload, comparison, and reporting options. +struct Args { + /// Run the Rust kernel alone instead of launching the CuTeDSL comparison. + rust_only: bool, + /// Compare every supported mode combination for the selected workload. + all_cases: bool, + /// Optional Python interpreter for the comparison process. + python: Option, + /// Requested runtime or precompiled sequence dimensions. + shape_mode: ShapeMode, + /// Attention dimensions and mask after CLI validation. + shape: Shape, + /// Power-of-two multiplier applied to Q and K, leaving V unchanged. + input_scale: u32, + /// Use an all-zero Q tensor to exercise uniform softmax. + zero_query: bool, + /// CUDA device ordinal after visible-device remapping. + device: usize, + /// Untimed invocations before collecting samples. + warmup: usize, + /// Number of timed samples; zero requests correctness checks only. + iters: usize, + /// Optional path for raw little-endian FP16 output in standalone mode. + output: Option, + /// Optional path for the standalone or comparison JSON report. + json: Option, +} + +impl Args { + /// Read and validate options from the process command line. + fn parse() -> Result { + Self::parse_from(std::env::args().skip(1)) + } + + /// Parse prefill options, enforce supported shapes, and handle the help request. + fn parse_from(mut cli: impl Iterator) -> Result { + let mut args = Self { + rust_only: false, + all_cases: false, + python: None, + shape_mode: ShapeMode::Runtime, + shape: Shape { + seq_q: 1024, + seq_k: 1024, + heads: 8, + causal: true, + }, + input_scale: 1, + zero_query: false, + device: 0, + warmup: 100, + iters: 31, + output: None, + json: None, + }; + while let Some(key) = cli.next() { + match key.as_str() { + "--help" | "-h" => { + let mode = "Causal mode defaults to Sq=Sk=1024; --noncausal also permits Sq=1."; + println!( + "{} [--seq-q {}] [--seq-k 1024] [--heads 8]\n [--causal | --noncausal] [--shape-mode runtime|static] [--input-scale 1|4|16] [--zero-query]\n [--device 0] [--warmup 100] [--iters 31] [--json result.json] [--python PATH] [--rust-only | --all-cases]\nDefault: verify and compare CuTeDSL and cute-rs-preview; prints two performance lines.\n--all-cases prints a four-case runtime/static x causal/noncausal comparison table; requires Sq=Sk=1024, H=8.\nIt overrides the selected --shape-mode and mask; cannot be combined with --rust-only.\nLive timings use CUDA events; archived Nsight measurements are a separate metric.\n--json saves the combined report; --rust-only saves standalone JSON and permits --output o.f16.\nPython selection: --python, CUDA_OXIDE_CUTE_PYTHON, example .venv, cute-rs-preview/.venv, then python3.\nB=1, D=64, FP16, softmax scale=0.125, 64x64 tiles, 128 threads.\nEvery output is checked against a CPU FP64 oracle before timing.\n--input-scale scales Q/K only; --zero-query exercises uniform softmax.\n--iters 0 verifies two launches without timing.\n{mode}\nStatic shapes require Sq=Sk=1024 and H=8; runtime is the default.\nKV length must be a positive multiple of 64; causal mode requires Sq=Sk.\nDevice ordinal is after CUDA_VISIBLE_DEVICES remapping.", + env!("CARGO_PKG_NAME"), + 1024, + ); + std::process::exit(0); + } + "--rust-only" => args.rust_only = true, + "--all-cases" => args.all_cases = true, + "--causal" => args.shape.causal = true, + "--noncausal" => args.shape.causal = false, + "--zero-query" => args.zero_query = true, + _ => { + let value = cli + .next() + .ok_or_else(|| format!("missing value for {key}"))?; + match key.as_str() { + "--shape-mode" => { + args.shape_mode = match value.as_str() { + "runtime" => ShapeMode::Runtime, + "static" => ShapeMode::Static, + _ => return Err("--shape-mode must be runtime or static".into()), + } + } + "--seq-q" => args.shape.seq_q = value.parse()?, + "--seq-k" => args.shape.seq_k = value.parse()?, + "--heads" => args.shape.heads = value.parse()?, + "--input-scale" => args.input_scale = value.parse()?, + "--device" => args.device = value.parse()?, + "--warmup" => args.warmup = value.parse()?, + "--iters" => args.iters = value.parse()?, + "--python" => args.python = Some(value.into()), + "--output" => args.output = Some(value.into()), + "--json" => args.json = Some(value.into()), + _ => return Err(format!("unknown argument {key}").into()), + } + } + } + } + crate::comparison::validate_mode( + args.rust_only, + args.python.as_ref(), + args.output.as_ref(), + )?; + if args.all_cases && args.rust_only { + return Err("--all-cases cannot be combined with --rust-only".into()); + } + if args.all_cases + && (args.shape.seq_q != 1024 || args.shape.seq_k != 1024 || args.shape.heads != 8) + { + return Err("--all-cases requires Sq=Sk=1024 and H=8".into()); + } + args.shape.validate()?; + if args.shape_mode == ShapeMode::Static + && (args.shape.seq_q != 1024 || args.shape.seq_k != 1024 || args.shape.heads != 8) + { + return Err("--shape-mode static requires Sq=Sk=1024 and H=8".into()); + } + if ![1, 4, 16].contains(&args.input_scale) { + return Err("--input-scale must be 1, 4, or 16".into()); + } + Ok(args) + } +} + +/// Dispatch the comparison process or a standalone prefill run. +pub fn run() -> Result<()> { + let args = Args::parse()?; + if !args.rust_only { + return crate::comparison::run( + "prefill", + args.python.as_ref(), + args.all_cases, + args.shape, + args.input_scale, + args.zero_query, + args.device, + args.warmup, + args.iters, + "--shape-mode", + if args.shape_mode == ShapeMode::Static { + "static" + } else { + "runtime" + }, + args.json.as_ref(), + ); + } + if args.shape.causal { + run_mask::(args) + } else { + run_mask::(args) + } +} + +/// Select runtime or static dimensions for the chosen causal-mask specialization. +fn run_mask(args: Args) -> Result<()> { + match args.shape_mode { + ShapeMode::Runtime => run_shape::(args), + ShapeMode::Static => run_shape::(args), + } +} + +/// Verify the kernel and replay, collect CUDA-event samples, and write the report. +fn run_shape( + args: Args, +) -> Result<()> { + let shape = args.shape; + let (q, k, v) = make_inputs(shape, args.input_scale, args.zero_query); + let expected = reference(shape, &q, &k, &v); + let ctx = CudaContext::new(args.device)?; + let capability = ctx.compute_capability()?; + if capability.0 < 12 { + return Err("this Rust specialization requires SM120 or newer".into()); + } + let stream = ctx.default_stream(); + // All AOT mask/shape specializations live in this package's CUTLASS CUBIN. + let image = cuda_host::load_embedded_module(&ctx, env!("CARGO_PKG_NAME"))?; + // SAFETY: this bundle contains the typed kernels and their launch contracts. + let module = unsafe { kernel::kernels::from_module(image) }?; + let q_dev = DeviceBuffer::from_host(&stream, &q)?; + let k_dev = DeviceBuffer::from_host(&stream, &k)?; + let v_dev = DeviceBuffer::from_host(&stream, &v)?; + let poison = vec![f16::NAN; shape.q_elements()]; + let mut output = DeviceBuffer::from_host(&stream, &poison)?; + let blocks = (shape.seq_q.div_ceil(64) * shape.heads) as u32; + let launch = module.prepare_fmha_prefill::(LaunchConfig1D::new( + blocks, + THREADS, + SHARED_BYTES, + ))?; + let registers = launch.function().num_registers()?; + let run = |output: &mut DeviceBuffer| -> Result<()> { + // Q/K/V have the checked contiguous BSHD extents; each CTA + // owns one disjoint 64-row/head output tile. All allocations remain + // live until synchronization, and the launch reserves all shared data. + module.fmha_prefill::( + &stream, + &launch, + &q_dev, + &k_dev, + &v_dev, + output, + shape.seq_q as u32, + shape.seq_k as u32, + shape.heads as u32, + )?; + Ok(()) + }; + run(&mut output)?; + let checked = output.to_host_vec(&stream)?; + let validation = verify(&checked, &expected)?; + let checked_bits: Vec = checked.iter().map(|x| x.to_bits()).collect(); + // Re-poison before replay so incomplete writes cannot preserve a pass. + output.copy_from_host(&stream, &poison)?; + run(&mut output)?; + let replay = output.to_host_vec(&stream)?; + verify(&replay, &expected)?; + if !replay + .iter() + .map(|x| x.to_bits()) + .eq(checked_bits.iter().copied()) + { + return Err("output changed across identical launches".into()); + } + let warmup = if args.iters == 0 { 0 } else { args.warmup }; + for _ in 0..warmup { + run(&mut output)?; + } + stream.synchronize()?; + let events = (0..args.iters) + .map(|_| { + Ok(( + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + )) + }) + .collect::>>()?; + let mut samples = Vec::with_capacity(args.iters); + for (start, end) in &events { + start.record(&stream)?; + run(&mut output)?; + end.record(&stream)?; + end.synchronize()?; + let elapsed = f64::from(start.elapsed_ms(end)?) * 1000.0; + if !elapsed.is_finite() || elapsed <= 0.0 { + return Err(format!("invalid event timing: {elapsed}").into()); + } + samples.push(elapsed); + } + let after = output.to_host_vec(&stream)?; + verify(&after, &expected)?; + if !after + .iter() + .map(|x| x.to_bits()) + .eq(checked_bits.iter().copied()) + { + return Err("output changed during timing".into()); + } + if let Some(path) = args.output { + let bytes: Vec = after + .iter() + .flat_map(|value| value.to_bits().to_le_bytes()) + .collect(); + std::fs::write(path, bytes)?; + } + let mut sorted = samples.clone(); + sorted.sort_by(f64::total_cmp); + let median = if sorted.is_empty() { + None + } else { + Some((sorted[(sorted.len() - 1) / 2] + sorted[sorted.len() / 2]) * 0.5) + }; + let report = json!({ + "implementation": "cute-rs", "workload": "prefill-fmha", "shape": [1, shape.seq_q, shape.seq_k, shape.heads, HEAD_DIM], + "causal": shape.causal, "input_scale": args.input_scale, "zero_query": args.zero_query, + "input_hashes": {"q": format!("{:016x}", bits_hash(&q)), "k": format!("{:016x}", bits_hash(&k)), "v": format!("{:016x}", bits_hash(&v))}, + "output_hash": format!("{:016x}", bits_hash(&after)), "verification": "cpu-pass", + "validation": {"elements": validation.elements, "max_abs_error": validation.max_abs_error, + "max_tolerance_ratio": validation.max_tolerance_ratio, "atol": fixture::ATOL, "rtol": fixture::RTOL, + "oracle": "cpu-fp64-stable-softmax"}, + "replay_bitwise_equal": true, "metric_value": median, "metric_unit": "us", "samples_us": samples, + "provenance": {"gpu": ctx.device_name()?, "compute_capability": [capability.0, capability.1], + "device": args.device, "tool": "cuda-events-single-kernel", "timing_mode": "direct", + "warmup": warmup, "iters": args.iters, "grid": [blocks, 1, 1], "threads": THREADS, "shared_memory_bytes": SHARED_BYTES, + "dtype": "float16", "registers_per_thread": registers, + "causal_specialization": "static", + "shape_specialization": if args.shape_mode == ShapeMode::Static { "static" } else { "runtime" }} + }); + let rendered = serde_json::to_string_pretty(&report)?; + if let Some(path) = args.json { + std::fs::write(path, format!("{rendered}\n"))?; + } + println!("{rendered}"); + Ok(()) +} + +/// Tests for prefill CLI modes and exact static-shape dispatch. +#[cfg(test)] +mod tests { + use super::*; + + /// Parse borrowed CLI words without reading process arguments. + fn parse(values: &[&str]) -> Result { + Args::parse_from(values.iter().map(|s| s.to_string())) + } + + /// Check that raw output is standalone-only and Python selection is comparison-only. + #[test] + fn comparison_cli_requires_standalone_mode_for_output() { + assert!(parse(&["--python", "/tmp/python"]).is_ok()); + assert!(parse(&["--output", "o.f16"]).is_err()); + assert!(parse(&["--rust-only", "--output", "o.f16"]).is_ok()); + assert!(parse(&["--rust-only", "--python", "/tmp/python"]).is_err()); + } + + /// Keep runtime shape selection as the default, including asymmetric noncausal shapes. + #[test] + fn runtime_shapes_remain_the_default() { + assert_eq!(parse(&[]).unwrap().shape_mode, ShapeMode::Runtime); + assert!(parse(&["--seq-q", "64", "--seq-k", "128", "--noncausal"]).is_ok()); + } + + /// Accept either mask and shape mode when requesting the full comparison matrix. + #[test] + fn all_cases_accepts_both_selected_masks_and_shape_modes() { + for mode in ["runtime", "static"] { + for mask in ["--causal", "--noncausal"] { + assert!(parse(&["--all-cases", "--shape-mode", mode, mask]).is_ok()); + } + } + } + + /// Reject standalone-only options in a multi-case comparison. + #[test] + fn all_cases_rejects_standalone_execution_and_output() { + for options in [ + vec!["--all-cases", "--rust-only"], + vec!["--rust-only", "--all-cases"], + vec!["--all-cases", "--output", "o.f16"], + ] { + assert!(parse(&options).is_err()); + } + } + + /// Require the exact static shape for the runtime-versus-static comparison matrix. + #[test] + fn all_cases_requires_dimensions_supported_by_static_kernels() { + for (dimension, value) in [("--seq-q", "64"), ("--seq-k", "64"), ("--heads", "3")] { + assert!(parse(&["--all-cases", "--noncausal", dimension, value]).is_err()); + } + } + + /// Check static dispatch for both causal and noncausal masks. + #[test] + fn static_dispatch_accepts_both_masks_for_its_exact_shape() { + for mask in ["--causal", "--noncausal"] { + let args = parse(&["--shape-mode", "static", mask]).unwrap(); + assert_eq!(args.shape_mode, ShapeMode::Static); + } + } + + /// Reject static shapes with a mismatched query length, KV length, or head count. + #[test] + fn static_dispatch_rejects_every_mismatched_dimension() { + for (dimension, value) in [("--seq-q", "64"), ("--seq-k", "64"), ("--heads", "3")] { + assert!(parse(&["--shape-mode", "static", "--noncausal", dimension, value]).is_err()); + } + assert!(parse(&["--shape-mode", "unknown"]).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/test_default_compare.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/test_default_compare.py new file mode 100644 index 0000000000..9eb95e3e16 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/test_default_compare.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only checks for the default FMHA comparison and its runner protocol.""" + +import contextlib +import importlib.util +import io +import json +from pathlib import Path +import subprocess +import tempfile +import unittest +from unittest import mock + + +HERE = Path(__file__).resolve().parent + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +fixture = load("default_compare_fixture", HERE / "reference/launch.py") +comparison = load("default_compare", HERE / "compare.py") if (HERE / "compare.py").exists() else None + + +class DefaultComparisonTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(comparison, "default FMHA comparison is not implemented") + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + self.binary = self.directory / "rust" + self.binary.write_text("test runner") + self.binary.chmod(0o755) + self.mutate = lambda report: None + + def run_runner(self, command, **kwargs): + rust = command[0] == str(self.binary) + if rust: + self.assertEqual(command[1], "--rust-only", "Rust subprocess must not recurse") + cli = command[2:] + else: + self.assertNotIn("--rust-only", command) + cli = command[2:] + launch = (HERE if self.workload == "prefill" else HERE.parent / "fmha_decode_sm120a_cute") / "reference/launch.py" + args = load("test_reference_runner", launch).parse_args(cli) + self.assertEqual((args.seq_q, args.seq_k, args.heads, args.input_scale, args.zero_query), + (1, 64 if self.workload == "prefill" else 65, 1, 4, True)) + inputs = fixture.make_inputs(1, args.seq_k, 1, 4, True) + expected = fixture.reference_cpu(*inputs, 1, args.seq_k, 1, False) + output = fixture.decode_f16(fixture.encode_f16(expected)) + samples = ([3.0, 1.0, 2.0] if rust else [6.0, 2.0, 4.0]) if args.iters else [] + report = { + "implementation": "cute-rs" if rust else "cutlass-cute-dsl", + "workload": args.workload, "shape": [1, 1, args.seq_k, 1, 64], + "causal": False, "input_scale": 4, "zero_query": True, + "input_hashes": dict(zip("qkv", map(fixture.bits_hash, inputs))), + "output_hash": fixture.bits_hash(output), "verification": "cpu-pass", + "replay_bitwise_equal": True, "validation": fixture.verify(output, expected), + "metric_value": (2.0 if rust else 4.0) if samples else None, + "metric_unit": "us", "samples_us": samples, + "provenance": { + "gpu": "test GPU", "compute_capability": [12, 0], "device": args.device, + "warmup": args.warmup if args.iters else 0, "iters": args.iters, + "dtype": "float16", "timing_mode": "direct", + "tool": "cuda-events-single-kernel" if self.workload == "prefill" else "cuda-events-kernel-sequence", + "shape_specialization": "runtime", "causal_specialization": "static", + "algorithm": "single-query-split-kv", "kv_chunk": 64, + "split_count": 2, "kernel_count": 2, + "kernel_sequence": ["fmha_decode_partial", "fmha_decode_merge"], + }, + } + if rust: + self.mutate(report) + args.json.write_text(json.dumps(report)) + args.output.write_bytes(fixture.encode_f16(output)) + return subprocess.CompletedProcess(command, 0, stdout="runner diagnostics\n", stderr="") + + def invoke(self, workload, iters=3): + self.workload = workload + result_path = self.directory / "combined.json" + argv = ["--workload", workload, "--rust-binary", str(self.binary), + "--seq-q", "1", "--seq-k", "64" if workload == "prefill" else "65", + "--heads", "1", "--noncausal", "--input-scale", "4", "--zero-query", + "--device", "2", "--warmup", "2", "--iters", str(iters), + "--json", str(result_path)] + stdout, stderr = io.StringIO(), io.StringIO() + with mock.patch.object(comparison.subprocess, "run", side_effect=self.run_runner), \ + contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + status = comparison.main(argv) + return status, stdout.getvalue(), stderr.getvalue(), result_path + + def test_prints_both_medians_and_saves_checked_reports_for_each_workload(self): + for workload in ("prefill", "decode"): + with self.subTest(workload=workload): + status, stdout, _, path = self.invoke(workload) + self.assertEqual(status, 0) + self.assertIn("CUDA-event median", stdout) + self.assertIn("runtime / noncausal", stdout) + self.assertIn("CuTeDSL: 4.000 us\ncute-rs: 2.000 us", stdout) + report = json.loads(path.read_text()) + self.assertEqual(report["verification"], "both-pass") + self.assertEqual(report["cuda_oxide"]["metric_value"], 2) + self.assertEqual(report["python_reference"]["provenance"]["device"], 2) + self.assertEqual(report["output_comparison"]["elements"], 64) + + def test_zero_iterations_verifies_both_without_inventing_timings(self): + for workload in ("prefill", "decode"): + status, stdout, _, path = self.invoke(workload, iters=0) + self.assertEqual(status, 0) + self.assertIn("CuTeDSL: verified (timing disabled)\n" + "cute-rs: verified (timing disabled)", stdout) + self.assertIsNone(json.loads(path.read_text())["cuda_oxide"]["metric_value"]) + + def test_rejects_invalid_timing_or_mismatched_device_before_printing_summary(self): + for field, value in (("metric_value", 999), ("samples_us", [0, 1, 2]), + ("metric_unit", "ms"), ("input_hashes", {})): + with self.subTest(field=field): + self.mutate = lambda report: report.update({field: value}) + status, stdout, stderr, _ = self.invoke("prefill") + self.assertEqual(status, 1) + self.assertEqual(stdout, "") + self.assertIn("FAIL", stderr) + self.mutate = lambda report: report["provenance"].update(gpu="different GPU") + status, stdout, _, _ = self.invoke("decode") + self.assertEqual(status, 1) + self.assertEqual(stdout, "") + + def test_child_failure_surfaces_diagnostics_without_success_summary(self): + self.run_runner = lambda *args, **kwargs: subprocess.CompletedProcess(args[0], 1, "", "missing GPU dependency") + status, stdout, stderr, _ = self.invoke("prefill") + self.assertEqual(status, 1) + self.assertEqual(stdout, "") + self.assertIn("missing GPU dependency", stderr) + + +class AllCasesTests(unittest.TestCase): + workload = "prefill" + dimensions = (1024, 1024, 8) + selection = ["--shape-mode", "static", "--noncausal"] + expected_cases = [("runtime", True), ("runtime", False), ("static", True), ("static", False)] + unsupported_shapes = [("--shape-mode", "runtime", "--seq-q", "64", "--seq-k", "64"), + ("--shape-mode", "runtime", "--heads", "3")] + + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.path = Path(self.temporary.name) / "all.json" + self.seen = [] + + def compare_case(self, options, cli): + _, launch, _ = comparison.load_reference(self.workload) + args = launch.parse_args(cli) + mode = args.shape_mode if self.workload == "prefill" else args.specialization + self.seen.append((mode, args.causal)) + # The selected mode/mask must be replaced without losing other settings. + self.assertEqual((args.seq_q, args.seq_k, args.heads), self.dimensions) + self.assertEqual((args.device, args.warmup, args.input_scale, args.zero_query), (2, 7, 4, True)) + return { + "workload": args.workload, "verification": "both-pass", + "cuda_oxide": {"metric_value": 2.0 if args.iters else None, + "shape": [1, *self.dimensions, 64], "causal": args.causal, + "provenance": {"warmup": args.warmup, "iters": args.iters, + "shape_specialization": mode}}, + "python_reference": {"metric_value": 4.0 if args.iters else None}, + "timing": {"mode": "direct", "statistic": "median", "unit": "us"}, + } + + def invoke(self, *extra, failure=None): + argv = ["--workload", self.workload, "--rust-binary", "/unused-rust", + "--all-cases", *self.selection, + "--device", "2", "--warmup", "7", "--iters", "3", + "--input-scale", "4", "--zero-query", "--json", str(self.path), *extra] + stdout, stderr = io.StringIO(), io.StringIO() + with mock.patch.object(comparison, "compare", side_effect=failure or self.compare_case) as runner, \ + contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + status = comparison.main(argv) + return status, stdout.getvalue(), stderr.getvalue(), runner + + def test_all_modes_are_measured_and_saved_despite_single_case_flags(self): + status, stdout, _, runner = self.invoke() + self.assertEqual(status, 0) + self.assertEqual(runner.call_count, len(self.expected_cases)) + expected = self.expected_cases + self.assertEqual(self.seen, expected) + report = json.loads(self.path.read_text()) + self.assertEqual(report["verification"], "all-pass") + self.assertEqual([(case["shape_mode"], case["causal"]) for case in report["cases"]], expected) + self.assertIn("CUDA-event median", stdout) + self.assertIn("Nsight kernel-active", stdout) + for mode, causal in expected: + self.assertIn(f"| {mode.title()} | {'Causal' if causal else 'Noncausal'} | 2.000 | 4.000 | -50.0% |", stdout) + + def test_verification_only_table_has_no_fabricated_times_or_relative_deltas(self): + status, stdout, _, _ = self.invoke("--iters", "0") + self.assertEqual(status, 0) + self.assertEqual(stdout.count("| verified | verified | n/a |"), len(self.expected_cases)) + self.assertNotIn("-50.0%", stdout) + + def test_unsupported_shape_fails_before_running_any_case(self): + for extra in self.unsupported_shapes: + with self.subTest(extra=extra): + status, stdout, stderr, runner = self.invoke(*extra) + self.assertEqual(status, 1) + runner.assert_not_called() + self.assertEqual(stdout, "") + self.assertIn("--all-cases", stderr) + self.assertFalse(self.path.exists()) + + def test_failed_case_cannot_produce_success_table_or_json(self): + def fail_second(options, cli): + if self.seen: + raise RuntimeError("second case failed validation") + return self.compare_case(options, cli) + + status, stdout, stderr, runner = self.invoke(failure=fail_second) + self.assertEqual(status, 1) + self.assertEqual(runner.call_count, 2) + self.assertEqual(stdout, "") + self.assertIn("failed validation", stderr) + self.assertFalse(self.path.exists()) + + +class DecodeAllCasesTests(AllCasesTests): + workload = "decode" + dimensions = (1, 1024, 8) + selection = ["--specialization", "auto", "--noncausal"] + expected_cases = [("runtime", False), ("static", False)] + unsupported_shapes = [("--specialization", "runtime", "--seq-k", "66"), + ("--specialization", "runtime", "--heads", "3")] + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/validate.py b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/validate.py new file mode 100644 index 0000000000..e2efb8d2a7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/validate.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate all outputs for multiple tiles, masks, heads, and logit ranges. + +Select the intended GPU through CUDA_VISIBLE_DEVICES. No timing is performed. +--reference also runs the unchanged public CuTeDSL kernel on identical inputs. +""" + +import argparse +import hashlib +import json +import math +from pathlib import Path +import subprocess +import sys +import tempfile + +sys.path.insert(0, str(Path(__file__).parent / "reference")) +from launch import decode_f16, verify + +CASES = ( + ("causal-one-tile", 64, 64, 2, True, 1, False), + ("noncausal-one-tile", 64, 64, 2, False, 1, False), + ("causal-two-tiles", 128, 128, 8, True, 1, False), + ("noncausal-two-tiles", 128, 128, 8, False, 1, False), + ("causal-three-tiles", 192, 192, 3, True, 1, False), + ("noncausal-single-query", 1, 192, 3, False, 1, False), + ("more-kv-tiles", 64, 128, 3, False, 1, False), + ("more-query-tiles", 128, 64, 3, False, 1, False), + ("causal-prefill", 1024, 1024, 8, True, 1, False), + ("noncausal-fmha", 1024, 1024, 8, False, 1, False), + ("causal-large-logits", 128, 128, 2, True, 16, False), + ("noncausal-large-logits", 128, 128, 2, False, 16, False), + ("moderate-logits", 64, 128, 2, False, 4, False), + ("causal-uniform-softmax", 128, 128, 3, True, 1, True), + ("noncausal-uniform-softmax", 128, 128, 3, False, 1, True), +) + + +def checked_report(path, seq_q, seq_k, heads, causal, scale, zero_query): + report = json.loads(path.read_text()) + expected = {"shape": [1, seq_q, seq_k, heads, 64], "causal": causal, + "input_scale": scale, "zero_query": zero_query, + "verification": "cpu-pass", "replay_bitwise_equal": True} + for key, value in expected.items(): + if report.get(key) != value: + raise ValueError(f"wrong report {key}: {report.get(key)!r} != {value!r}") + validation = report["validation"] + if validation["elements"] != seq_q * heads * 64: + raise ValueError("report does not cover every output") + for field in ("max_abs_error", "max_tolerance_ratio"): + if not math.isfinite(validation[field]) or validation[field] < 0: + raise ValueError(f"invalid {field}") + if validation["max_tolerance_ratio"] > 1: + raise ValueError("verification tolerance exceeded") + return report + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path, required=True) + parser.add_argument("--reference", action="store_true") + parser.add_argument("--shape-mode", choices=("runtime", "static"), default="runtime", + help="static checks the supported 1024-token, 8-head specialization") + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--timeout", type=float, default=300) + args = parser.parse_args() + binary = args.binary.resolve() + directory = args.output_dir or Path(tempfile.mkdtemp(prefix="fmha-prefill-validation-")) + directory.mkdir(parents=True, exist_ok=True) + summary = {"binary": str(binary), "binary_sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + "status": "running", "checks": []} + report_path = directory / "report.json" + cases = CASES + if args.shape_mode == "static": + cases = ( + ("static-causal", 1024, 1024, 8, True, 1, False), + ("static-noncausal", 1024, 1024, 8, False, 1, False), + ("static-causal-large-logits", 1024, 1024, 8, True, 16, False), + ("static-noncausal-uniform", 1024, 1024, 8, False, 1, True), + ) + summary["shape_mode"] = args.shape_mode + try: + for name, sq, sk, heads, causal, scale, zero_query in cases: + common = ["--seq-q", str(sq), "--seq-k", str(sk), "--heads", str(heads), + "--causal" if causal else "--noncausal", "--input-scale", str(scale), + "--warmup", "0", "--iters", "0"] + common += ["--shape-mode", args.shape_mode] + if zero_query: + common.append("--zero-query") + implementations = [("rust", [str(binary), "--rust-only"])] + if args.reference: + script = Path(__file__).parent / "reference" / "launch.py" + implementations.append(("cutlass", [args.python, str(script)])) + reports, outputs = {}, {} + for label, prefix in implementations: + result_path = directory / f"{name}-{label}.json" + output_path = directory / f"{name}-{label}.f16" + result_path.unlink(missing_ok=True) + output_path.unlink(missing_ok=True) + command = [*prefix, *common, "--json", str(result_path), "--output", str(output_path)] + entry = {"name": name, "implementation": label, "command": command, "status": "running"} + summary["checks"].append(entry) + try: + completed = subprocess.run(command, capture_output=True, text=True, timeout=args.timeout) + (directory / f"{name}-{label}.log").write_text(completed.stdout + completed.stderr) + except subprocess.TimeoutExpired as error: + entry["status"] = "timeout" + raise RuntimeError(f"{name}/{label} exceeded {args.timeout:g} seconds") from error + entry["exit_code"] = completed.returncode + if completed.returncode: + entry["status"] = "failed" + raise RuntimeError(f"{name}/{label} exited {completed.returncode}; see its log") + reports[label] = checked_report(result_path, sq, sk, heads, causal, scale, zero_query) + if reports[label].get("provenance", {}).get("shape_specialization") != args.shape_mode: + raise ValueError(f"wrong shape specialization for {name}/{label}") + outputs[label] = decode_f16(output_path.read_bytes()) + if len(outputs[label]) != sq * heads * 64 or not all(map(math.isfinite, outputs[label])): + raise ValueError(f"invalid or incomplete output for {name}/{label}") + entry["status"] = "passed" + if args.reference: + if reports["rust"]["input_hashes"] != reports["cutlass"]["input_hashes"]: + raise ValueError(f"Rust/CuTeDSL input mismatch for {name}") + parity = verify(outputs["rust"], outputs["cutlass"]) + summary["checks"].append({"name": name, "implementation": "cross-check", "status": "passed", **parity}) + report_path.write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n") + print(f"PASS {name}", flush=True) + summary["status"] = "passed" + except (OSError, ValueError, KeyError, TypeError, RuntimeError) as error: + summary["status"] = "failed" + summary["error"] = str(error) + print(f"FAIL: {error}", file=sys.stderr) + report_path.write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n") + print(f"Validation report: {report_path}") + return 0 if summary["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock similarity index 95% rename from cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock rename to cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock index 0df03ea524..f88a780692 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock @@ -55,6 +55,26 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "cexpr" version = "0.6.0" @@ -115,6 +135,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "thiserror", +] + [[package]] name = "cuda-artifact-finalizer" version = "0.2.1" @@ -128,7 +157,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,18 +169,20 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", + "bytemuck", "cuda-bindings", "cuda-core-derive", + "cuda-repr", "half", "oxide-artifacts", ] [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -162,7 +193,9 @@ dependencies = [ name = "cuda-device" version = "0.2.1" dependencies = [ + "bytemuck", "cuda-macros", + "cuda-repr", ] [[package]] @@ -171,6 +204,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -191,6 +225,10 @@ dependencies = [ "syn", ] +[[package]] +name = "cuda-repr" +version = "0.5.0" + [[package]] name = "cuda-target-spec" version = "0.1.0" @@ -233,6 +271,17 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "fp16_gemm_256x352_ctm" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "serde_json", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -255,6 +304,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if", "crunchy", "zerocopy", @@ -562,15 +612,6 @@ dependencies = [ "digest", ] -[[package]] -name = "shared_ptr_api" -version = "0.1.0" -dependencies = [ - "cuda-core", - "cuda-device", - "cuda-host", -] - [[package]] name = "shlex" version = "1.3.0" diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.toml new file mode 100644 index 0000000000..759d89fb8c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "fp16_gemm_256x352_ctm" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "SM100 two-CTA FP16 GEMM using the typed Rust primitives API" + +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core" } +serde_json = "1" diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/README.md b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/README.md new file mode 100644 index 0000000000..b4d6d232dc --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/README.md @@ -0,0 +1,84 @@ +# FP16 GEMM with Rust CTM — SM100a (B200) + +This example computes a matrix product with optional row bias: + +```text +C[M, N] = A[M, K] @ B[N, K].T + bias[M, 1] +``` + +Inputs and output are FP16; accumulation and bias addition use FP32 before +rounding to FP16. M must be a positive multiple of **256**, and N/K of **8**. +The `256x352` name describes an output tile; `--mnk` sets the full matrix size. + +## Tile flow + +Two CTAs cooperate on each tile, splitting its 352 columns into 192 and 160: + +```text +A/B global → TMA → shared [5 stages] → paired MMA → TMEM +C global ← TMA ← shared [2 stages] ← FP16 + optional row bias +``` + +`ctm-rs` expresses this schedule through typed instructions and physical formats, +with explicit barriers and storage lifetimes. The Rust kernel uses no CuTe +operations or layout algebra. The [CuTe version](../fp16_gemm_256x352_cute/README.md) +expresses the same schedule with typed layouts, tensor operations and pipelines. + +## Run it + +You need an SM100 GPU and a compatible CUDA driver. The comparison uses the +checked-in [Python CTM reference](../fp16_gemm_256x352_cute/public_reference/README.md): +CUTLASS instruction primitives with CuTe compilation and epilogue helpers. +From `cuda-oxide/`, once per checkout: + +```bash +cargo oxide toolchain install cutlass +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Then: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_MLIR_PROFILE=cutlass-primitives-mlir22 \ + cargo oxide build fp16_gemm_256x352_ctm --arch sm_100a + +cute-rs-preview/.venv/bin/python \ + cute-rs-preview/examples/fp16_gemm_256x352_cute/compare.py \ + --rust-bin cute-rs-preview/examples/fp16_gemm_256x352_ctm/target/release/fp16_gemm_256x352_ctm \ + --mnk 8192,8096,8192 --warmup 2000 --iters 301 --json /tmp/ctm-gemm.json +``` + +Both outputs are checked against an independent PyTorch oracle on identical +inputs. The command ends with: + +```text +CUTLASS (CuTe + CTM): + +The Python reference has an unresolved group-two TMA memcheck diagnostic; +see [validation and known limitation](public_reference/README.md#validation-and-known-limitation). + +- [Layouts, pipeline, descriptor ABI and Rust API](public_reference/architecture.md) +- [Python setup, timing options and validation commands](public_reference/README.md) +- [Public reference provenance](public_reference/SOURCE.md) diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/compare.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/compare.py new file mode 100644 index 0000000000..78650e2f17 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/compare.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Verify and time Rust and the public Python CTM FP16 reference on identical tensors.""" + +import argparse +import hashlib +import importlib +import json +import os +from pathlib import Path +import socket +import statistics +import subprocess +import sys +import tempfile + + +SETUP_HINT = ("from cuda-oxide/ run `python3 -m venv cute-rs-preview/.venv && " + "cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt`, " + "then rerun; or set CUDA_OXIDE_CUTE_PYTHON to an environment that has them") + + +def shape(value): + """Parse M,N,K and enforce the device launch's dimension and index bounds.""" + try: + dims = tuple(int(part) for part in value.split(",")) + except ValueError as exc: + raise argparse.ArgumentTypeError("expected M,N,K integers") from exc + if len(dims) != 3 or any(v <= 0 for v in dims): + raise argparse.ArgumentTypeError("expected three positive M,N,K integers") + if dims[0] % 256 or dims[1] % 8 or dims[2] % 8: + raise argparse.ArgumentTypeError("M must be divisible by 256; N and K by 8") + i32_max = (1 << 31) - 1 + m, n, k = dims + ntiles = (n + 351) // 352 + if max(dims) > i32_max or (ntiles - 1) * 352 + 320 > i32_max: + raise argparse.ArgumentTypeError("dimensions and padded TMA coordinates must fit i32") + if (m // 256) * ntiles > i32_max // 2: + raise argparse.ArgumentTypeError("two-CTA problem grid exceeds the CUDA grid-x limit") + if max(m*k, n*k, m*n) > sys.maxsize // 2: + raise argparse.ArgumentTypeError("FP16 allocation exceeds the host addressable limit") + return dims + + +def measure(torch, oracle, kernel, inputs, output, warmup, iters, graph_launches): + """Time a prepared callable with one launch per event pair, then verify it.""" + if graph_launches != 1: + raise ValueError("comparison requires one direct kernel launch per sample") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + # Keep initialization, launches, verification and events on one explicit + # stream. The prepared callable honors the supplied stream argument. + with torch.cuda.stream(stream): + for _ in range(warmup): + kernel(*inputs, output, stream=stream) + stream.synchronize() + # Verify the timed output rather than adding a correctness launch. + # Poison after warmup so a no-op sample cannot reuse warmup's result. + output.fill_(float("nan")) + stream.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for start, end in zip(starts, ends): + start.record(stream) + kernel(*inputs, output, stream=stream) + end.record(stream) + end.synchronize() + stream.synchronize() + oracle.verify("fp16", *inputs, output) + return statistics.median(s.elapsed_time(e) * 1000.0 for s, e in zip(starts, ends)) + + +def main(): + """Compare both implementations using serialized copies of the same inputs.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rust-bin", type=Path, required=True, help="compiled Rust host executable") + parser.add_argument("--mnk", type=shape, default=(256, 352, 64)) + parser.add_argument("--has-bias", "--has_bias", action="store_true") + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--timing-mode", choices=("direct",), default="direct", + help="one kernel per CUDA-event sample; use --rust-only for graph timing") + parser.add_argument("--graph-launches", type=int, default=1, help="must be 1 for comparison") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--clusters", type=int, help="override Rust persistent cluster count") + parser.add_argument("--json", type=Path, help="write combined report") + args = parser.parse_args() + if args.warmup < 0 or args.iters < 1 or args.graph_launches < 1 or (args.clusters is not None and args.clusters < 1): + parser.error("warmup must be nonnegative; iters, graph-launches and clusters must be positive") + if args.graph_launches != 1: + parser.error("comparison requires direct timing and --graph-launches 1; use --rust-only for graph timing") + binary = args.rust_bin.resolve(strict=True) + reference_dir = Path(__file__).resolve().parent / "public_reference" + try: + import torch + import numpy as np + from public_reference import oracle, gemm + except ImportError as exc: + sys.exit(f"comparison dependencies could not be imported by {sys.executable}: {exc}; {SETUP_HINT}") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + raise RuntimeError("this comparison requires an SM100 GPU") + # Both processes select device ordinal zero. CUDA_VISIBLE_DEVICES is the + # common way to choose a physical GPU without mixing devices in one run. + torch.cuda.set_device(0) + inputs, outputs = oracle.prepare_tensors("fp16", mnk=args.mnk, has_bias=args.has_bias, seed=args.seed) + python_output = outputs[0] + # Eager compilation stays outside event timing even with --warmup 0. + callable_kernel = gemm.prepare_callable( + has_bias=args.has_bias, inputs=inputs, output=python_output + ) + python_capacity = int(gemm.cluster_capacity()) + python_clusters = min((args.mnk[0] // 256) * ((args.mnk[1] + 351) // 352), python_capacity) + python_us = measure(torch, oracle, callable_kernel, inputs, python_output, + args.warmup, args.iters, args.graph_launches) + m, n, k = args.mnk + with tempfile.TemporaryDirectory(prefix="fp16-gemm-compare-") as temporary: + directory = Path(temporary) + # The oracle's B is the [1,K,N] view; the kernel stores [N,K]. + tensors = {"a": inputs[0][0], "b": inputs[1][0].transpose(0, 1)} + if args.has_bias: + tensors["bias"] = inputs[2] + digests = {} + for name, tensor in tensors.items(): + raw = tensor.detach().contiguous().cpu().numpy().astype(", 0, RowMajor<128, 64>>; +type B0Smem = ComposedLayout, 0, RowMajor<96, 64>>; +type B1Smem = ComposedLayout, 0, RowMajor<80, 64>>; +type TiledMma = Sm100TiledMma; +type MainloopPipeline = PipelineTmaUmma<5, 77_824>; +``` + +All four typed TMA descriptors use `#[grid_constant]` parameters. The kernel +takes typed references: + +```rust +#[grid_constant] a: &TensorMap +``` + +The generated host launcher accepts each descriptor **by value**, passing +128 bytes directly without allocating a descriptor buffer in device global +memory: + +```text +host descriptor -> 128 launch-argument bytes -> shared read-only descriptor +``` + +Each thread borrows the same descriptor storage. Element and layout types stay +explicit in the signature. Tensor data stays in device memory and must remain +live until GPU work completes; the unsafe host launch makes this caller +obligation explicit. + +## Compiler flow and source map + +The compiler preserves layout and pipeline information through verification +and lowering: + +![The cute-rs stack from Rust types through verification to a CUDA cubin](../../../../cuda-oxide-book/cute/images/cute-rs-preview-stack-and-flow.svg) + +- [SM100 APIs](../../../cute-rs/src/sm100/mod.rs): `Sm100SharedTile`, + `Sm100ClusterTmaCopy`, `Sm100TiledMma`, `Sm100TmemEpilogue`, and pipelines. +- [Host setup](../src/main.rs): typed TMA descriptors, cluster launch, and validation. +- [Dialect operations](../../../dialect-cute/src/ops/sm100/mod.rs): layout, + collector, and pipeline verification before lowering. +- [Example README](../README.md): computation, run commands, and performance. diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/gemm.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/gemm.py new file mode 100644 index 0000000000..0c12fa43de --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/gemm.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Original SM100 FP16 GEMM reference for cuda-oxide. + +The two CTAs compute 256x352 outputs using 256x192 and 256x160 MMAs. +This module uses the public, experimental CUTLASS Python DSL 4.7.0 APIs. +Inputs and output remain owned by the prepared callable until it is released. +See the sibling documentation for sources and the hardware validation status. +""" + +import functools +import importlib.metadata +import operator + +import cutlass +from cutlass import cute, utils +from cutlass.cute.runtime import from_dlpack, make_fake_stream, make_fake_tensor +from cutlass.experimental import cuda, primitives as prim +from cuda.bindings import driver + +# Byte offsets describe disjoint regions in each CTA's shared allocation. +_STAGES = 5 +_A_BYTES = 128 * 64 * 2 +_B0_BYTES = 96 * 64 * 2 +_B1_BYTES = 80 * 64 * 2 +_B_OFFSET = _STAGES * _A_BYTES +_C_OFFSET = _B_OFFSET + _STAGES * (_B0_BYTES + _B1_BYTES) +_C_BYTES = 128 * 32 * 2 +_CONTROL = _C_OFFSET + 2 * _C_BYTES +SMEM_BYTES = _CONTROL + 128 +assert SMEM_BYTES == 211072 +_GROUP = prim.CTAGroup.CTA_2 + + +@cute.jit +def _raw(ptr: cute.Pointer): + """Adapt a CuTe pointer through public APIs for experimental primitives.""" + return cutlass.Pointer(ptr.to_llvm_ptr(), dtype=ptr.dtype, space=ptr.memspace) + + +@cute.jit +def _wait(bar: cutlass.Pointer, phase: cutlass.Int32, cluster: cutlass.Constexpr = False): + """Acquire a completed barrier phase; remote epilogue arrivals need cluster scope.""" + if cutlass.const_expr(cluster): + while not prim.mbarrier_try_wait_parity( + bar, phase, scope=prim.MBarrierScope.CLUSTER, order=prim.MemOrder.ACQUIRE + ): + pass + else: + while not prim.mbarrier_try_wait_parity(bar, phase): + pass + + +@cute.jit +def _operand(base: cute.Pointer): + """Describe a 64-wide K-major FP16 tile, with 128-byte swizzle and K16 steps.""" + return prim.Tcgen05SmemDesc.build( + _raw(base), leading_byte_offset=16, stride_byte_offset=8 * 64 * 2, + layout=prim.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + + +class _Gemm: + """Static launch configuration; descriptors derive their bounds from tensor views.""" + + def __init__(self, has_bias, clusters): + self.has_bias = has_bias + self.clusters = clusters + + @cute.jit + def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, + bias: cute.Tensor, stream: driver.CUstream): + """Build four checked tensor maps and enqueue exactly one clustered kernel.""" + ad = cuda.create_tensor_map_tiled_from_view( + a, (128, 64), stride_order=(1, 0), swizzle=cuda.TensorMapSwizzle.s128b + ) + bd0 = cuda.create_tensor_map_tiled_from_view( + b, (96, 64), stride_order=(1, 0), swizzle=cuda.TensorMapSwizzle.s128b + ) + bd1 = cuda.create_tensor_map_tiled_from_view( + b, (80, 64), stride_order=(1, 0), swizzle=cuda.TensorMapSwizzle.s128b + ) + cd = cuda.create_tensor_map_tiled_from_view( + c, (128, 32), stride_order=(1, 0), swizzle=cuda.TensorMapSwizzle.s64b + ) + assert ad.shared_storage_bytes == _A_BYTES + assert bd0.shared_storage_bytes == _B0_BYTES + assert bd1.shared_storage_bytes == _B1_BYTES + assert cd.shared_storage_bytes == _C_BYTES + self.kernel(ad, bd0, bd1, cd, bias, a.shape[0], c.shape[1], a.shape[1]).launch( + grid=(2 * self.clusters, 1, 1), block=(192, 1, 1), cluster=(2, 1, 1), + smem=SMEM_BYTES, stream=stream, + ) + + @cute.kernel + def kernel(self, ad: cutlass.GridConstant[cuda.TensorMap], + bd0: cutlass.GridConstant[cuda.TensorMap], + bd1: cutlass.GridConstant[cuda.TensorMap], + cd: cutlass.GridConstant[cuda.TensorMap], bias: cute.Tensor, + m: cutlass.Constexpr, n: cutlass.Constexpr, k: cutlass.Constexpr): + """Five-stage TMA/MMA ring with a separate, two-stage TMA store epilogue. + + Each role executes convergently within its assigned warps. Only CTA 0's + elected MMA lane issues MMAs. Every epilogue warp participates in TMEM + loads; all lanes of warp 0 participate in allocation and deallocation. + """ + tid, _, _ = cute.arch.thread_idx() + warp = tid // 32 + rank = prim.cluster_ctarank() + base = cute.arch.get_dyn_smem(cutlass.Uint8, alignment=1024) + full = _raw(cute.recast_ptr(base + _CONTROL, dtype=cutlass.Int64)) + empty = full + _STAGES + acc_empty, acc_full, dealloc = full + 10, full + 11, full + 12 + token = cutlass.Pointer((full + 13).to_llvm_ptr(), dtype=cutlass.Int32, space=cutlass.AddressSpace.smem) + if tid == 0: + for slot in cutlass.range_constexpr(_STAGES): + prim.mbarrier_init(full + slot, 1) + prim.mbarrier_init(empty + slot, 1) + prim.mbarrier_init(acc_empty, 2) + prim.mbarrier_init(acc_full, 1) + prim.mbarrier_init(dealloc, 32) + prim.fence_mbarrier_init() + prim.barrier_cluster_arrive_relaxed_aligned() + prim.barrier_cluster_wait_aligned() + if warp == 4: + if prim.elect_sync(): + prim.prefetch_tensormap(ad.get_ptr()) + prim.prefetch_tensormap(bd0.get_ptr()) + prim.prefetch_tensormap(bd1.get_ptr()) + prim.prefetch_tensormap(cd.get_ptr()) + + first = cute.arch.block_idx()[0] // 2 + stride = cute.arch.grid_dim()[0] // 2 + mtiles = m // 256 + tile_count = mtiles * ((n + 351) // 352) + ktiles = (k + 63) // 64 + + if warp == 5: + # Producers in both CTAs deliver distinct halves. CTA-group-2 + # redirects all six completion byte counts to CTA 0's full barrier. + slot, phase, tile = cutlass.Int32(0), cutlass.Int32(1), first + while tile < tile_count: + row = (tile % mtiles) * 256 + rank * 128 + col = (tile // mtiles) * 352 + kt = cutlass.Int32(0) + while kt < ktiles: + _wait(empty + slot, phase) + if prim.elect_sync(): + if rank == 0: + transaction_bytes = 2 * (ad.global_tx_bytes() + + bd0.global_tx_bytes() + bd1.global_tx_bytes()) + prim.mbarrier_arrive_expect_tx(full + slot, transaction_bytes) + prim.cp_async_bulk_tensor_shared_cluster_global( + _raw(base + slot * _A_BYTES), ad.get_ptr(), [kt * 64, row], + cutlass.Array(full + slot, (1,), dtype=cutlass.Int64), [], group=_GROUP, + ) + bstage = base + _B_OFFSET + slot * (_B0_BYTES + _B1_BYTES) + prim.cp_async_bulk_tensor_shared_cluster_global( + _raw(bstage), bd0.get_ptr(), [kt * 64, col + rank * 96], + cutlass.Array(full + slot, (1,), dtype=cutlass.Int64), [], group=_GROUP, + ) + prim.cp_async_bulk_tensor_shared_cluster_global( + _raw(bstage + _B0_BYTES), bd1.get_ptr(), + [kt * 64, col + 192 + rank * 80], + cutlass.Array(full + slot, (1,), dtype=cutlass.Int64), [], group=_GROUP, + ) + slot += 1 + if slot == _STAGES: + slot = 0 + phase ^= 1 + kt += 1 + tile += stride + # Wait for the last used slot, accounting for a short first K tile + # and a phase wrap. No producer can outlive its shared storage. + for _ in cutlass.range_constexpr(_STAGES - 1): + slot += 1 + if slot == _STAGES: + slot = 0 + phase ^= 1 + _wait(empty + slot, phase) + + elif warp == 4: + prim.barrier_cta_sync(2, thread_count=160) + tmem = token.load() + d0 = prim.make_tmem_ptr(tmem, cutlass.Float32) + d1 = prim.make_tmem_ptr(tmem + 192, cutlass.Float32) + id0 = prim.Tcgen05InstrDesc.build( + a_dtype=cutlass.Float16, b_dtype=cutlass.Float16, + c_dtype=cutlass.Float32, m_dim=256, n_dim=192, + ) + id1 = prim.Tcgen05InstrDesc.build( + a_dtype=cutlass.Float16, b_dtype=cutlass.Float16, + c_dtype=cutlass.Float32, m_dim=256, n_dim=160, + ) + slot, phase = cutlass.Int32(0), cutlass.Int32(0) + released, tile = cutlass.Int32(1), first + if rank == 0: + while tile < tile_count: + _wait(acc_empty, released, True) + kt = cutlass.Int32(0) + while kt < ktiles: + _wait(full + slot, phase) + prim.tcgen05_fence(prim.Tcgen05Fence.AFTER_THREAD_SYNC) + if prim.elect_sync(): + da = _operand(base + slot * _A_BYTES) + bstage = base + _B_OFFSET + slot * (_B0_BYTES + _B1_BYTES) + db0, db1 = _operand(bstage), _operand(bstage + _B0_BYTES) + for part in cutlass.range_constexpr(4): + use_d = kt > 0 + if cutlass.const_expr(part > 0): + use_d = cutlass.Boolean(True) + a16 = da.advance_start_address(part * 32) + prim.tcgen05_mma( + prim.Tcgen05MMAKind.F16, _GROUP, d0, a16, + db0.advance_start_address(part * 32), id0, use_d, + collector_op=prim.Tcgen05MMACollectorOp.FILL, + ) + prim.tcgen05_mma( + prim.Tcgen05MMAKind.F16, _GROUP, d1, a16, + db1.advance_start_address(part * 32), id1, use_d, + collector_op=prim.Tcgen05MMACollectorOp.LASTUSE, + ) + prim.tcgen05_commit(empty + slot, multicast_mask=3, group=_GROUP) + slot += 1 + if slot == _STAGES: + slot = 0 + phase ^= 1 + kt += 1 + if prim.elect_sync(): + prim.tcgen05_commit(acc_full, multicast_mask=3, group=_GROUP) + released ^= 1 + tile += stride + _wait(acc_empty, released, True) + else: + if warp == 0: + prim.tcgen05_alloc(token, 512, group=_GROUP) + prim.tcgen05_relinquish_alloc_permit(group=_GROUP) + prim.barrier_cta_sync(2, thread_count=160) + tmem = token.load() + # CuTe derives the physical shared addresses from the same B64 + # layout represented by the output tensor map, without XOR formulas. + shared_layout = cute.make_composed_layout( + cute.make_swizzle(2, 3, 3), 0, + cute.make_layout((128, 32), stride=(32, 1)), + ) + phase, output_slot, tile = cutlass.Int32(0), cutlass.Int32(0), first + while tile < tile_count: + row = (tile % mtiles) * 256 + rank * 128 + col = (tile // mtiles) * 352 + row_bias = cutlass.Float32(0) + if cutlass.const_expr(self.has_bias): + row_bias = bias[row + tid].to(cutlass.Float32) + _wait(acc_full, phase) + prim.tcgen05_fence(prim.Tcgen05Fence.AFTER_THREAD_SYNC) + for piece in range(11): + output_slot ^= 1 + dst = base + _C_OFFSET + output_slot * _C_BYTES + # A 32x32b.x32 load gives each lane 32 columns of one row. + address = tmem + (warp * 32 << 16) + piece * 32 + source = prim.make_tmem_ptr(address, cutlass.Float32) + values = prim.tcgen05_ld(prim.Tcgen05LdStShape.SHAPE_32X32B, source, num=32) + prim.tcgen05_wait(prim.Tcgen05Wait.LOAD) + narrowed = cute.make_rmem_tensor(32, cutlass.Float16) + for j in cutlass.range_constexpr(32): + narrowed[j] = (values[j] + row_bias).to(cutlass.Float16) + shared = cute.make_tensor(cute.recast_ptr(dst, dtype=cutlass.Float16), shared_layout) + row_view = shared[tid, None] + row_view = cute.make_tensor(row_view.iterator.align(16), row_view.layout) + cute.copy(cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float16, + num_bits_per_copy=128), narrowed, row_view) + prim.fence_proxy(prim.Proxy.ASYNC_SHARED, space=prim.SharedSpace.shared_cta) + prim.barrier_cta_sync(1, thread_count=128) + if tid == 0: + prim.cp_async_bulk_tensor_global_shared_cta( + cd.get_ptr(), _raw(dst), [col + piece * 32, row] + ) + prim.cp_async_bulk_commit_group() + prim.cp_async_bulk_wait_group(1, read=True) + prim.barrier_cta_sync(1, thread_count=128) + prim.tcgen05_fence(prim.Tcgen05Fence.BEFORE_THREAD_SYNC) + prim.barrier_cta_sync(1, thread_count=128) + if tid == 0: + # One arrival per CTA, after all four epilogue warps drain. + peer_empty = prim.mapa(acc_empty, 0) + prim.mbarrier_arrive(peer_empty, scope=prim.MemScope.CLUSTER) + phase ^= 1 + tile += stride + if warp == 0: + if tid == 0: + prim.cp_async_bulk_wait_group(0, read=True) + prim.barrier_cta_sync(3, thread_count=32) + peer_dealloc = prim.mapa(dealloc, rank ^ 1) + prim.mbarrier_arrive(peer_dealloc, scope=prim.MemScope.CLUSTER) + _wait(dealloc, cutlass.Int32(0), True) + pointer = prim.make_tmem_ptr(tmem, cutlass.Float32) + prim.tcgen05_dealloc(pointer, 512, group=_GROUP) + + +def _dimensions(m, n, k): + """Check positive signed-32-bit dimensions and an overflow-safe tile schedule.""" + dims = tuple(operator.index(x) for x in (m, n, k)) + if any(isinstance(x, bool) for x in (m, n, k)) or any(x <= 0 or x >= 2**31 for x in dims): + raise ValueError('M, N, and K must be positive signed-32-bit integers') + m, n, k = dims + if m % 256 or n % 8 or k % 8: + raise ValueError('M must be a multiple of 256; N and K must be multiples of 8') + if ((n + 351) // 352 - 1) * 352 + 320 >= 2**31: + raise ValueError('padded N coordinates must fit signed 32 bits') + if (m // 256) * ((n + 351) // 352) >= 2**31: + raise ValueError('persistent tile count must fit signed 32 bits') + return dims + + +def _version(): + """Reject unvalidated API versions rather than silently changing code generation.""" + if importlib.metadata.version('nvidia-cutlass-dsl') != '4.7.0': + raise RuntimeError('this reference requires nvidia-cutlass-dsl==4.7.0') + + +@functools.lru_cache(maxsize=128) +def compile_kernel(m, n, k, has_bias=False, clusters=1): + """Compile fake tensor contracts for SM100 without launching or needing SM100 hardware.""" + _version() + m, n, k = _dimensions(m, n, k) + clusters = operator.index(clusters) + if clusters <= 0 or clusters > (2**31 - 1) // 2 or clusters > (m // 256) * ((n + 351) // 352): + raise ValueError('clusters must be positive and no greater than the output tile count') + if (m // 256) * ((n + 351) // 352) + clusters - 1 >= 2**31: + raise ValueError('persistent tile advancement must fit signed 32 bits') + args = ( + make_fake_tensor(cutlass.Float16, (m, k), (k, 1), assumed_align=16), + make_fake_tensor(cutlass.Float16, (n, k), (k, 1), assumed_align=16), + make_fake_tensor(cutlass.Float16, (m, n), (n, 1), assumed_align=16), + make_fake_tensor(cutlass.Float16, (m,), (1,), assumed_align=16), + make_fake_stream(), + ) + return cute.compile(_Gemm(bool(has_bias), clusters), *args, options='--gpu-arch sm_100a') + + +def _sm100(device): + """Check the actual launch target; offline compilation intentionally skips this check.""" + import torch + if torch.device(device).type != 'cuda' or torch.cuda.get_device_capability(device) != (10, 0): + raise RuntimeError('this reference requires an SM100 GPU') + + +def cluster_capacity(): + """Return the public hardware-query capacity for two-CTA clusters on the current SM100 GPU.""" + import torch + device = torch.cuda.current_device() + _sm100(device) + return int(utils.HardwareInfo(device).get_max_active_clusters(2)) + + +def _arguments(tensors, has_bias): + """Validate Torch storage, exact strides, alignment, and common CUDA device.""" + import torch + if len(tensors) != (4 if has_bias else 3): + raise TypeError('expected A, B, bias, C' if has_bias else 'expected A, B, C') + a, b, c = tensors[0], tensors[1], tensors[-1] + if any(not isinstance(t, torch.Tensor) for t in tensors): + raise TypeError('arguments must be torch.Tensor objects') + if a.ndim != 3 or b.ndim != 3 or c.ndim != 3 or a.shape[0] != 1: + raise ValueError('expected A[1,M,K], B[1,K,N], and C[1,M,N]') + m, n, k = _dimensions(a.shape[1], b.shape[2], a.shape[2]) + if tuple(b.shape) != (1, k, n) or tuple(c.shape) != (1, m, n): + raise ValueError('matrix shapes do not match') + expected = [(m * k, k, 1), (n * k, 1, k), (m * n, n, 1)] + if any(tuple(t.stride()) != s for t, s in zip((a, b, c), expected)): + raise ValueError('A/C must be contiguous; B must transpose contiguous [1,N,K]') + if has_bias and (tuple(tensors[2].shape) != (m,) or tuple(tensors[2].stride()) != (1,)): + raise ValueError('bias must be contiguous with shape [M]') + if any(t.dtype != torch.float16 for t in tensors): + raise ValueError('all tensors must be FP16 tensors') + if any(t.device != a.device for t in tensors) or any(t.data_ptr() % 16 for t in tensors): + raise ValueError('tensors must share a device and have 16-byte-aligned addresses') + # Output must not overlap any input. Check the exact dense storage spans. + clo, chi = c.data_ptr(), c.data_ptr() + c.numel() * 2 + for t in tensors[:-1]: + if max(clo, t.data_ptr()) < min(chi, t.data_ptr() + t.numel() * 2): + raise ValueError('output must not overlap an input') + _sm100(a.device) + return m, n, k + + +def prepare_callable(has_bias=False, *, inputs=None, output=None, clusters=None): + """Prepare ``call(A, B, [bias,] C, stream=...)`` with owned, cached DLPack views. + + Supplying inputs/output performs validation, fake-tensor JIT compilation, + and view conversion immediately, with zero GPU kernel launches. Repeated + calls with the same tensor objects reuse those views during graph capture. + ``stream`` accepts a Torch CUDA stream or an explicit driver CUstream. + A raw CUstream must belong to the tensors' device; Torch streams are checked. + """ + _version() + cache = {} + + def fingerprint(tensors): + """Snapshot CPU storage metadata so mutable Tensor objects cannot reuse stale views.""" + return tuple((t.data_ptr(), tuple(t.shape), tuple(t.stride()), t.dtype, t.device) + for t in tensors) + + def bind(tensors): + """Own tensors and their 2-D DSL views for the lifetime of this callable.""" + import torch + key = tuple(id(t) for t in tensors) + if key in cache: + if fingerprint(tensors) != cache[key][3]: + raise ValueError('a prepared tensor changed its storage, shape, strides, dtype, or device') + else: + m, n, k = _arguments(tensors, has_bias) + with torch.cuda.device(tensors[0].device): + capacity = int(utils.HardwareInfo(tensors[0].device.index).get_max_active_clusters(2)) + requested = capacity if clusters is None else operator.index(clusters) + if requested <= 0: + raise ValueError('clusters must be positive') + count = min(requested, capacity, (m // 256) * ((n + 351) // 352)) + with torch.cuda.device(tensors[0].device): + compiled = compile_kernel(m, n, k, has_bias, count) + a, b, c = tensors[0], tensors[1], tensors[-1] + bias = tensors[2] if has_bias else a.reshape(-1)[:m] + views = tuple(from_dlpack(t, assumed_align=16) for t in (a[0], b[0].T, c[0], bias)) + cache[key] = (tensors, views, compiled, fingerprint(tensors)) + return cache[key] + + def call(*tensors, stream=None): + """Enqueue one GEMM on the explicitly requested stream, or the device's current stream.""" + import torch + owners, views, compiled, _ = bind(tensors) + if stream is None: + stream = torch.cuda.current_stream(owners[0].device) + if isinstance(stream, torch.cuda.Stream): + if stream.device != owners[0].device: + raise ValueError('stream and tensors must share a CUDA device') + stream = driver.CUstream(stream.cuda_stream) + if not isinstance(stream, driver.CUstream): + raise TypeError('stream must be torch.cuda.Stream or cuda.bindings.driver.CUstream') + with torch.cuda.device(owners[0].device): + compiled(*views, stream) + return tensors[-1] + + if (inputs is None) != (output is None): + raise ValueError('inputs and output must be supplied together') + if inputs is not None: + bind((*inputs, output)) + return call diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/oracle.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/oracle.py new file mode 100644 index 0000000000..7afa266f4a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/oracle.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Independent FP16 GEMM correctness oracle using public PyTorch operations. + +The reference accumulates in FP32, adds optional row bias in FP32, then rounds +once to FP16. Input integers keep ordinary benchmark sums exactly representable +in FP32. This module does not implement or time the CuTe device kernel. +""" + +from datetime import datetime, timezone +from numbers import Integral +import subprocess + +import torch + + +def prepare_tensors(family, mnk=(256, 352, 64), has_bias=False, seed=0, device="cuda"): + """Return owned inputs and one uninitialized output, on ``device``. + + A is contiguous [1,M,K]; B is a [1,K,N] transpose view of contiguous + [1,N,K] storage. Optional bias is contiguous [M]. All values are FP16 + integers in [-2,2). A private CPU generator preserves caller RNG state + and gives the same seeded inputs on CPU and CUDA. Output is [1,M,N]. + Only the ``fp16`` family and positive integer dimensions are supported. + """ + if family != "fp16": + raise ValueError("only the fp16 family is supported") + if len(mnk) != 3 or any(not isinstance(x, Integral) or isinstance(x, bool) or x <= 0 for x in mnk): + raise ValueError("mnk must contain three positive integers") + m, n, k = mnk + generator = torch.Generator(device="cpu").manual_seed(seed) + + def integers(shape): + """Allocate seeded FP16 integer data without mutating global RNG state.""" + return torch.randint(-2, 2, shape, generator=generator).to(device=device, dtype=torch.float16) + + inputs = [integers((1, m, k)), integers((1, n, k)).transpose(1, 2)] + if has_bias: + inputs.append(integers((m,))) + return inputs, [torch.empty((1, m, n), device=device, dtype=torch.float16)] + + +def verify(family, *tensors): + """Assert ``A @ B + bias[:,None]`` matches the final tensor. + + Accept A, B, C or A, B, bias, C in FP16 on one CPU/CUDA device; A and B + have a single batch dimension. Wrong shapes, dtypes, nonfinite values, + and differences outside atol=0.1/rtol=1e-5 raise AssertionError. TF32 and + autocast are disabled for FP32 reference accumulation. The caller's + float32 matmul precision is restored, including when checking fails. + """ + if family != "fp16": + raise ValueError("only the fp16 family is supported") + if len(tensors) not in (3, 4): + raise AssertionError("expected A, B, [bias,] C") + if any(not isinstance(t, torch.Tensor) for t in tensors): + raise AssertionError("all inputs and output must be tensors") + a, b = tensors[:2] + c = tensors[-1] + bias = tensors[2] if len(tensors) == 4 else None + if any(t.dtype != torch.float16 or t.device != a.device for t in tensors): + raise AssertionError("all inputs and output must be FP16 on one device") + if a.device.type not in ("cpu", "cuda"): + raise AssertionError("only CPU and CUDA tensors are supported") + if a.ndim != 3 or b.ndim != 3 or a.shape[0] != 1 or b.shape[0] != 1: + raise AssertionError("A and B must have shape [1,M,K] and [1,K,N]") + _, m, k = a.shape + n = b.shape[2] + if min(m, n, k) <= 0 or b.shape[1] != k or c.shape != (1, m, n): + raise AssertionError("incompatible A, B or C shape") + if bias is not None and bias.shape != (m,): + raise AssertionError("bias must have shape [M]") + if any(not bool(torch.isfinite(t).all()) for t in tensors): + raise AssertionError("inputs and output must contain only finite values") + precision = torch.get_float32_matmul_precision() + try: + torch.set_float32_matmul_precision("highest") + with torch.no_grad(), torch.autocast(device_type=a.device.type, enabled=False): + expected = torch.bmm(a.float(), b.float()) + if bias is not None: + expected = expected + bias.reshape(1, m, 1).float() + expected = expected.to(dtype=torch.float16) + if not bool(torch.isfinite(expected).all()): + raise AssertionError("FP16 reference result is not finite") + torch.testing.assert_close(c, expected, atol=0.1, rtol=1e-5) + finally: + torch.set_float32_matmul_precision(precision) + + +def clock_state(device="cuda:0"): + """Return a best-effort post-timing clock snapshot, never a timing claim. + + Query nvidia-smi by the selected device UUID, avoiding physical-index + assumptions under CUDA_VISIBLE_DEVICES. Missing UUIDs, nvidia-smi, or + unavailable telemetry produce an explicit unavailable status. + """ + snapshot = {"phase": "post-timing", "timestamp_utc": datetime.now(timezone.utc).isoformat()} + try: + uuid = getattr(torch.cuda.get_device_properties(device), "uuid", None) + if uuid is None: + raise RuntimeError("selected GPU UUID is unavailable") + uuid = str(uuid) + # PyTorch can expose the bare UUID; nvidia-smi expects its GPU prefix. + if not uuid.startswith(("GPU-", "MIG-")): + uuid = "GPU-" + uuid + result = subprocess.run( + ["nvidia-smi", "--id=" + uuid, + "--query-gpu=uuid,clocks.current.sm,clocks.current.memory,power.draw", + "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5, check=True, + ) + fields = [field.strip() for field in result.stdout.strip().split(",")] + if len(fields) != 4: + raise ValueError("unexpected nvidia-smi telemetry format") + snapshot.update(status="available", gpu_uuid=fields[0], + sm_clock_mhz=float(fields[1]), memory_clock_mhz=float(fields[2]), + power_w=float(fields[3])) + except (OSError, RuntimeError, AssertionError, ValueError, subprocess.SubprocessError) as exc: + snapshot.update(status="unavailable", reason=str(exc)) + return {"clock_snapshot": snapshot} diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/requirements.txt b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/requirements.txt new file mode 100644 index 0000000000..b2cb40f71c --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/requirements.txt @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# This comparison uses the shared environment pinned in cute-rs-preview/requirements.txt. +-r ../../../requirements.txt diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_compare.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_compare.py new file mode 100644 index 0000000000..b1f26718f4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_compare.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dependency-free checks for public comparison imports and CLI help.""" + +import os +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile +import unittest + + +class ComparisonCliTests(unittest.TestCase): + """Keep ignored local references outside the public comparison import path.""" + + def test_help_needs_no_runtime_dependencies(self): + result = subprocess.run( + [sys.executable, "-S", str(Path(__file__).resolve().parent.parent / "compare.py"), "--help"], + capture_output=True, text=True, check=False, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("--rust-bin", result.stdout) + + def test_ignored_reference_cannot_replace_public_oracle(self): + # The actual comparison script must select the public package even if + # old module names exist both beside it and on PYTHONPATH. + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + shutil.copyfile(Path(__file__).resolve().parent.parent / "compare.py", root / "compare.py") + public = root / "public_reference" + public.mkdir() + (public / "__init__.py").write_text("") + (public / "oracle.py").write_text("raise RuntimeError('public oracle selected')\n") + (public / "gemm.py").write_text("") + old = root / "reference" + old.mkdir() + for name in ("reference.py", "fp16_gemm_3_256x352.py"): + (old / name).write_text("raise RuntimeError('ignored private reference selected')\n") + for name in ("torch.py", "numpy.py"): + (root / name).write_text("") + environment = os.environ.copy() + environment["PYTHONPATH"] = str(old) + result = subprocess.run( + [sys.executable, "-S", str(root / "compare.py"), "--rust-bin", sys.executable], + env=environment, capture_output=True, text=True, check=False, + ) + self.assertNotEqual(result.returncode, 0) + self.assertIn("public oracle selected", result.stderr) + self.assertNotIn("ignored private reference selected", result.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_gemm.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_gemm.py new file mode 100644 index 0000000000..e541956dc7 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_gemm.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for the public kernel's validation and prepared-call ownership.""" + +import contextlib +import unittest +from unittest import mock + +try: + import torch + try: + from . import gemm + except ImportError: + import gemm +except ImportError: + gemm = None + + +@unittest.skipIf(gemm is None, 'install the optional pinned public-reference dependencies') +class GemmContractTests(unittest.TestCase): + """Exercise actual tensor metadata and cache behavior without an SM100 device.""" + + def test_dimension_and_launch_bounds_reject_before_compilation(self): + """Malformed dimensions and padded coordinate overflow never reach the compiler.""" + for dims in [(0, 8, 8), (128, 8, 8), (256, 7, 8), (256, 8, 7), + (256, 8, 2**31), (True, 8, 8), (256, 2**31 - 8, 8)]: + with self.subTest(dims=dims), self.assertRaises((TypeError, ValueError)): + gemm.compile_kernel(*dims) + for clusters in (0, -1, 2**30): + with self.subTest(clusters=clusters), self.assertRaises(ValueError): + gemm.compile_kernel(256, 352, 64, clusters=clusters) + + def test_exact_strides_and_output_alias_are_rejected(self): + """The host contract catches wrong B storage and destructive output overlap.""" + a = torch.empty((1, 256, 8), dtype=torch.float16) + b = torch.empty((1, 8, 8), dtype=torch.float16) + c = torch.empty_like(a) + with mock.patch.object(gemm, '_sm100'): + with self.assertRaisesRegex(ValueError, 'transpose'): + gemm._arguments((a, b, c), False) + with self.assertRaisesRegex(ValueError, 'overlap'): + gemm._arguments((a, b.transpose(1, 2), a), False) + + def test_preparation_launches_nothing_and_calls_reuse_owned_views(self): + """Eager JIT/view preparation queues no GPU work; matching calls reuse the views.""" + self._exercise_prepared_cache(mutate=False) + + def test_prepared_cache_rejects_storage_replacement(self): + """Tensor identity must not hide a replacement backing allocation.""" + self._exercise_prepared_cache(mutate=True) + + def _exercise_prepared_cache(self, mutate): + """Use real CPU tensors; replace only CUDA-specific target and execution services.""" + a = torch.empty((1, 256, 8), dtype=torch.float16) + b = torch.empty((1, 8, 8), dtype=torch.float16).transpose(1, 2) + c = torch.empty_like(a) + launches, conversions = [], [] + stream = gemm.driver.CUstream(0) + + def enqueue(*args): + launches.append(args) + + def view(tensor, **kwargs): + conversions.append(tensor) + return tensor + + with contextlib.ExitStack() as stack: + stack.enter_context(mock.patch.object(gemm, '_sm100')) + stack.enter_context(mock.patch.object(torch.cuda, 'device', side_effect=lambda _: contextlib.nullcontext())) + stack.enter_context(mock.patch.object(gemm.utils, 'HardwareInfo')) + gemm.utils.HardwareInfo.return_value.get_max_active_clusters.return_value = 1 + compiler = stack.enter_context(mock.patch.object(gemm, 'compile_kernel', return_value=enqueue)) + stack.enter_context(mock.patch.object(gemm, 'from_dlpack', side_effect=view)) + prepared = gemm.prepare_callable(inputs=(a, b), output=c) + self.assertEqual(launches, []) + self.assertEqual(len(conversions), 4) + self.assertIs(prepared(a, b, c, stream=stream), c) + self.assertIs(prepared(a, b, c, stream=stream), c) + self.assertEqual(len(launches), 2) + self.assertEqual(len(conversions), 4) + compiler.assert_called_once() + if mutate: + a.set_(torch.empty_like(a)) + with self.assertRaisesRegex(ValueError, 'changed its storage'): + prepared(a, b, c, stream=stream) + self.assertEqual(len(launches), 2) + self.assertEqual(len(conversions), 4) + compiler.assert_called_once() + + +if __name__ == '__main__': + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_setup.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_setup.py new file mode 100644 index 0000000000..302d56ac46 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_setup.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise setup recovery and offline reuse without downloading GPU packages.""" + +import importlib.util +from pathlib import Path +import subprocess +import tempfile +import unittest +from unittest import mock + + +class SetupTests(unittest.TestCase): + """Mock only package-manager processes; keep filesystem/cache behavior real.""" + + def setUp(self): + """Give each test an isolated cute-rs-preview tree, including spaces in its path.""" + helper = Path(__file__).resolve().parent.parent / "setup_reference.py" + self.assertTrue(helper.is_file(), "the example needs automatic Python setup") + spec = importlib.util.spec_from_file_location("setup_reference", helper) + self.setup = importlib.util.module_from_spec(spec) + spec.loader.exec_module(self.setup) + self.temporary = tempfile.TemporaryDirectory(prefix="gemm setup ") + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.example = self.root / "examples/gemm" + (self.example / "public_reference").mkdir(parents=True) + self.requirements = self.root / "requirements.txt" + self.requirements.write_text("nvidia-cutlass-dsl[cu13]==4.7.0\ntorch\nnumpy\n") + self.python = self.root / ".venv/bin/python" + self.python.parent.mkdir(parents=True) + self.python.touch() + (self.root / ".venv/pyvenv.cfg").write_text("include-system-site-packages = false\n") + self.installs = 0 + self.creations = 0 + self.healthy = False + self.fail_install = False + self.pip_missing = False + self.fail_seed = False + self.run_patch = mock.patch.object(self.setup.subprocess, "run", side_effect=self.run_process) + self.run_patch.start() + self.addCleanup(self.run_patch.stop) + self.uv_patch = mock.patch.object(self.setup.shutil, "which", return_value=None) + self.uv_patch.start() + self.addCleanup(self.uv_patch.stop) + + def run_process(self, command, **kwargs): + """Stand in for pip and dependency imports, which require large downloads.""" + command = [str(part) for part in command] + if command[1:3] == ["-m", "venv"] or command[1:2] == ["venv"]: + self.assertEqual(command[-1], str(self.root / ".venv")) + self.creations += 1 + self.python.parent.mkdir(parents=True, exist_ok=True) + self.python.touch() + (self.root / ".venv/pyvenv.cfg").write_text("include-system-site-packages = false\n") + return subprocess.CompletedProcess(command, 0) + if command[1:3] == ["-m", "ensurepip"]: + if self.fail_seed: + raise subprocess.CalledProcessError(1, command) + self.pip_missing = False + return subprocess.CompletedProcess(command, 0) + if command[1:3] == ["pip", "install"]: + self.assertEqual(command[3:5], ["--python", str(self.python)]) + command = [str(self.python), "-m", "pip", "install", *command[5:]] + if command[1:3] == ["-m", "pip"]: + if "install" in command: + self.installs += 1 + self.assertEqual(command[0], str(self.python)) + self.assertEqual(command[-2:], ["-r", str(self.requirements)]) + if self.fail_install: + raise subprocess.CalledProcessError(1, command) + self.healthy = True + return subprocess.CompletedProcess(command, 1 if self.pip_missing else 0) + if command[1:2] == ["-c"]: + return subprocess.CompletedProcess(command, 0 if self.healthy else 1, + stdout="", stderr="missing dependency") + self.fail(f"unexpected subprocess: {command}") + + def test_missing_environment_is_created_before_install(self): + """The default path must provision a fresh checkout, not just repair one.""" + self.python.unlink() + (self.root / ".venv/pyvenv.cfg").unlink() + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertTrue(self.python.is_file()) + self.assertEqual(self.creations, 1) + self.assertEqual(self.installs, 1) + + def test_uv_can_create_environment_without_ensurepip(self): + """The optional uv path must install into the shared isolated interpreter.""" + self.python.unlink() + (self.root / ".venv/pyvenv.cfg").unlink() + with mock.patch.object(self.setup.shutil, "which", return_value="/tools/uv"): + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertEqual(self.creations, 1) + self.assertEqual(self.installs, 1) + + def test_partial_environment_with_missing_pip_can_retry(self): + """An existing Python left by a failed venv bootstrap must not prevent recovery.""" + self.pip_missing = self.fail_seed = True + with self.assertRaisesRegex(RuntimeError, "venv/ensurepip"): + self.setup.ensure_environment(self.example) + self.fail_seed = False + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertFalse(self.pip_missing) + self.assertEqual(self.installs, 1) + + def test_first_setup_then_cached_run_needs_no_installer(self): + """A completed install must be reusable without network or package-manager work.""" + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertEqual(self.installs, 1) + self.fail_install = True + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertEqual(self.installs, 1) + + def test_changed_requirements_are_installed_again(self): + """An existing environment must not hide a newly required dependency.""" + self.setup.ensure_environment(self.example) + self.requirements.write_text(self.requirements.read_text() + "new-dependency\n") + self.setup.ensure_environment(self.example) + self.assertEqual(self.installs, 2) + + def test_failed_install_is_retried_without_deleting_environment(self): + """An interrupted install cannot leave a success marker or destroy local files.""" + preserved = self.root / ".venv/preserved" + preserved.write_text("keep") + self.fail_install = True + with self.assertRaises(RuntimeError): + self.setup.ensure_environment(self.example) + self.fail_install = False + self.assertEqual(self.setup.ensure_environment(self.example), self.python) + self.assertEqual(self.installs, 2) + self.assertEqual(preserved.read_text(), "keep") + + def test_broken_cached_dependencies_are_rechecked(self): + """Removing an installed package must invalidate otherwise unchanged setup.""" + self.setup.ensure_environment(self.example) + self.healthy = False + self.setup.ensure_environment(self.example) + self.assertEqual(self.installs, 2) + + def test_failed_cuda_build_check_never_marks_setup_complete(self): + """A successful pip exit alone is insufficient if imports or CUDA support fail.""" + real_run = self.run_process + + def cpu_torch(command, **kwargs): + """Simulate an installed CPU-only wheel at the external probe boundary.""" + result = real_run(command, **kwargs) + if str(command[1]) == "-c": + return subprocess.CompletedProcess(command, 1, stdout="", + stderr="CUDA-enabled PyTorch is required") + return result + + with mock.patch.object(self.setup.subprocess, "run", side_effect=cpu_torch): + with self.assertRaisesRegex(RuntimeError, "CUDA-enabled"): + self.setup.ensure_environment(self.example) + self.setup.ensure_environment(self.example) + self.assertEqual(self.installs, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_validate.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_validate.py new file mode 100644 index 0000000000..aec5080f55 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/test_validate.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dependency-free CLI tests for early rejection of unsupported hardware.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest + + +EXAMPLE = Path(__file__).resolve().parent.parent +"""Example directory containing the validation package.""" + + +class ValidationCliTests(unittest.TestCase): + """Check that CLI and hardware errors cannot accidentally start compilation.""" + + def invoke(self, *args: str, torch_stub: str | None = None): + """Run without site packages, optionally supplying the CUDA query boundary.""" + with tempfile.TemporaryDirectory() as directory: + environment = os.environ.copy() + environment["PYTHONPATH"] = os.pathsep.join((directory, str(EXAMPLE))) + if torch_stub is not None: + Path(directory, "torch.py").write_text(torch_stub) + return subprocess.run( + [sys.executable, "-S", "-m", "public_reference.validate", *args], + env=environment, + cwd=directory, + text=True, + capture_output=True, + check=False, + ) + + def test_help_works_without_cuda_packages(self): + """A top-level torch or CuTe import would break dependency-free help.""" + result = self.invoke("--help") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("--case", result.stdout) + + def test_missing_torch_is_reported_without_a_traceback(self): + """Missing runtime dependencies must produce a readable failed report.""" + result = self.invoke("--case", "full") + self.assertEqual(result.returncode, 1, result.stderr) + report = json.loads(result.stdout) + self.assertEqual(report["status"], "failed") + self.assertIn("torch", report["error"]) + self.assertNotIn("Traceback", result.stderr) + + def test_unavailable_cuda_does_not_query_or_compile_a_kernel(self): + """Querying an absent device or importing CuTe would replace this error.""" + result = self.invoke( + "--case", "full", + torch_stub=( + "class cuda:\n" + " @staticmethod\n" + " def is_available(): return False\n" + ), + ) + self.assertEqual(result.returncode, 1, result.stderr) + report = json.loads(result.stdout) + self.assertIn("CUDA is unavailable", report["error"]) + self.assertEqual(report["cases"], []) + + def test_sm120_is_rejected_before_loading_kernel_or_oracle(self): + """SM120 must fail at the architecture gate even without CuTe installed.""" + result = self.invoke( + "--case", "full", + torch_stub=( + "class cuda:\n" + " @staticmethod\n" + " def is_available(): return True\n" + " @staticmethod\n" + " def current_device(): return 0\n" + " @staticmethod\n" + " def get_device_capability(device): return (12, 0)\n" + " @staticmethod\n" + " def get_device_name(device): return 'test SM120'\n" + ), + ) + self.assertEqual(result.returncode, 1, result.stderr) + report = json.loads(result.stdout) + self.assertIn("SM100", report["error"]) + self.assertIn("12.0", report["error"]) + self.assertEqual(report["cases"], []) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/validate.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/validate.py new file mode 100644 index 0000000000..dfb0458a73 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/public_reference/validate.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Check the public FP16 reference on SM100, including stream and graph reuse. + +Run with ``python -m public_reference.validate`` from the example directory, +or wrap that command with compute-sanitizer. This runner reports +correctness only; use the example's compare.py for performance measurements. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import json +from pathlib import Path +import sys + + +@dataclass(frozen=True) +class Case: + """One GEMM extent, optional row bias, and resident-cluster override.""" + + name: str + """CLI selector for this coverage case.""" + mnk: tuple[int, int, int] + """Logical M, N, K dimensions, all measured in elements.""" + has_bias: bool = False + """Whether to add FP16 bias[M] before the final FP16 rounding.""" + clusters: int | None = None + """Resident cluster count; None lets the launcher select occupancy.""" + + +CASES = ( + Case("full", (256, 352, 64)), + Case("smalltails", (256, 8, 8), has_bias=True), + Case("persistent", (512, 712, 72), has_bias=True, clusters=1), + Case("5stages", (512, 704, 320), clusters=1), + Case("wrap", (512, 712, 392), has_bias=True, clusters=1), + Case("multicluster", (768, 1064, 704), has_bias=True, clusters=2), +) +"""Cover full tiles, N/K tails, pipeline depth, wraparound, and cluster reuse.""" + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse case selection without importing CUDA or kernel dependencies.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--case", action="append", choices=[case.name for case in CASES], + help="run only this case (repeatable; default: all six)", + ) + parser.add_argument("--seed", type=int, default=0, help="input seed (default: 0)") + parser.add_argument( + "--output", type=Path, help="also save the JSON report to this path", + ) + return parser.parse_args(argv) + + +def require_sm100(torch) -> dict: + """Reject missing CUDA and every capability except 10.0 before JIT imports.""" + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is unavailable; this validation requires SM100 (10.0)") + device = torch.cuda.current_device() + capability = torch.cuda.get_device_capability(device) + name = torch.cuda.get_device_name(device) + if capability != (10, 0): + raise RuntimeError( + f"SM100 (10.0) is required; device {device} is {name} " + f"with capability {capability[0]}.{capability[1]}" + ) + return {"index": device, "name": name, "capability": list(capability)} + + +def validate_case(torch, oracle, prepare_callable, case: Case, seed: int, record: dict): + """Compile once, then verify a stream launch and two poisoned graph replays. + + Inputs remain owned here until all launches complete. A copy queued on the + supplied stream is the observed output, so verification consumes that + stream's result after synchronization. Every invocation must overwrite NaNs. + """ + inputs, outputs = oracle.prepare_tensors( + "fp16", mnk=case.mnk, has_bias=case.has_bias, seed=seed, + ) + output = outputs[0] + launch = prepare_callable( + has_bias=case.has_bias, inputs=inputs, output=output, clusters=case.clusters, + ) + stream = torch.cuda.Stream(device=output.device) + stream.wait_stream(torch.cuda.current_stream(output.device)) + + # Keep the current stream unchanged at the call boundary: launch must use + # the stream argument. The consumer is queued on that supplied stream. + with torch.cuda.stream(stream): + output.fill_(float("nan")) + launch(*inputs, output, stream=stream) + with torch.cuda.stream(stream): + observed = output.clone() + stream.synchronize() + oracle.verify("fp16", *inputs, observed) + record["checks"].append("explicit-stream") + + # Eager compilation and the checked launch above populate cached tensor + # views before capture. Capture records only the launch, never compilation. + stream.wait_stream(torch.cuda.current_stream(output.device)) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + output.fill_(float("nan")) + with torch.cuda.graph(graph, stream=stream): + launch(*inputs, output, stream=stream) + for replay in (1, 2): + with torch.cuda.stream(stream): + output.fill_(float("nan")) + graph.replay() + observed = output.clone() + stream.synchronize() + oracle.verify("fp16", *inputs, observed) + record["checks"].append(f"graph-replay-{replay}") + + +def main(argv: list[str] | None = None) -> int: + """Emit one JSON report, returning nonzero on hardware or numerical failure.""" + args = parse_args(argv) + selected = set(args.case) if args.case else {case.name for case in CASES} + report = {"status": "failed", "dtype": "fp16", "seed": args.seed, "cases": []} + active = None + try: + import torch + + report["device"] = require_sm100(torch) + # The architecture check intentionally precedes CuTe and oracle imports. + from public_reference import oracle, gemm + + for case in CASES: + if case.name not in selected: + continue + active = { + "name": case.name, "mnk": list(case.mnk), "has_bias": case.has_bias, + "clusters": case.clusters, "seed": args.seed, "status": "running", + "checks": [], + } + report["cases"].append(active) + validate_case(torch, oracle, gemm.prepare_callable, case, args.seed, active) + active["status"] = "passed" + report["status"] = "passed" + except Exception as error: + report["error"] = f"{type(error).__name__}: {error}" + if active is not None: + active["status"] = "failed" + active["error"] = report["error"] + + encoded = json.dumps(report, indent=2) + "\n" + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(encoded) + sys.stdout.write(encoded) + return 0 if report["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/setup_reference.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/setup_reference.py new file mode 100644 index 0000000000..fbf8cf12be --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/setup_reference.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prepare the shared cute-rs-preview/.venv before comparison timing. + +Only that ignored virtualenv is managed, from cute-rs-preview/requirements.txt. +Successful setup prints its Python path to stdout; progress and package-manager +output go to stderr. The launcher bypasses this helper when the caller selects +an explicit Python environment or a prepared virtualenv already exists. +""" + +import fcntl +import hashlib +from pathlib import Path +import shutil +import subprocess +import sys + + +_PROBE = """ +import torch +import numpy +from public_reference import gemm +gemm._version() +if torch.version.cuda is None: + raise RuntimeError('CUDA-enabled PyTorch is required; the installed wheel is CPU-only') +""" +"""Check imports and the pinned DSL without initializing or requiring a GPU.""" + + +def probe(python, directory): + """Return import diagnostics from the managed interpreter, without installing.""" + return subprocess.run( + [str(python), "-c", _PROBE], cwd=directory, + capture_output=True, text=True, check=False, + ) + + +def install(command): + """Run a package-manager command with its progress outside benchmark stdout.""" + subprocess.run(command, check=True, stdout=sys.stderr) + + +def ensure_environment(directory): + """Create or repair the shared .venv, retry failed installs, and reuse healthy cached setup. + +An OS file lock serializes concurrent setup attempts and is released on process +exit. The requirements stamp is written only after installation and imports +succeed. Existing environments and their other files are never deleted. +""" + directory = Path(directory).resolve() + preview = directory.parent.parent + requirements = preview / "requirements.txt" + digest = hashlib.sha256(requirements.read_bytes()).hexdigest() + environment = preview / ".venv" + python = environment / "bin/python" + environment.mkdir(exist_ok=True) + with (environment / ".setup.lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + stamp = environment / ".requirements.sha256" + if python.is_file() and stamp.is_file() and stamp.read_text().strip() == digest: + if probe(python, directory).returncode == 0: + return python + # Never trust an earlier success after a failed repair or changed pin. + stamp.unlink(missing_ok=True) + print(f"Setting up the shared cute-rs comparison environment in {environment}", file=sys.stderr) + uv = shutil.which("uv") + try: + if not python.is_file() or not (environment / "pyvenv.cfg").is_file(): + if uv: + install([uv, "venv", "--allow-existing", "--python", sys.executable, + str(environment)]) + else: + install([sys.executable, "-m", "venv", str(environment)]) + if uv: + install([uv, "pip", "install", "--python", str(python), + "-r", str(requirements)]) + else: + pip = subprocess.run([str(python), "-m", "pip", "--version"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + if pip.returncode: + install([str(python), "-m", "ensurepip", "--upgrade"]) + install([str(python), "-m", "pip", "install", "--disable-pip-version-check", + "-r", str(requirements)]) + except (OSError, subprocess.CalledProcessError) as error: + raise RuntimeError( + "Could not set up the comparison environment. Install Python's venv/ensurepip " + "support (python3-venv on Debian/Ubuntu), or put uv on PATH; check the " + "package-manager error above and network access, then rerun. " + "You can also select a prepared environment with --python or " + "CUDA_OXIDE_CUTE_PYTHON." + ) from error + result = probe(python, directory) + if result.returncode: + raise RuntimeError( + f"Comparison dependencies installed, but validation failed:\n{result.stderr}" + f"Install a compatible CUDA-enabled PyTorch build in {environment}, " + "or select a prepared environment with --python." + ) + pending = stamp.with_suffix(".tmp") + pending.write_text(digest + "\n") + pending.replace(stamp) + return python + + +def main(): + """Print the configured interpreter, or an actionable setup failure.""" + if sys.argv[1:] in (["--help"], ["-h"]): + print(__doc__) + return 0 + if sys.argv[1:]: + print("Usage: python3 setup_reference.py", file=sys.stderr) + return 2 + try: + print(ensure_environment(Path(__file__).resolve().parent)) + except (OSError, RuntimeError) as error: + print(f"FP16 comparison setup failed: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/kernel.rs b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/kernel.rs new file mode 100644 index 0000000000..f3cf0dc555 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/kernel.rs @@ -0,0 +1,332 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! SM100 CuTe GEMM, emitted through the CUTLASS MLIR translation backend. +//! The two CTAs collectively compute 256x352, split into 256x192 and 256x160 +//! MMA instructions sharing the A operand collector. Five AB stages overlap +//! TMA with MMA; four epilogue warps drain TMEM through two shared C stages. + +use ctm_rs::memory::DynamicSharedArray; +use ctm_rs::{ + barrier, cluster, fence, + groups::CtaGroup2, + mbarrier::{self, Barrier, GeneralBarrier, ManagedBarrier, Ready, Uninit}, + memory::SharedAddress, + tcgen05::{ + self, + tmem::{TmemGuard, TmemReady, TmemUninit}, + }, + thread, tma, warp, +}; +use ctm_rs::{cluster_launch, cuda_module, kernel, launch_bounds, launch_contract}; +use cute_rs::sm100::PipelineTmaStore; +use cute_rs::{ + ComposedLayout, OperandA, OperandB, PipelineTmaUmma, PipelineUmmaAsync, RowMajor, + Sm100ClusterTmaCopy, Sm100SharedTile, Sm100TiledMma, Sm100TmemEpilogue, Swizzle, TensorMap, +}; + +// These element layouts also encode the host TMA descriptors. FP16 shifts +// the swizzle's low-bit position by one when lowering to byte addresses. +/// Per-CTA 128x64 FP16 A tile with the shared swizzle used by its tensor map. +pub type ASmem = ComposedLayout, 0, RowMajor<128, 64>>; +/// Per-CTA 96x64 FP16 B tile feeding the 192-column MMA partition. +pub type B0Smem = ComposedLayout, 0, RowMajor<96, 64>>; +/// Per-CTA 80x64 FP16 B tile feeding the 160-column MMA partition. +pub type B1Smem = ComposedLayout, 0, RowMajor<80, 64>>; +/// Per-CTA 128x32 FP16 output tile with the B64 shared swizzle. +pub type CSmem = ComposedLayout, 0, RowMajor<128, 32>>; + +/// Threads required in each CTA by the kernel launch contract. +pub const THREADS: u32 = 192; +/// Number of A/B shared-memory stages overlapped by TMA and MMA. +const STAGES: usize = 5; +/// Shared bytes occupied by one CTA's A stage. +const A_BYTES: usize = 128 * 64 * 2; +/// Shared bytes occupied by both B partitions in one CTA stage. +const B_BYTES: usize = 176 * 64 * 2; +/// Byte offset of the staged B ring after all staged A tiles. +const B_OFFSET: usize = A_BYTES * STAGES; +/// Byte offset of the two shared output stages after the A/B rings. +const C_OFFSET: usize = B_OFFSET + B_BYTES * STAGES; +/// Shared bytes occupied by one FP16 output stage. +const C_STAGE_BYTES: usize = 128 * 32 * 2; +/// Byte offset of load, accumulator, and deallocation barriers after output storage. +const BARRIER_OFFSET: usize = C_OFFSET + 2 * C_STAGE_BYTES; +/// Total dynamic shared-memory allocation in bytes required per CTA. +pub const DYNAMIC_SMEM_BYTES: usize = BARRIER_OFFSET + 128; +/// Compile-time consistency check for the schedule's dynamic shared-memory reservation. +const _: () = assert!(DYNAMIC_SMEM_BYTES == 211_072); + +/// Five-stage two-CTA TMA/MMA pipeline expecting 77,824 transaction bytes per stage. +type MainloopPipeline = PipelineTmaUmma<5, 77_824>; +/// Accumulator handoff between the MMA producer and four epilogue warps per CTA. +type AccumulatorPipeline = PipelineUmmaAsync<4>; +/// Paired 256x192 and 256x160 MMA instructions sharing the A collector. +type TiledMma = Sm100TiledMma; + +// Wait for both CTAs before freeing tensor memory. The CuTe pipelines below +// track when input tiles and accumulators can be reused. +/// Wait for both CTAs before freeing TMEM; bar must be the live initialized deallocation barrier. +#[inline(always)] +unsafe fn wait_deallocation(bar: *const Barrier) { + unsafe { + let bar = ManagedBarrier::::adopt_ready(bar.cast_mut()); + while !mbarrier::try_wait_parity_cluster(&bar, 0) {} + } +} + +/// Two-CTA SM100 GEMM entry point and its explicit launch/storage contract. +#[cuda_module] +pub mod kernels { + + use super::*; + + /// Compute output tiles with a (2,1,1) cluster and four typed tensor maps. + /// + /// Each map is passed by value; all threads borrow its read-only bytes. + /// The tensor allocations must remain live until the stream finishes. + #[launch_bounds(192)] + #[cluster_launch(2, 1, 1)] + #[launch_contract( + domain = 1, + block = (192, 1, 1), + dynamic_shared = 211072, + dynamic_shared_alignment = 1024, + min_compute_capability = (10, 0), + )] + #[kernel] + #[allow(clippy::too_many_arguments)] // Keep each tensor map explicit in the signature. + pub fn fp16_gemm_256x352( + #[grid_constant] a: &TensorMap, + #[grid_constant] b0: &TensorMap, + #[grid_constant] b1: &TensorMap, + #[grid_constant] c: &TensorMap, + bias: *const f16, + m: u64, + n: u64, + k: u64, + has_bias: u64, + ) { + // The CUTLASS launch ABI uses 64-bit scalar arguments. The host checks + // that each dimension fits before the kernel converts it to u32. + let m = m as u32; + let n = n as u32; + let k = k as u32; + let has_bias = has_bias != 0; + let tid = thread::threadIdx_x(); + let warp_id = tid / 32; + let rank = cluster::cluster_ctarank(); + let smem = DynamicSharedArray::::get_raw(); + // Shared storage: A[5], B[5], C[2], full[5], empty[5], acc_empty, + // acc_full, dealloc, tmem_token. Tile bases are 1024-byte aligned. + unsafe { + let full = smem.add(BARRIER_OFFSET).cast::(); + let empty = full.add(STAGES); + let acc_empty = empty.add(STAGES); + let acc_full = acc_empty.add(1); + let dealloc = acc_full.add(1); + let tmem_token = dealloc.add(1).cast::(); + let pipeline = MainloopPipeline::from_raw(full, empty); + let accumulator = AccumulatorPipeline::from_raw(acc_empty, acc_full); + let mma = TiledMma::new(); + let stores = PipelineTmaStore::<2>::new(); + if warp_id == 4 { + // CuTe maps retain their layout types; prefetch only needs + // their ABI-compatible opaque descriptor bytes. + tma::prefetch_tensormap((a as *const TensorMap).cast()); + tma::prefetch_tensormap((b0 as *const TensorMap).cast()); + tma::prefetch_tensormap((b1 as *const TensorMap).cast()); + tma::prefetch_tensormap((c as *const TensorMap).cast()); + } + if warp_id == 0 && warp::is_elected_sync(u32::MAX) { + ManagedBarrier::::from_static(dealloc).init_instruction(32); + accumulator.init(); + pipeline.init(); + } + fence::fence_mbarrier_init(); + cluster::barrier_cluster_arrive_relaxed(); + + // Each cluster visits output tiles in M-first order. + let mtiles = m / 256; + let total_tiles = mtiles * n.div_ceil(352); + let mut tile = thread::blockIdx_x() / 2; + let tile_stride = thread::gridDim_x() / 2; + let ktiles = k.div_ceil(64); + cluster::barrier_cluster_wait(); + + if warp_id == 5 { + let mut stage = 0usize; + let mut phase = 1u32; + while tile < total_tiles { + let row = (tile % mtiles) * 256 + rank * 128; + let col = (tile / mtiles) * 352; + let mut kt = 0; + while kt < ktiles { + pipeline.producer_acquire(stage as u32, phase); + if rank == 0 && warp::is_elected_sync(u32::MAX) { + pipeline.producer_expect_tx(stage as u32); + } + // Native CuTe cluster TMA performs its own warp election. + // Every producer lane must participate in each copy. + let barrier = pipeline.producer_get_barrier(stage as u32); + let mut a_tile = Sm100SharedTile::::from_raw( + smem.add(stage * A_BYTES).cast(), + ); + let mut b0_tile = Sm100SharedTile::::from_raw( + smem.add(B_OFFSET + stage * B_BYTES).cast(), + ); + let mut b1_tile = Sm100SharedTile::::from_raw( + smem.add(B_OFFSET + stage * B_BYTES + 96 * 64 * 2).cast(), + ); + Sm100ClusterTmaCopy::copy(a, &mut a_tile, row, kt * 64, barrier, rank); + Sm100ClusterTmaCopy::copy( + b0, + &mut b0_tile, + col + rank * 96, + kt * 64, + barrier, + rank, + ); + Sm100ClusterTmaCopy::copy( + b1, + &mut b1_tile, + col + 192 + rank * 80, + kt * 64, + barrier, + rank, + ); + stage += 1; + if stage == STAGES { + stage = 0; + phase ^= 1; + } + kt += 1; + } + tile += tile_stride; + } + // The next stage plus STAGES-1 identifies the last slot and + // its release phase, including wraparound and short K tails. + let mut tail = 0; + while tail < STAGES - 1 { + stage += 1; + if stage == STAGES { + stage = 0; + phase ^= 1; + } + tail += 1; + } + if warp::is_elected_sync(u32::MAX) { + pipeline.producer_acquire(stage as u32, phase); + } + } else if warp_id == 4 { + barrier::barrier_cta_sync(2, 160); + let tmem = *tmem_token; + let mut stage = 0usize; + let mut phase = 0u32; + let mut acc_phase = 1u32; + while tile < total_tiles { + if rank == 0 { + accumulator.producer_acquire(acc_phase); + let mut accumulate = false; + let mut kt = 0; + while kt < ktiles { + pipeline.consumer_wait(stage as u32, phase); + let a_tile = Sm100SharedTile::::from_raw( + smem.add(stage * A_BYTES).cast(), + ); + let b0_tile = Sm100SharedTile::::from_raw( + smem.add(B_OFFSET + stage * B_BYTES).cast(), + ); + let b1_tile = Sm100SharedTile::::from_raw( + smem.add(B_OFFSET + stage * B_BYTES + 96 * 64 * 2).cast(), + ); + if warp::is_elected_sync(u32::MAX) { + // Each of the four K16 steps reuses A across + // the two N partitions, then releases it. + mma.gemm(tmem, &a_tile, &b0_tile, &b1_tile, accumulate); + pipeline.consumer_release(stage as u32); + } + accumulate = true; + stage += 1; + if stage == STAGES { + stage = 0; + phase ^= 1; + } + kt += 1; + } + if warp::is_elected_sync(u32::MAX) { + accumulator.producer_commit(); + } + } + acc_phase ^= 1; + tile += tile_stride; + } + if rank == 0 && warp::is_elected_sync(u32::MAX) { + accumulator.producer_acquire(acc_phase); + } + } else { + if warp_id == 0 { + TmemGuard::::from_raw_slot(tmem_token) + .alloc_collective(); + tcgen05::relinquish_alloc_permit_cg2(); + } + barrier::barrier_cta_sync(2, 160); + let tmem = *tmem_token; + let mut phase = 0u32; + let mut epi_stage = 0usize; + while tile < total_tiles { + let row = (tile % mtiles) * 256 + rank * 128; + let col = (tile / mtiles) * 352; + let row_bias = if has_bias { + *bias.add((row + tid) as usize) as f32 + } else { + 0.0 + }; + accumulator.consumer_wait(phase); + let mut subtile = 0u32; + while subtile < 11 { + epi_stage ^= 1; + let dst = smem.add(C_OFFSET + epi_stage * C_STAGE_BYTES); + Sm100TmemEpilogue::::store( + Sm100TmemEpilogue::::slice_address(tmem, warp_id, subtile), + dst.cast(), + tid, + row_bias, + ); + fence::fence_proxy_async_shared_cta(); + barrier::barrier_cta_sync(1, 128); + if tid == 0 { + Sm100TmemEpilogue::::store_async( + c, + dst.cast(), + row, + col + subtile * 32, + ); + stores.producer_commit(); + stores.producer_acquire(); + } + barrier::barrier_cta_sync(1, 128); + subtile += 1; + } + accumulator.consumer_release(); + phase ^= 1; + tile += tile_stride; + } + if warp_id == 0 { + // Thread 0 owns every TMA store group. + if tid == 0 { + stores.producer_tail(); + } + barrier::barrier_cta_sync(3, 32); + let peer = cluster::mapa(SharedAddress::from_raw(dealloc), rank ^ 1); + cluster::arrive_remote(peer); + wait_deallocation(dealloc); + TmemGuard::::adopt_ready(tmem_token) + .dealloc_collective(); + } + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/main.rs new file mode 100644 index 0000000000..d9460539b3 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/main.rs @@ -0,0 +1,677 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Host input preparation, validation, and timing for the two-CTA 256x352 FP16 GEMM. + +#![feature(f16)] + +mod kernel; +mod python; + +use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D, sys}; +use cute_rs::copy::tma::create_tensor_map_tiled; +use serde_json::json; +use std::{ + error::Error, + path::{Path, PathBuf}, + sync::Arc, +}; + +/// Host-side result type carrying CUDA, file, and argument-validation errors. +type Result = std::result::Result>; + +/// Choose per-launch event timing or repeated launches captured in a CUDA graph. +#[derive(Clone, Copy, PartialEq, Eq)] +enum TimingMode { + /// Time individual kernel launches between CUDA events. + Direct, + /// Time replayed CUDA graphs and report latency per captured kernel launch. + Graph, +} + +impl TimingMode { + /// Stable timing-mode label included in output reports. + fn name(self) -> &'static str { + match self { + Self::Direct => "direct", + Self::Graph => "graph", + } + } + + /// Report the CUDA timing mechanism used for the selected mode. + fn tool(self) -> &'static str { + match self { + Self::Direct => "cuda-events-single-kernel", + Self::Graph => "cuda-graph-events-per-kernel", + } + } +} + +/// Validated matrix shape, timing controls, and input/output paths for the benchmark. +struct Args { + /// Run the Rust kernel alone instead of invoking the CUTLASS comparison harness. + rust_only: bool, + /// Optional Python executable used by the comparison harness. + python: Option, + /// Number of output rows; must be a positive multiple of 256. + m: u32, + /// Number of output columns; must be positive and divisible by eight. + n: u32, + /// Reduction extent in FP16 elements; must be positive and divisible by eight. + k: u32, + /// Add one FP16 bias value per output row when enabled. + has_bias: bool, + /// Untimed warmup iterations before collecting timing samples. + warmup: usize, + /// Number of timed iterations used for the latency report. + iters: usize, + /// Kernel submissions captured into each timed graph replay. + graph_launches: usize, + /// Whether timings measure direct launches or graph replay. + timing_mode: TimingMode, + /// Optional persistent two-CTA cluster count; otherwise chosen from occupancy. + clusters: Option, + /// Optional directory containing little-endian FP16 input files. + input_dir: Option, + /// Optional destination for the little-endian FP16 result matrix. + output: Option, + /// Optional destination for the machine-readable benchmark report. + json: Option, + /// Use externally verified input/output files instead of running the CPU reference. + skip_cpu_check: bool, +} + +impl Args { + /// Parse command-line options and validate shape, launch limits, timing, and oracle-mode requirements. + fn parse() -> Result { + let mut args = Self { + rust_only: false, + python: None, + m: 256, + n: 352, + k: 64, + has_bias: false, + warmup: 25, + iters: 100, + graph_launches: 10, + timing_mode: TimingMode::Graph, + clusters: None, + input_dir: None, + output: None, + json: None, + skip_cpu_check: false, + }; + let mut timing_mode_explicit = false; + let mut graph_launches_explicit = false; + let mut cli = std::env::args().skip(1); + while let Some(key) = cli.next() { + match key.as_str() { + "--help" | "-h" => { + println!( + "fp16_gemm_256x352_cute [--mnk M,N,K] [--has-bias] [--warmup 25] [--iters 100]\n [--clusters COUNT] [--json RESULT.json] [--python PATH]\nDefault: verify and compare CUTLASS (CuTe + CTM) and cute-rs with direct CUDA-event timing.\nPrints two performance lines; --json saves the combined report.\nPython selection: --python, CUDA_OXIDE_CUTE_PYTHON, example .venv, cute-rs-preview/.venv, otherwise set up cute-rs-preview/.venv automatically.\nRequires the local compare.py and public_reference/ assets.\n\n--rust-only runs cute-rs alone with CPU verification and JSON output:\n [--timing-mode direct|graph] [--graph-launches 10] [--input-dir DIR] [--output C.f16]\n [--skip-cpu-check] (requires --input-dir and --output; external oracle mode)\nRust-only timing defaults to graph; comparison supports direct timing only.\nDirect --warmup 0 --iters 1 launches each selected implementation exactly once.\nM must be a positive multiple of 256, N and K of 8. SM100 GPU required." + ); + std::process::exit(0); + } + "--has-bias" | "--has_bias" => args.has_bias = true, + "--skip-cpu-check" => args.skip_cpu_check = true, + "--rust-only" => args.rust_only = true, + _ => { + let value = cli + .next() + .ok_or_else(|| format!("missing value for {key}"))?; + match key.as_str() { + "--mnk" => { + let shape = value + .split(',') + .map(str::parse::) + .collect::, _>>()?; + if shape.len() != 3 { + return Err("--mnk requires M,N,K".into()); + } + (args.m, args.n, args.k) = (shape[0], shape[1], shape[2]); + } + "--warmup" => args.warmup = value.parse()?, + "--iters" => args.iters = value.parse()?, + "--graph-launches" => { + args.graph_launches = value.parse()?; + graph_launches_explicit = true; + } + "--timing-mode" => { + timing_mode_explicit = true; + args.timing_mode = match value.as_str() { + "direct" => TimingMode::Direct, + "graph" => TimingMode::Graph, + _ => return Err("--timing-mode requires direct or graph".into()), + }; + } + "--clusters" => args.clusters = Some(value.parse()?), + "--input-dir" => args.input_dir = Some(value.into()), + "--output" => args.output = Some(value.into()), + "--json" => args.json = Some(value.into()), + "--python" => args.python = Some(value.into()), + _ => return Err(format!("unknown argument {key}").into()), + } + } + } + } + if !args.rust_only { + if !timing_mode_explicit { + args.timing_mode = TimingMode::Direct; + } + if !graph_launches_explicit { + args.graph_launches = 1; + } + if args.timing_mode != TimingMode::Direct || args.graph_launches != 1 { + return Err("CUTLASS (CuTe + CTM) comparison requires direct timing and --graph-launches 1; use --rust-only for graph timing".into()); + } + if args.input_dir.is_some() || args.output.is_some() || args.skip_cpu_check { + return Err("--input-dir, --output, and --skip-cpu-check require --rust-only; comparison supplies identical inputs and verifies both outputs".into()); + } + } else if args.python.is_some() { + return Err("--python applies only to the CUTLASS (CuTe + CTM) comparison".into()); + } + if args.m == 0 + || args.n == 0 + || args.k == 0 + || !args.m.is_multiple_of(256) + || !args.n.is_multiple_of(8) + || !args.k.is_multiple_of(8) + { + return Err("require positive M divisible by 256, N and K by 8".into()); + } + if args.m > i32::MAX as u32 || args.n > i32::MAX as u32 || args.k > i32::MAX as u32 { + return Err("TMA coordinates must fit i32".into()); + } + // The last N tile issues all eleven stores even when some are + // out-of-bounds. Their starting coordinates must remain positive + // after the kernel casts to i32. + let ntiles = args.n.div_ceil(352); + let last_store_col = u64::from(ntiles - 1) * 352 + 320; + if last_store_col > i32::MAX as u64 { + return Err("padded N-tile TMA coordinates must fit i32".into()); + } + let tiles = (args.m / 256) + .checked_mul(ntiles) + .ok_or("tile count overflow")?; + // The occupancy probe uses the full problem grid. This bound also + // prevents tile += tile_stride from wrapping in the persistent loop. + if tiles > i32::MAX as u32 / 2 { + return Err("two-CTA problem grid must fit the CUDA grid-x limit".into()); + } + for count in [ + u64::from(args.m) * u64::from(args.k), + u64::from(args.n) * u64::from(args.k), + u64::from(args.m) * u64::from(args.n), + ] { + if count > isize::MAX as u64 / 2 { + return Err("FP16 allocation size exceeds the host addressable limit".into()); + } + } + if args.iters == 0 || args.clusters == Some(0) { + return Err("iters and clusters must be positive".into()); + } + if args.timing_mode == TimingMode::Graph && args.graph_launches == 0 { + return Err("graph-launches must be positive in graph timing mode".into()); + } + if args.skip_cpu_check && (args.input_dir.is_none() || args.output.is_none()) { + return Err( + "--skip-cpu-check requires --input-dir and --output for external verification" + .into(), + ); + } + Ok(args) + } +} + +/// Run the local Python comparison harness with matching matrix shape and direct timing settings. +fn compare(args: &Args) -> Result<()> { + let directory = Path::new(env!("CARGO_MANIFEST_DIR")); + let script = directory.join("compare.py"); + if !script.is_file() + || !directory.join("public_reference/__init__.py").is_file() + || !directory.join("public_reference/gemm.py").is_file() + || !directory.join("public_reference/oracle.py").is_file() + { + return Err(format!( + "CUTLASS (CuTe + CTM) comparison requires the local compare.py and public_reference/ assets in {}; use --rust-only for standalone execution", + directory.display() + ) + .into()); + } + let python = python::select( + directory, + args.python.as_deref(), + std::env::var_os("CUDA_OXIDE_CUTE_PYTHON"), + )?; + // Start Python before creating a CUDA context. It runs this executable + // with --rust-only, using the same input bytes for both implementations. + let mut command = std::process::Command::new(&python); + command + .arg(script) + .arg("--rust-bin") + .arg(std::env::current_exe()?) + .args(["--mnk", &format!("{},{},{}", args.m, args.n, args.k)]) + .args(["--warmup", &args.warmup.to_string()]) + .args(["--iters", &args.iters.to_string()]) + .args(["--timing-mode", "direct", "--graph-launches", "1"]); + if args.has_bias { + command.arg("--has-bias"); + } + if let Some(clusters) = args.clusters { + command.args(["--clusters", &clusters.to_string()]); + } + if let Some(path) = &args.json { + command.arg("--json").arg(path); + } + let status = command.status().map_err(|error| { + format!( + "could not start CUTLASS (CuTe + CTM) comparison with {}: {error}; select its Python environment with --python or CUDA_OXIDE_CUTE_PYTHON", + python.display() + ) + })?; + if !status.success() { + return Err(format!("CUTLASS (CuTe + CTM) comparison failed ({status})").into()); + } + Ok(()) +} + +/// Convert a CUDA driver result into a host error annotated with the failed operation. +fn cuda_status(status: sys::CUresult, operation: &str) -> Result<()> { + if status == sys::cudaError_enum_CUDA_SUCCESS { + Ok(()) + } else { + Err(format!("{operation}: {}", cuda_core::DriverError(status)).into()) + } +} + +// The executable owns graph resources only; main keeps every captured +// module/allocation alive until after this object is dropped. +/// Own an executable CUDA graph; captured modules and allocations remain owned by main. +struct KernelGraph { + /// Instantiated CUDA graph handle destroyed when this wrapper is dropped. + executable: sys::CUgraphExec, + /// Stream and context on which the graph was captured and is replayed. + stream: Arc, +} + +impl KernelGraph { + /// Capture repeated calls to run, instantiate the graph, and clean up partial capture on failure. + fn capture( + stream: &Arc, + repetitions: usize, + run: impl Fn() -> Result<()>, + ) -> Result { + stream.context().bind_to_thread()?; + // All input copies and warmups have completed before capture begins. + cuda_status( + unsafe { + sys::cuStreamBeginCapture_v2( + stream.cu_stream(), + sys::CUstreamCaptureMode_enum_CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, + ) + }, + "begin graph capture", + )?; + let captured = (|| { + for _ in 0..repetitions { + run()?; + } + Ok(()) + })(); + let mut graph = std::ptr::null_mut(); + // Always terminate capture, including when a kernel submission fails. + let ended = unsafe { sys::cuStreamEndCapture(stream.cu_stream(), &mut graph) }; + if let Err(error) = captured { + if !graph.is_null() { + unsafe { + sys::cuGraphDestroy(graph); + } + } + return Err(error); + } + if let Err(error) = cuda_status(ended, "end graph capture") { + if !graph.is_null() { + unsafe { + sys::cuGraphDestroy(graph); + } + } + return Err(error); + } + let mut executable = std::ptr::null_mut(); + let instantiated = unsafe { sys::cuGraphInstantiateWithFlags(&mut executable, graph, 0) }; + // Instantiation owns its own graph state; the source graph is no + // longer needed regardless of whether instantiation succeeded. + unsafe { + sys::cuGraphDestroy(graph); + } + cuda_status(instantiated, "instantiate graph")?; + Ok(Self { + executable, + stream: Arc::clone(stream), + }) + } + + /// Bind the stream's context and enqueue the captured graph once. + fn replay(&self) -> Result<()> { + self.stream.context().bind_to_thread()?; + cuda_status( + unsafe { sys::cuGraphLaunch(self.executable, self.stream.cu_stream()) }, + "replay graph", + ) + } +} + +impl Drop for KernelGraph { + fn drop(&mut self) { + if self.stream.context().bind_to_thread().is_ok() { + unsafe { + sys::cuGraphExecDestroy(self.executable); + } + } + } +} + +/// Read exactly count little-endian FP16 bit patterns, rejecting a mismatched file size. +fn load_f16(path: &Path, count: usize) -> Result> { + let bytes = std::fs::read(path)?; + if bytes.len() != count.checked_mul(2).ok_or("input size overflow")? { + return Err(format!( + "{}: expected {count} FP16 values, got {} bytes", + path.display(), + bytes.len() + ) + .into()); + } + Ok(bytes + .as_chunks::<2>() + .0 + .iter() + .map(|v| u16::from_le_bytes(*v)) + .collect()) +} + +/// Write FP16 bit patterns as little-endian bytes and flush the output file. +fn store_f16(path: &Path, data: &[u16]) -> Result<()> { + use std::io::Write; + let mut file = std::io::BufWriter::new(std::fs::File::create(path)?); + for value in data { + file.write_all(&value.to_le_bytes())?; + } + file.flush()?; + Ok(()) +} + +// The standalone run uses reproducible integers in [-2, 2), exact in FP16. +// compare.py supplies the same seeded PyTorch inputs to both implementations. +/// Generate reproducible integers in [-2, 2) and return their exact FP16 bit patterns. +fn integer_input(count: usize, mut state: u32) -> Vec { + (0..count) + .map(|_| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + (((state >> 16) % 4) as i32 - 2) as f16 + }) + .map(f16::to_bits) + .collect() +} + +/// Compare GPU output with FP32 CPU accumulation rounded to FP16, including optional row bias. +fn verify(args: &Args, a: &[u16], b: &[u16], bias: &[u16], got: &[u16]) -> Result<()> { + let (m, n, k) = (args.m as usize, args.n as usize, args.k as usize); + let mut failures = 0; + for row in 0..m { + for col in 0..n { + let mut sum = 0.0f32; + for inner in 0..k { + sum += f16::from_bits(a[row * k + inner]) as f32 + * f16::from_bits(b[col * k + inner]) as f32; + } + if args.has_bias { + sum += f16::from_bits(bias[row]) as f32; + } + let actual = f16::from_bits(got[row * n + col]) as f32; + // The Python oracle rounds its FP32 accumulation to the output + // dtype before comparison. Apply the identical acceptance band. + let expected = (sum as f16) as f32; + if !actual.is_finite() || (actual - expected).abs() > 0.1 + 1.0e-5 * expected.abs() { + if failures < 5 { + eprintln!("C[{row},{col}]: expected {expected}, got {actual}"); + } + failures += 1; + } + } + } + if failures != 0 { + return Err(format!("CPU verification failed: {failures} / {} values", got.len()).into()); + } + eprintln!("CPU verification: PASS ({} values)", got.len()); + Ok(()) +} + +/// Validate inputs, launch the SM100 kernel, verify results, and report direct or graph timing. +fn main() -> Result<()> { + let args = Args::parse()?; + if !args.rust_only { + return compare(&args); + } + let single_launch = + args.timing_mode == TimingMode::Direct && args.warmup == 0 && args.iters == 1; + let ctx = CudaContext::new(0)?; + let capability = ctx.compute_capability()?; + if capability != (10, 0) { + return Err(format!( + "this example requires SM100; found sm_{}{}", + capability.0, capability.1 + ) + .into()); + } + // Stream capture cannot start on CUDA's legacy default stream. + let stream = ctx.new_stream()?; + let (m, n, k) = (args.m as usize, args.n as usize, args.k as usize); + let (a, b, bias) = if let Some(dir) = &args.input_dir { + ( + load_f16(&dir.join("a.f16"), m * k)?, + load_f16(&dir.join("b.f16"), n * k)?, + if args.has_bias { + load_f16(&dir.join("bias.f16"), m)? + } else { + vec![0; m] + }, + ) + } else { + ( + integer_input(m * k, 1), + integer_input(n * k, 7), + (0..m) + .map(|i| (((i % 17) as f32 - 8.0) / 8.0) as f16) + .map(f16::to_bits) + .collect(), + ) + }; + let a_dev = DeviceBuffer::from_host(&stream, &a)?; + let b_dev = DeviceBuffer::from_host(&stream, &b)?; + let bias_dev = DeviceBuffer::from_host(&stream, &bias)?; + // Poison the output so missing tiles, including N tails, fail verification. + let mut c_dev = DeviceBuffer::from_host(&stream, &vec![f16::NAN.to_bits(); m * n])?; + // These layout types set both the descriptor's tile shape and swizzle + // and the kernel's shared-memory layout. + let a_desc = create_tensor_map_tiled::( + a_dev.cu_deviceptr() as *mut core::ffi::c_void, + u64::from(args.m), + u64::from(args.k), + u64::from(args.k), + )?; + let b0_desc = create_tensor_map_tiled::( + b_dev.cu_deviceptr() as *mut core::ffi::c_void, + u64::from(args.n), + u64::from(args.k), + u64::from(args.k), + )?; + let b1_desc = create_tensor_map_tiled::( + b_dev.cu_deviceptr() as *mut core::ffi::c_void, + u64::from(args.n), + u64::from(args.k), + u64::from(args.k), + )?; + let c_desc = create_tensor_map_tiled::( + c_dev.cu_deviceptr() as *mut core::ffi::c_void, + u64::from(args.m), + u64::from(args.n), + u64::from(args.n), + )?; + // SAFETY: this executable exclusively owns the module and all its buffers. + let module = unsafe { kernel::kernels::load(&ctx) }?; + let tiles = (args.m / 256) + .checked_mul(args.n.div_ceil(352)) + .ok_or("tile count overflow")?; + let blocks = tiles.checked_mul(2).ok_or("grid size overflow")?; + let probe = module.prepare_fp16_gemm_256x352(LaunchConfig1D::new( + 2, + kernel::THREADS, + kernel::DYNAMIC_SMEM_BYTES as u32, + ))?; + let capacity = probe.function().max_active_clusters( + (blocks, 1, 1), + (kernel::THREADS, 1, 1), + kernel::DYNAMIC_SMEM_BYTES as u32, + (2, 1, 1), + )?; + let clusters = args.clusters.unwrap_or(capacity.min(tiles)); + if clusters == 0 || clusters > capacity || clusters > tiles { + return Err(format!( + "requested {clusters} clusters; resident capacity={capacity}, tiles={tiles}" + ) + .into()); + } + let registers = probe.function().num_registers()?; + drop(probe); + let launch = module.prepare_fp16_gemm_256x352(LaunchConfig1D::new( + clusters * 2, + kernel::THREADS, + kernel::DYNAMIC_SMEM_BYTES as u32, + ))?; + let run = || -> Result<()> { + // SAFETY: each descriptor is copied into read-only kernel parameters. + // We keep all tensor buffers alive until the stream finishes. The + // prepared launch checks block, cluster, and shared-memory sizes. + // Output tiles do not overlap, and the tensor shapes were checked above. + unsafe { + module.fp16_gemm_256x352( + &stream, + &launch, + a_desc, + b0_desc, + b1_desc, + c_desc, + bias_dev.cu_deviceptr() as *const f16, + u64::from(args.m), + u64::from(args.n), + u64::from(args.k), + u64::from(args.has_bias), + )?; + } + Ok(()) + }; + // A single direct sample verifies its timed output without a preliminary launch. + let got = if single_launch { + None + } else { + run()?; + let got = c_dev.to_host_vec(&stream)?; + if !args.skip_cpu_check { + verify(&args, &a, &b, &bias, &got)?; + } + Some(got) + }; + for _ in 0..args.warmup { + run()?; + } + stream.synchronize()?; + let graph = if args.timing_mode == TimingMode::Graph { + let graph = KernelGraph::capture(&stream, args.graph_launches, run)?; + // Keep graph setup and the first replay outside timing. Poison the + // output so an empty or wrong-stream graph cannot preserve a pass. + c_dev.copy_from_host(&stream, &vec![f16::NAN.to_bits(); m * n])?; + graph.replay()?; + if Some(c_dev.to_host_vec(&stream)?) != got { + return Err("captured graph output differs from the checked direct launch".into()); + } + for _ in 0..args.warmup { + graph.replay()?; + } + stream.synchronize()?; + Some(graph) + } else { + None + }; + let launches_per_sample = if graph.is_some() { + args.graph_launches + } else { + 1 + }; + let events = (0..args.iters) + .map(|_| { + Ok(( + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event_with_flags(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + )) + }) + .collect::>>()?; + for (start, end) in &events { + start.record(&stream)?; + if let Some(graph) = &graph { + graph.replay()?; + } else { + run()?; + } + end.record(&stream)?; + if args.timing_mode == TimingMode::Direct { + // Match blockscale_gemm_cute: complete each single-kernel sample + // before submitting the next one, reusing the same input buffers. + end.synchronize()?; + } + } + stream.synchronize()?; + let after_timing = c_dev.to_host_vec(&stream)?; + if let Some(got) = got { + if after_timing != got { + return Err("output changed across identical repeated launches".into()); + } + } else if !args.skip_cpu_check { + verify(&args, &a, &b, &bias, &after_timing)?; + } + if let Some(path) = &args.output { + store_f16(path, &after_timing)?; + } + let mut times = events + .iter() + .map(|(start, end)| { + Ok(f64::from(start.elapsed_ms(end)?) * 1000.0 / launches_per_sample as f64) + }) + .collect::>>()?; + times.sort_by(f64::total_cmp); + let mid = times.len() / 2; + let median = if times.len() % 2 == 0 { + (times[mid - 1] + times[mid]) / 2.0 + } else { + times[mid] + }; + let result = json!({"implementation": "cute-rs", "mnk": [m,n,k], "has_bias": args.has_bias, + "metric_value": median, "metric_unit": "us", "tflops": 2.0*m as f64*n as f64*k as f64/(median*1.0e6), + "verification": if args.skip_cpu_check { "external-oracle-required" } else { "cpu-pass" }, + "provenance": {"gpu": ctx.device_name()?, "compute_capability": [capability.0,capability.1], + "tool": args.timing_mode.tool(), "timing_mode": args.timing_mode.name(), + "warmup": args.warmup, "iters": args.iters, "single_launch": single_launch, + "graph_launches": launches_per_sample, + "clusters": clusters, "resident_cluster_capacity": capacity, "threads": kernel::THREADS, + "shared_memory_bytes": kernel::DYNAMIC_SMEM_BYTES, "registers_per_thread": registers}}); + let rendered = serde_json::to_string_pretty(&result)?; + if let Some(path) = args.json { + std::fs::write(path, &rendered)?; + } + println!("{rendered}"); + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/python.rs b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/python.rs new file mode 100644 index 0000000000..73e570882b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/src/python.rs @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Select caller-managed Python or a prepared virtualenv, or set up the shared one. + +use std::{ + ffi::OsString, + path::{Path, PathBuf}, + process::{Command, Stdio}, +}; + +/// Honor explicit interpreters and prepared virtualenvs; otherwise finish setup before any GPU timing. +pub(super) fn select( + directory: &Path, + explicit: Option<&Path>, + environment: Option, +) -> super::Result { + if let Some(python) = explicit { + return Ok(python.to_owned()); + } + if let Some(python) = environment { + return Ok(python.into()); + } + if let Some(python) = prepared(directory) { + return Ok(python); + } + let setup = directory.join("setup_reference.py"); + if !setup.is_file() { + return Err(format!("missing comparison setup helper: {}", setup.display()).into()); + } + let output = Command::new("python3") + .arg(setup) + .stderr(Stdio::inherit()) + .output() + .map_err(|error| format!("could not start comparison setup with python3: {error}"))?; + if !output.status.success() { + return Err(format!("comparison environment setup failed ({})", output.status).into()); + } + let python = PathBuf::from(String::from_utf8(output.stdout)?.trim()); + if !python.is_file() { + return Err(format!("setup returned an unavailable Python: {}", python.display()).into()); + } + Ok(python) +} + +/// The example's own `.venv`, then the shared `cute-rs-preview/.venv`. +fn prepared(directory: &Path) -> Option { + [ + directory.join(".venv/bin/python"), + directory.join("../../.venv/bin/python"), + ] + .into_iter() + .find(|path| path.is_file()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prepared_virtualenvs_are_used_before_setup() { + let temp = std::env::temp_dir().join(format!("fp16-gemm-python-{}", std::process::id())); + let example = temp.join("examples/gemm"); + let shared = temp.join(".venv/bin/python"); + let local = example.join(".venv/bin/python"); + std::fs::create_dir_all(shared.parent().unwrap()).unwrap(); + std::fs::create_dir_all(local.parent().unwrap()).unwrap(); + std::fs::write(&shared, "").unwrap(); + // The example directory has no setup helper, so only a prepared virtualenv can succeed. + assert_eq!( + select(&example, None, None).unwrap(), + example.join("../../.venv/bin/python") + ); + std::fs::write(&local, "").unwrap(); + assert_eq!(select(&example, None, None).unwrap(), local); + std::fs::remove_dir_all(&temp).unwrap(); + } + + #[test] + fn explicit_python_overrides_environment_without_running_setup() { + let selected = select( + Path::new("/no-setup-assets"), + Some(Path::new("/custom python/bin/python")), + Some("/other/python".into()), + ) + .unwrap(); + assert_eq!(selected, Path::new("/custom python/bin/python")); + } + + #[test] + fn environment_python_bypasses_local_setup() { + let selected = select( + Path::new("/no-setup-assets"), + None, + Some("/configured/python".into()), + ) + .unwrap(); + assert_eq!(selected, Path::new("/configured/python")); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/test_reference.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/test_reference.py new file mode 100644 index 0000000000..268100e17b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/test_reference.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contract tests for the public FP16 correctness oracle.""" + +import shutil +import subprocess +import unittest + +import torch + +from public_reference import oracle +import compare + + +class ReferenceTests(unittest.TestCase): + """Exercise layouts and numerics independently of the CuTe kernel.""" + + def test_seed_preserves_rng_and_produces_shared_integer_inputs(self): + state = torch.random.get_rng_state().clone() + first, outputs = oracle.prepare_tensors("fp16", (3, 5, 4), True, seed=17, device="cpu") + second, _ = oracle.prepare_tensors("fp16", (3, 5, 4), True, seed=17, device="cpu") + other, _ = oracle.prepare_tensors("fp16", (3, 5, 4), True, seed=18, device="cpu") + self.assertTrue(torch.equal(state, torch.random.get_rng_state())) + self.assertTrue(all(torch.equal(a, b) for a, b in zip(first, second))) + self.assertFalse(torch.equal(first[0], other[0])) + for tensor in first: + self.assertEqual(tensor.dtype, torch.float16) + self.assertTrue(bool(((tensor >= -2) & (tensor < 2)).all())) + self.assertTrue(torch.equal(tensor, tensor.round())) + a, b, bias = first + self.assertEqual(a.shape, (1, 3, 4)) + self.assertTrue(a.is_contiguous()) + self.assertEqual(b.shape, (1, 4, 5)) + self.assertEqual(b.stride(), (20, 1, 4)) + self.assertTrue(b.transpose(1, 2).is_contiguous()) + self.assertEqual(bias.shape, (3,)) + self.assertEqual(len(outputs), 1) + self.assertEqual(outputs[0].shape, (1, 3, 5)) + self.assertEqual(outputs[0].dtype, torch.float16) + + def test_hand_computed_product_and_row_bias(self): + a = torch.tensor([[[1, 2], [3, 4]]], dtype=torch.float16) + b = torch.tensor([[[5, 7, 9], [6, 8, 10]]], dtype=torch.float16) + c = torch.tensor([[[17, 23, 29], [39, 53, 67]]], dtype=torch.float16) + oracle.verify("fp16", a, b, c) + bias = torch.tensor([10, -10], dtype=torch.float16) + biased = torch.tensor([[[27, 33, 39], [29, 43, 57]]], dtype=torch.float16) + oracle.verify("fp16", a, b, bias, biased) + with self.assertRaises(AssertionError): + oracle.verify("fp16", a, b, bias, c) + + def test_bias_is_added_before_half_rounding(self): + a = torch.tensor([[[1, 1]]], dtype=torch.float16) + b = torch.tensor([[[2048], [1]]], dtype=torch.float16) + bias = torch.tensor([1], dtype=torch.float16) + oracle.verify("fp16", a, b, bias, torch.tensor([[[2050]]], dtype=torch.float16)) + with self.assertRaises(AssertionError): + oracle.verify("fp16", a, b, bias, torch.tensor([[[2048]]], dtype=torch.float16)) + + def test_nonfinite_and_invalid_results_fail(self): + a = torch.ones((1, 2, 3), dtype=torch.float16) + b = torch.ones((1, 3, 4), dtype=torch.float16) + for bad in (torch.full((1, 2, 4), float("nan"), dtype=torch.float16), + torch.full((1, 2, 4), float("inf"), dtype=torch.float16), + torch.ones((1, 2, 4), dtype=torch.float16), + torch.full((2, 4), 3, dtype=torch.float16), + torch.full((1, 2, 4), 3, dtype=torch.float32)): + with self.subTest(shape=bad.shape, dtype=bad.dtype): + with self.assertRaises(AssertionError): + oracle.verify("fp16", a, b, bad) + a[0, 0, 0] = float("nan") + with self.assertRaises(AssertionError): + oracle.verify("fp16", a, b, torch.full((1, 2, 4), float("nan"), dtype=torch.float16)) + + def test_only_fp16_positive_shapes_are_supported(self): + for family, dims in (("fp32", (2, 3, 4)), ("fp16", (0, 3, 4)), + ("fp16", (2, 3)), ("fp16", (2, 3, 1.5))): + with self.subTest(family=family, dims=dims): + with self.assertRaises(ValueError): + oracle.prepare_tensors(family, dims, device="cpu") + with self.assertRaises(ValueError): + oracle.verify("fp32") + + def test_incompatible_inputs_and_bias_fail(self): + a = torch.ones((1, 2, 3), dtype=torch.float16) + b = torch.ones((1, 3, 4), dtype=torch.float16) + c = torch.full((1, 2, 4), 3, dtype=torch.float16) + for inputs in ((a.float(), b, c), (a[0], b, c), + (a, b[:, :2], c), (a, b, torch.ones(4, dtype=torch.float16), c)): + with self.subTest(shapes=[x.shape for x in inputs]): + with self.assertRaises(AssertionError): + oracle.verify("fp16", *inputs) + + def test_outer_autocast_does_not_change_reference_rounding(self): + a = torch.tensor([[[1, 1]]], dtype=torch.float16) + b = torch.tensor([[[2048], [1]]], dtype=torch.float16) + expected = torch.tensor([[[2050]]], dtype=torch.float16) + bias = torch.tensor([1], dtype=torch.float16) + with torch.autocast("cpu", dtype=torch.bfloat16): + oracle.verify("fp16", a, b, bias, expected) + + def test_non_direct_measurement_is_rejected_before_launch(self): + with self.assertRaises(ValueError): + compare.measure(torch, oracle, None, [], None, 0, 1, 2) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA unavailable") + def test_measured_noop_cannot_reuse_a_valid_output(self): + inputs, outputs = oracle.prepare_tensors("fp16", (2, 3, 4), device="cuda") + outputs[0].copy_(torch.bmm(*inputs)) + + def noop(*args, stream): + """Deliberately leave a poisoned result unwritten.""" + + with self.assertRaises(AssertionError): + compare.measure(torch, oracle, noop, inputs, outputs[0], 0, 1, 1) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA unavailable") + def test_direct_measurement_uses_explicit_stream_and_verifies_result(self): + inputs, outputs = oracle.prepare_tensors("fp16", (2, 3, 4), device="cuda") + observed_streams = [] + + def product(a, b, c, stream): + """Use a real GPU product to exercise timing and stream ordering.""" + observed_streams.append((stream.cuda_stream, torch.cuda.current_stream().cuda_stream)) + torch.bmm(a, b, out=c) + + latency = compare.measure(torch, oracle, product, inputs, outputs[0], 2, 3, 1) + self.assertGreater(latency, 0) + self.assertEqual(len(observed_streams), 5) + self.assertTrue(all(explicit == current for explicit, current in observed_streams)) + oracle.verify("fp16", *inputs, outputs[0]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA unavailable") + def test_seeded_inputs_are_identical_across_devices(self): + cpu, _ = oracle.prepare_tensors("fp16", (3, 5, 4), True, seed=42, device="cpu") + gpu, _ = oracle.prepare_tensors("fp16", (3, 5, 4), True, seed=42, device="cuda") + self.assertTrue(all(torch.equal(a, b.cpu()) for a, b in zip(cpu, gpu))) + + def test_clock_snapshot_is_explicit_and_best_effort(self): + snapshot = oracle.clock_state("cpu")["clock_snapshot"] + self.assertEqual(snapshot["status"], "unavailable") + self.assertEqual(snapshot["phase"], "post-timing") + self.assertIn("timestamp_utc", snapshot) + + @unittest.skipUnless(torch.cuda.is_available() and shutil.which("nvidia-smi"), "GPU telemetry unavailable") + def test_clock_snapshot_accepts_torch_gpu_uuid_format(self): + uuid = getattr(torch.cuda.get_device_properties(0), "uuid", None) + if uuid is None: + self.skipTest("PyTorch device UUID unavailable") + query_uuid = str(uuid) + if not query_uuid.startswith(("GPU-", "MIG-")): + query_uuid = "GPU-" + query_uuid + try: + result = subprocess.run( + ["nvidia-smi", "--id=" + query_uuid, + "--query-gpu=uuid,clocks.current.sm,clocks.current.memory,power.draw", + "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5, check=True, + ) + fields = result.stdout.strip().split(",") + for field in fields[1:]: + float(field) + except (subprocess.SubprocessError, ValueError): + self.skipTest("nvidia-smi clock telemetry unavailable") + snapshot = oracle.clock_state()["clock_snapshot"] + self.assertEqual(snapshot["status"], "available", snapshot) + self.assertEqual(snapshot["gpu_uuid"], fields[0].strip()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA unavailable") + def test_cuda_reference_restores_precision_after_success_and_failure(self): + old = torch.get_float32_matmul_precision() + try: + torch.set_float32_matmul_precision("high") + a = torch.tensor([[[1.0009765625, 1.0009765625]]], device="cuda", dtype=torch.float16) + b = torch.tensor([[[1024], [1024]]], device="cuda", dtype=torch.float16) + expected = torch.tensor([[[2050]]], device="cuda", dtype=torch.float16) + oracle.verify("fp16", a, b, expected) + self.assertEqual(torch.get_float32_matmul_precision(), "high") + with self.assertRaises(AssertionError): + oracle.verify("fp16", a, b, torch.zeros_like(expected)) + self.assertEqual(torch.get_float32_matmul_precision(), "high") + finally: + torch.set_float32_matmul_precision(old) + + +if __name__ == "__main__": + unittest.main() diff --git a/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/validate.py b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/validate.py new file mode 100644 index 0000000000..00ba5592de --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_cute/validate.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Check FP16 CuTe GEMM tails, persistent tiles, and five-stage pipeline reuse. + +Runs on Linux with an SM100 GPU. Every case uses the executable's exhaustive +CPU oracle and captured-graph replay check. Optional CUDA sanitizer runs add +memory/TMA and synchronization checks; these short runs are not benchmarks. +""" + +import argparse +import hashlib +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import tempfile + + +CASES = ( + ("full", "256,352,64", ()), + ("small-tails-bias", "256,8,8", ("--has-bias",)), + ("persistent-tails-bias", "512,712,72", ("--has-bias", "--clusters", "1")), + ("five-stages", "512,704,320", ("--clusters", "1")), + ("stage-wrap-tails-bias", "512,712,392", ("--has-bias", "--clusters", "1")), + ("multicluster-phase-wraps", "768,1064,704", ("--has-bias", "--clusters", "2")), +) + + +def run_command(command, timeout): + """Bound the entire Linux process group, including sanitizer children.""" + with subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) as process: + try: + output, _ = process.communicate(timeout=timeout) + return process.returncode, output + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + output, _ = process.communicate() + return None, output + + +def validate(binary, directory, timeout, sanitizer): + report = { + "binary": str(binary), + "binary_sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + "timeout_seconds": timeout, + "checks": [], + "status": "running", + } + + def save_report(): + (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + def check(name, command): + try: + status, output = run_command(command, timeout) + except OSError as error: + status, output = -1, str(error) + (directory / f"{name}.log").write_text(output) + entry = {"name": name, "command": command, "exit_code": status} + report["checks"].append(entry) + if status != 0: + entry["status"] = "timeout" if status is None else "failed" + report["status"] = "failed" + save_report() + reason = f"timeout after {timeout:g}s" if status is None else f"exit {status}" + print(f"FAIL {name}: {reason}\n{output}", file=sys.stderr, flush=True) + return None + return entry + + for name, shape, extra in CASES: + result_path = directory / f"{name}.json" + # A failed or malformed invocation must not reuse an older result. + result_path.unlink(missing_ok=True) + command = [ + str(binary), "--rust-only", "--mnk", shape, *extra, "--warmup", "1", "--iters", "2", + "--graph-launches", "2", "--json", str(result_path), + ] + entry = check(name, command) + if entry is None: + return False + try: + result = json.loads(result_path.read_text()) + if result["verification"] != "cpu-pass": + raise ValueError(f'expected cpu-pass, got {result["verification"]!r}') + if result["mnk"] != [int(value) for value in shape.split(",")]: + raise ValueError("result dimensions differ from the requested case") + except (OSError, ValueError, KeyError, TypeError) as error: + entry["status"] = "failed" + report["status"] = "failed" + save_report() + print(f"FAIL {name}: invalid verification report: {error}", file=sys.stderr) + return False + entry.update(status="passed", result=result) + save_report() + print(f"PASS {name}: {shape}", flush=True) + + if sanitizer: + for tool in ("memcheck", "synccheck"): + command = ["compute-sanitizer", "--tool", tool, "--error-exitcode", "1"] + if tool == "memcheck": + command += ["--check-tensor-ops", "yes"] + command += [ + str(binary), "--rust-only", "--mnk", "512,712,392", "--has-bias", "--clusters", "1", + "--warmup", "0", "--iters", "1", "--graph-launches", "1", + ] + entry = check(tool, command) + if entry is None: + return False + entry["status"] = "passed" + save_report() + print(f"PASS {tool}", flush=True) + + report["status"] = "passed" + save_report() + return True + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--binary", type=Path, + default=Path(__file__).resolve().parent / "target/release/fp16_gemm_256x352_cute", + help="compiled executable; override when Cargo uses a custom target directory", + ) + parser.add_argument("--output-dir", type=Path, help="preserve per-case logs, JSON, and report.json") + parser.add_argument("--timeout", type=float, default=60, help="seconds per check (default: 60)") + parser.add_argument("--sanitizer", action="store_true", help="also run memcheck and synccheck") + args = parser.parse_args() + if not sys.platform.startswith("linux"): + parser.error("this CUDA runner requires Linux process-group handling") + if not 0 < args.timeout < float("inf"): + parser.error("--timeout must be a finite positive number") + try: + binary = args.binary.resolve(strict=True) + except OSError as error: + parser.error(f"build the example first or pass --binary: {error}") + if args.output_dir is not None: + directory = args.output_dir.resolve() + directory.mkdir(parents=True, exist_ok=True) + passed = validate(binary, directory, args.timeout, args.sanitizer) + print(f"Validation artifacts: {directory}", flush=True) + else: + with tempfile.TemporaryDirectory(prefix="fp16-cute-validate-") as temporary: + passed = validate(binary, Path(temporary), args.timeout, args.sanitizer) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.lock new file mode 100644 index 0000000000..5bbb281406 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.lock @@ -0,0 +1,805 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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 = "nvfp4_gemv_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.toml new file mode 100644 index 0000000000..c0c163c98d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "nvfp4_gemv_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "Packed NVFP4 GEMV using typed CuTe semantic views" +publish = false + +# Standalone example, matching the other cuda-oxide example crates. +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs" } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core", features = ["f16"] } diff --git a/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/README.md b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/README.md new file mode 100644 index 0000000000..57a567cccb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/README.md @@ -0,0 +1,137 @@ +# Packed FP4 GEMV with typed scales + +GEMV means “matrix times vector”: + +```text +A[M,K] × x[K] = y[M] + +y[row] = A[row,0] × x[0] + + A[row,1] × x[1] + + ... +``` + +This example stores the matrix and vector as packed E2M1 values. E2M1 is a +four-bit floating-point format, so two values fit in one byte. Every group of +16 values has an eight-bit UE8M0 scale. UE8M0 stores a power-of-two multiplier. + +```text +16 packed values: v0 v1 v2 ... v15 + │ + └── one scale applies to the whole group + +value used by the dot product = E2M1 value × UE8M0 scale +``` + +The exact storage contract matters more than the short format name: this +example uses E2M1 data, UE8M0 scales, and groups of 16. + +## The Tensor flow + +One GPU thread computes one output row. `BlockScaledTensor` keeps the packed +values, their scales, and their layout together. + +```text +packed A bytes + A scales ──► BlockScaledTensor + │ + row view + │ + K=64 tile ──► load + convert ──┐ + │ +packed x bytes + x scales ──► BlockScaledTensor │ + │ │ + row view │ + │ │ + K=64 tile ──► load + convert ──┤ + ▼ + scaled dot product + │ + repeat across K tiles + │ + ▼ + one f16 y value +``` + +The important types are visible at the point where the views are made: + +```rust +type MatrixA<'a> = + BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, 16, KMajor>; +type VectorB<'a> = + BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, 16, KMajor>; +``` + +They record: + +- the packed value format (`Float4E2M1FN`); +- the scale format (`Float8E8M0FNU`); +- the 16-value scale group; +- the K-contiguous layout; +- whether a view plays the matrix or vector role. + +`dot_accumulate` accepts the matching matrix tile and vector tile, so a role +or format mismatch is rejected before the kernel runs. + +## Shape and device requirements + +The packed FP4 operations in this example target `sm_120a`. Use a compatible +Blackwell GPU to execute it. + +The command-line shapes must follow these rules: + +```text +M > 0 and M is a multiple of 128 +K > 0 and K is a multiple of 64 +L > 0 +``` + +`L` is the batch count. The default problem is `M=512`, `K=256`, `L=1`. + +## Build and run it + +From the `cuda-oxide/` directory, install the pinned official CUTLASS 4.7 compiler +once: + +```bash +cargo oxide toolchain install cutlass +``` + +Build or run this example through the translation backend: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide build nvfp4_gemv_cute --arch sm_120a + +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run nvfp4_gemv_cute --arch sm_120a -- \ + --m 512 --k 256 --l 1 +``` + +Use `--help` to list the shape and device options. Every run performs all +correctness checks. + +A successful run compares every GEMV output `f16` bit pattern with the host +calculation. This example is the semantic checkpoint for block-scaled views, +K tiles, packed loads, and a reduction. + +## Single-kernel comparison + +Measured on 2026-08-18 at commit `aafe6e47fb4e` on an RTX 5090 (`sm_120`), +Nsight Systems 2026.1.3 reported these kernel-active medians for `M=512`, +`K=256`, `L=1`: + +| Implementation | Median | +| :--- | ---: | +| CUTLASS 4.7 translation | **2.368 µs** | +| CuTeDSL 4.6.2 | **2.240 µs** | + +Translation is 0.128 µs, or 5.71%, longer. The measured translation cubin is +15,416 bytes with SHA-256 +`5ce377b00a7b73fb8ce1e083aeb83780bfaa4c4c8af3a987a00f1f545ec3b277`. +All 102 translation runs passed bit-exact validation with output hash +`f1c627726a82c6bd`, matching the separate CuTeDSL correctness run. + +Each observation is one main-kernel launch; setup and CuTeDSL conversion +kernels are excluded. The translation samples came from 102 independently +validated processes because this host harness launches once per process, +whereas CuTeDSL reused one process. That cache/context difference means this +GEMV comparison is not a perfectly matched hot-reuse experiment. diff --git a/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/src/main.rs new file mode 100644 index 0000000000..d6e0e406fd --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/src/main.rs @@ -0,0 +1,467 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Functional NVFP4 GEMV through retained CuTe semantic operations. +//! +//! The device algorithm intentionally mirrors the CuTeDSL reference: +//! - A is M x K packed E2M1, b is K packed E2M1; +//! - one E8M0 scale is applied to each group of 16 logical values; +//! - a 128-thread CTA computes 128 output rows, one row per thread; +//! - K advances in tiles of 64 and accumulation is f32; +//! - C is stored as f16. +//! +//! cute-rs binds each packed value tensor to its scales, element formats, and +//! layout. Thread-row and K-tile views make ownership explicit while each +//! backend chooses how the semantic load and dot operations are implemented. +//! PASS requires every output f16 bit pattern to match the host reference. + +#![feature(f16)] + +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; +use cuda_core::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cute_rs::{ + BlockScaledTensor, Float4E2M1FN, Float8E8M0FNU, KMajor, Mkl, Nkl, Sm1xxBlockScaleKMajor, +}; + +/// Rows computed per CTA, with one output row per thread. +const CTA_M: usize = 128; +/// Logical K values consumed by each block-scaled dot step. +const K_TILE: usize = 64; +/// Logical packed values sharing one UE8M0 scale byte. +const SF_VEC: usize = 16; +/// Default number of matrix rows. +const DEFAULT_M: usize = 512; +/// Default reduction length, divisible by the K64 tile. +const DEFAULT_K: usize = 256; +/// Default number of independent matrix/vector batches. +const DEFAULT_L: usize = 1; + +/// Packed E2M1 matrix with UE8M0 scales and K-major matrix coordinates. +type MatrixA<'a> = BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, SF_VEC, KMajor>; + +/// Packed E2M1 vector with UE8M0 scales and K-major vector coordinates. +type VectorB<'a> = BlockScaledTensor<'a, Float4E2M1FN, Float8E8M0FNU, SF_VEC, KMajor>; + +/// Device entry point for one-thread-per-row block-scaled GEMV. +#[cuda_module] +mod kernels { + use super::*; + + /// One thread computes one C row, matching the CuTeDSL reference kernel. + /// + /// # Safety + /// + /// Launch exactly `block=(128,1,1)` and `grid=(m/128,1,L)`, with positive + /// `m % 128 == 0` and `k % 64 == 0`. A/B must contain at least + /// `L*m*k/2` and `L*k/2` packed bytes, with 16-byte-aligned bases. SFA/SFB + /// must use the canonical layout and contain at least + /// `L*ceil(m/128)*ceil(k/64)*512` and `L*ceil(k/64)*512` bytes, with + /// four-byte-aligned bases. C must contain `L*m` disjoint f16 elements. + /// No other block/grid axes may contain work. + #[kernel] + pub fn nvfp4_gemv( + a: &[u8], + b: &[u8], + sfa: &[u8], + sfb: &[u8], + mut c: DisjointSlice, + m: usize, + k: usize, + ) { + let row = thread::blockIdx_x() as usize * CTA_M + thread::threadIdx_x() as usize; + let batch = thread::blockIdx_z() as usize; + + // Bind each value tensor to its scale tensor once. Element formats, + // logical mode order, scale-vector size, and physical scale layout + // are now part of these Rust types rather than comments on byte + // slices. + let a = MatrixA::from_slices(a, sfa, m, k); + let b = VectorB::from_slices(b, sfb, 1, k); + let a_row = a.thread_row(batch, row); + let b_row = b.thread_row(batch, 0); + + let mut acc = 0.0f32; + // Advance one logical K tile at a time. + let mut k_base = 0usize; + + while k_base < k { + let k_tile = k_base / K_TILE; + // k_tile creates typed views only. load performs two 16-byte value + // transactions plus one four-byte scale transaction per operand, + // then hoists each scale conversion out of the dot product. + let (a_tile, b_tile) = + unsafe { (a_row.k_tile(k_tile).load(), b_row.k_tile(k_tile).load()) }; + acc = a_tile.dot_accumulate(b_tile, acc); + k_base += K_TILE; + } + + // SAFETY: (batch, row) is unique for every live thread and the host + // allocated L*M output elements. + unsafe { + c.as_mut_ptr().add(batch * m + row).write(acc as f16); + } + } +} + +/// Validated dimensions and CUDA device selected by the CLI. +#[derive(Clone, Copy, Debug)] +struct Args { + /// Number of matrix rows and output values per batch. + m: usize, + /// Reduction length shared by the matrix and vector. + k: usize, + /// Number of independent matrix/vector batches. + l: usize, + /// CUDA device ordinal after visible-device remapping. + device: usize, +} + +/// Print the accepted CLI options and exit successfully. +fn usage() -> ! { + println!("Usage: nvfp4_gemv_cute [--m M] [--k K] [--l L] [--device ORDINAL]"); + std::process::exit(0); +} + +/// Read dimensions and device selection, rejecting unsupported shapes or overflows. +fn parse_args() -> Args { + let mut parsed = Args { + m: DEFAULT_M, + k: DEFAULT_K, + l: DEFAULT_L, + device: 0, + }; + let mut args = std::env::args().skip(1); + while let Some(flag) = args.next() { + if flag == "--help" || flag == "-h" { + usage(); + } + let value = args + .next() + .unwrap_or_else(|| panic!("missing value after {flag}")); + match flag.as_str() { + "--m" => parsed.m = value.parse().expect("invalid --m"), + "--k" => parsed.k = value.parse().expect("invalid --k"), + "--l" => parsed.l = value.parse().expect("invalid --l"), + "--device" => parsed.device = value.parse().expect("invalid --device"), + _ => panic!("unknown argument {flag}; use --help"), + } + } + assert!( + parsed.m > 0 && parsed.m.is_multiple_of(CTA_M), + "M must be a positive multiple of 128" + ); + assert!( + parsed.k > 0 && parsed.k.is_multiple_of(K_TILE), + "K must be a positive multiple of 64" + ); + assert!(parsed.l > 0, "L must be positive"); + validate_problem_sizes(parsed.m, parsed.k, parsed.l); + parsed +} + +/// Round a nonnegative extent up to a whole number of nonzero-sized tiles. +fn ceil_div(a: usize, b: usize) -> usize { + a.div_ceil(b) +} + +/// Multiply allocation dimensions, panicking with the buffer name on overflow. +fn checked_product(label: &str, factors: &[usize]) -> usize { + factors + .iter() + .try_fold(1usize, |product, &factor| product.checked_mul(factor)) + .unwrap_or_else(|| panic!("{label} size overflows usize")) +} + +/// Check allocation products and CUDA grid limits before creating buffers. +fn validate_problem_sizes(m: usize, k: usize, l: usize) { + let _ = checked_product("packed A", &[l, m, k / 2]); + let _ = checked_product("packed B", &[l, k / 2]); + let _ = checked_product("C", &[l, m]); + let _ = scale_storage_len(l, m, k); + let _ = scale_storage_len(l, 1, k); + let _ = u32::try_from(ceil_div(m, CTA_M)).expect("grid.x exceeds u32"); + let _ = u32::try_from(l).expect("grid.z exceeds u32"); +} + +/// Decode an E2M1 nibble, including its sign bit. +fn e2m1_to_f32(nibble: u8) -> f32 { + /// FP32 magnitudes for the eight unsigned E2M1 bit patterns. + const VALUES: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + let magnitude = VALUES[(nibble & 0x7) as usize]; + if nibble & 0x8 != 0 { + -magnitude + } else { + magnitude + } +} + +/// Decode a scale byte, including the smallest scale and the NaN encoding. +fn ue8m0_to_f32(byte: u8) -> f32 { + if byte == 0 { + f32::from_bits(0x0040_0000) + } else if byte == 0xff { + f32::from_bits(0x7fff_ffff) + } else { + f32::from_bits((byte as u32) << 23) + } +} + +/// Generate a deterministic matrix value from its batch, row, and K coordinate. +fn a_nibble(batch: usize, row: usize, k: usize) -> u8 { + (((batch % 16) * 5 + (row % 16) * 13 + (k % 16) * 7 + 3) % 16) as u8 +} + +/// Generate a deterministic vector value from its batch and K coordinate. +fn b_nibble(batch: usize, k: usize) -> u8 { + (((batch % 16) * 3 + (k % 16) * 11 + 1) % 16) as u8 +} + +/// Generate a matrix scale byte representing one-half, one, or two. +fn sfa_byte(batch: usize, row: usize, k_group: usize) -> u8 { + 126 + ((batch % 3 + row % 3 + k_group % 3) % 3) as u8 +} + +/// Generate a vector scale byte representing one-half, one, or two. +fn sfb_byte(batch: usize, k_group: usize) -> u8 { + 126 + ((batch % 3 + 2 * (k_group % 3)) % 3) as u8 +} + +/// Pack generated matrix values in batch/row/K order, low nibble first. +fn pack_a(m: usize, k: usize, l: usize) -> Vec { + let mut packed = Vec::with_capacity(checked_product("packed A", &[l, m, k / 2])); + for batch in 0..l { + for row in 0..m { + for byte_k in 0..k / 2 { + let low = a_nibble(batch, row, 2 * byte_k); + let high = a_nibble(batch, row, 2 * byte_k + 1); + packed.push(low | (high << 4)); + } + } + } + packed +} + +/// Pack generated vectors in batch/K order, low nibble first. +fn pack_b(k: usize, l: usize) -> Vec { + let mut packed = Vec::with_capacity(checked_product("packed B", &[l, k / 2])); + for batch in 0..l { + for byte_k in 0..k / 2 { + let low = b_nibble(batch, 2 * byte_k); + let high = b_nibble(batch, 2 * byte_k + 1); + packed.push(low | (high << 4)); + } + } + packed +} + +/// Return the padded byte capacity of the canonical K-major scale layout. +fn scale_storage_len(l: usize, mn: usize, k: usize) -> usize { + Sm1xxBlockScaleKMajor::::new(mn, k).storage_len(l) +} + +/// Map a batch, row, and K-group to its canonical scale-byte offset. +fn scale_index(batch: usize, row: usize, k_group: usize, mn: usize, k: usize) -> usize { + Sm1xxBlockScaleKMajor::::new(mn, k).crd2idx(batch, row, k_group) +} + +/// Fill valid scale coordinates and leave canonical-layout padding as raw zero bytes. +fn make_scales(m: usize, k: usize, l: usize) -> (Vec, Vec) { + // CuTeDSL's f32-zero initialization converts unused canonical-layout + // padding to raw UE8M0 byte zero. Valid coordinates overwrite it below. + let mut sfa = vec![0u8; scale_storage_len(l, m, k)]; + let mut sfb = vec![0u8; scale_storage_len(l, 1, k)]; + for batch in 0..l { + for row in 0..m { + for k_group in 0..k / SF_VEC { + sfa[scale_index(batch, row, k_group, m, k)] = sfa_byte(batch, row, k_group); + } + } + for k_group in 0..k / SF_VEC { + sfb[scale_index(batch, 0, k_group, 1, k)] = sfb_byte(batch, k_group); + } + } + (sfa, sfb) +} + +/// Compute every GEMV result in FP32 and round once to FP16. +fn reference(m: usize, k: usize, l: usize) -> Vec { + let mut expected = vec![0.0f16; checked_product("C", &[l, m])]; + for batch in 0..l { + for row in 0..m { + let mut acc = 0.0f32; + for kk in 0..k { + let k_group = kk / SF_VEC; + acc += ((e2m1_to_f32(a_nibble(batch, row, kk)) + * ue8m0_to_f32(sfa_byte(batch, row, k_group))) + * e2m1_to_f32(b_nibble(batch, kk))) + * ue8m0_to_f32(sfb_byte(batch, k_group)); + } + expected[batch * m + row] = acc as f16; + } + } + expected +} + +/// Require bit-exact agreement for every output and print its deterministic hash. +fn validate(got: &[f16], expected: &[f16]) { + assert_eq!(got.len(), expected.len()); + for (index, (&actual, &want)) in got.iter().zip(expected).enumerate() { + assert_eq!( + actual.to_bits(), + want.to_bits(), + "bit mismatch at output {index}: expected {} ({:#06x}), got {} ({:#06x})", + want as f32, + want.to_bits(), + actual as f32, + actual.to_bits(), + ); + } + println!( + "correctness: all {} f16 values bit-exact, output_hash={:016x}", + got.len(), + output_hash(got) + ); +} + +/// Hash FP16 bit patterns for reproducible output comparisons. +fn output_hash(values: &[f16]) -> u64 { + values.iter().fold(0xcbf2_9ce4_8422_2325, |hash, value| { + (hash ^ u64::from(value.to_bits())).wrapping_mul(0x0000_0100_0000_01b3) + }) +} + +/// Generate inputs, launch GEMV on the selected device, and verify every output bit. +fn main() { + let args = parse_args(); + println!( + "NVFP4 GEMV cute-rs typed block-scaled kernel: M={} K={} L={}", + args.m, args.k, args.l + ); + + let a = pack_a(args.m, args.k, args.l); + let b = pack_b(args.k, args.l); + let (sfa, sfb) = make_scales(args.m, args.k, args.l); + let expected = reference(args.m, args.k, args.l); + + let ctx = CudaContext::new(args.device).expect("failed to create CUDA context"); + let (major, minor) = ctx + .compute_capability() + .expect("failed to query compute capability"); + let device_name = ctx.device_name().expect("failed to query device name"); + println!( + "device={} name={device_name:?} compute_capability=sm_{}{}", + args.device, major, minor, + ); + let stream = ctx.default_stream(); + let module = kernels::load(&ctx).expect("failed to load embedded CUDA module"); + + let a_dev = DeviceBuffer::from_host(&stream, &a).expect("copy A"); + let b_dev = DeviceBuffer::from_host(&stream, &b).expect("copy b"); + let sfa_dev = DeviceBuffer::from_host(&stream, &sfa).expect("copy SFA"); + let sfb_dev = DeviceBuffer::from_host(&stream, &sfb).expect("copy SFB"); + let mut c_dev = DeviceBuffer::::zeroed(&stream, checked_product("C", &[args.l, args.m])) + .expect("allocate C"); + let config = LaunchConfig { + grid_dim: ( + u32::try_from(ceil_div(args.m, CTA_M)).expect("grid.x exceeds u32"), + 1, + u32::try_from(args.l).expect("grid.z exceeds u32"), + ), + block_dim: (CTA_M as u32, 1, 1), + shared_mem_bytes: 0, + }; + let max_grid = ctx + .launch_limits() + .expect("failed to query launch limits") + .max_grid_dim(); + assert!( + config.grid_dim.0 <= max_grid.0 && config.grid_dim.2 <= max_grid.2, + "grid {:?} exceeds device maximum {:?}", + config.grid_dim, + max_grid + ); + + // SAFETY: all buffers match the documented packed/canonical layouts; + // the exact grid covers M rows, and each thread owns one C value. + unsafe { + module.nvfp4_gemv( + &stream, config, &a_dev, &b_dev, &sfa_dev, &sfb_dev, &mut c_dev, args.m, args.k, + ) + } + .expect("GEMV launch failed"); + stream.synchronize().expect("GEMV synchronization failed"); + let got = c_dev.to_host_vec(&stream).expect("copy C to host"); + validate(&got, &expected); + println!("NVFP4 GEMV passed its bit-exact correctness check"); +} + +/// Tests for format decoding, canonical scale indexing, and nibble packing. +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeSet; + + /// Check all signed E2M1 values, including positive and negative zero. + #[test] + fn e2m1_decode_covers_all_nibbles() { + let positives = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]; + for (nibble, expected) in positives.into_iter().enumerate() { + assert_eq!(e2m1_to_f32(nibble as u8), expected); + assert_eq!(e2m1_to_f32(nibble as u8 | 8), -expected); + } + assert_eq!(e2m1_to_f32(0).to_bits(), 0x0000_0000); + assert_eq!(e2m1_to_f32(8).to_bits(), 0x8000_0000); + } + + /// Check scale-byte boundary values and the reserved NaN encoding. + #[test] + fn ue8m0_decode_covers_boundaries_and_nan() { + assert_eq!(ue8m0_to_f32(0).to_bits(), 0x0040_0000); + assert_eq!(ue8m0_to_f32(1).to_bits(), 0x0080_0000); + assert_eq!(ue8m0_to_f32(126), 0.5); + assert_eq!(ue8m0_to_f32(127), 1.0); + assert_eq!(ue8m0_to_f32(128), 2.0); + assert_eq!(ue8m0_to_f32(254).to_bits(), 0x7f00_0000); + assert!(ue8m0_to_f32(255).is_nan()); + } + + /// Check that logical scale coordinates map to distinct, allocated bytes. + #[test] + fn canonical_scale_indices_are_unique_and_in_bounds() { + let (m, k, l) = (256, 128, 2); + let len = scale_storage_len(l, m, k); + let mut indices = BTreeSet::new(); + for batch in 0..l { + for row in 0..m { + for group in 0..k / SF_VEC { + let index = scale_index(batch, row, group, m, k); + assert!(index < len); + assert!(indices.insert(index)); + } + } + } + assert_eq!(indices.len(), l * m * k / SF_VEC); + + // Boundaries of the public ((32,4), (SFVec,4)) 512-byte atom: + // row-32 rotates the second M mode, row-128 advances the M atom, + // group-4 advances the K atom, and batch follows both rest modes. + assert_eq!(scale_index(0, 0, 0, m, k), 0); + assert_eq!(scale_index(0, 31, 0, m, k), 496); + assert_eq!(scale_index(0, 32, 0, m, k), 4); + assert_eq!(scale_index(0, 127, 3, m, k), 511); + assert_eq!(scale_index(0, 0, 4, m, k), 512); + assert_eq!(scale_index(0, 128, 0, m, k), 1_024); + assert_eq!(scale_index(1, 0, 0, m, k), 2_048); + } + + /// Check the low-then-high nibble order for matrix and vector inputs. + #[test] + fn nvfp4_packing_uses_low_nibble_first() { + let packed = pack_b(64, 1); + assert_eq!(packed[0], b_nibble(0, 0) | (b_nibble(0, 1) << 4)); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.lock new file mode 100644 index 0000000000..4b77312b50 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rmsnorm_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.toml new file mode 100644 index 0000000000..e5da6fdf66 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "rmsnorm_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "FP32 and BF16 RMSNorm with explicit storage and collective contracts" + +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core" } diff --git a/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/README.md b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/README.md new file mode 100644 index 0000000000..0c9db36697 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/README.md @@ -0,0 +1,46 @@ +# RMSNorm with warp row partitions + +This cuda-oxide example normalizes each matrix row and applies a weight vector: + +```text +inverse_rms = 1 / sqrt(mean(row * row) + epsilon) +output[col] = row[col] * inverse_rms * weight[col] +``` + +It runs with FP32 and BF16 storage, accumulating the square sum in FP32. +For BF16, the normalized value rounds to BF16 before multiplication by the +weight; the final result rounds again. + +## The row flow + +```text +row → StridedTensor → local_partition::<32>(lane) → sum local squares + │ + warp_reduction_sum_f32 + │ +weights → matching partition → normalize and multiply ────────┘ + │ +output ← StridedTensorMut ← store owned columns +``` + +One warp (32 threads) owns one row; four warps share a block. Lanes own every +32nd column. Short rows and column tails contribute zero to the reduction, +so every lane still reaches the collective. No row alignment is required. + +## Run it + +From `cuda-oxide/`, with the pinned toolchain, CUDA toolkit, the CUTLASS compiler +(`cargo oxide toolchain install cutlass`) and an RTX 5090: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run rmsnorm_cute --arch sm_120a +``` + +A successful run prints `PASS: FP32/BF16 RMSNorm`, reports numerical errors, +and confirms the exact BF16 rounding fixture and bit-preserving copies. +It exercises 45 shapes, including empty rows and partial blocks. + +See [numerics and validation](reference/validation.md) for the checked host +boundary, tolerances, test commands, and sanitizer commands. These are +synthetic correctness checks, without checkpoint or performance claims. diff --git a/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/reference/validation.md b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/reference/validation.md new file mode 100644 index 0000000000..a47e7f6a51 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/reference/validation.md @@ -0,0 +1,50 @@ +# Contracts and validation + +[Back to the example](../README.md). + +## Arithmetic and launch contract + +Standalone cuda-oxide correctness example using one full warp per row, four rows +per block, and explicit BF16 storage. Inputs and weights retain their storage +dtype; both kernels accumulate squares and normalize in FP32. The BF16 kernel +rounds normalized values to BF16 **before** multiplying BF16 weights, then rounds +the result to BF16. This preserves the selected Qwen3-TTS reference's arithmetic +boundary without adding model-specific compiler operations. + +The checked host boundary rejects zero widths, mismatched allocation lengths, +shape overflow, invalid epsilon and grid overflow. Zero rows skip allocation and +launch. All lanes in each participating warp reach the reusable +`cute_rs::arch::warp_reduction_sum_f32` collective; lanes beyond a short row contribute zero. +Scalar strided row accesses also support unaligned row starts. + +## Host and GPU checks + +From `cuda-oxide/`: + +```bash +cargo test --offline --manifest-path cute-rs-preview/examples/rmsnorm_cute/Cargo.toml +CARGO_NET_OFFLINE=true CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run rmsnorm_cute --arch sm_120a +/usr/local/cuda/bin/compute-sanitizer --tool memcheck --error-exitcode 1 \ + cute-rs-preview/examples/rmsnorm_cute/target/release/rmsnorm_cute +/usr/local/cuda/bin/compute-sanitizer --tool synccheck --error-exitcode 1 \ + cute-rs-preview/examples/rmsnorm_cute/target/release/rmsnorm_cute +``` + +## Acceptance criteria + +Validation covers widths 1, 7, 31, 32, 33, 128, 1024, 2048 and 2051; row counts +0, 1, 3, 4 and 5; and zero, constant, signed and tiny-valued inputs. An independent +FP64 reference checks stored input values. BF16 comparisons explicitly reproduce +both rounding boundaries using a separate integer rounding oracle. Tolerances +are `atol=1e-6, rtol=2e-5` for FP32 and `atol=1e-7, rtol=0.016` for BF16. +An additional exact BF16 fixture rejects the incorrect cast-once ordering, +independently of those numerical tolerances. + +A separate BF16 copy kernel exercises typed 16-byte vector loads/stores and +scalar tails at lengths 0, 1, 7, 8, 9 and 1031. Copies must preserve every storage +bit, including signed zeros, infinities, subnormals and NaN payloads. + +The fixtures keep magnitudes bounded: FP64 reference arithmetic must not hide +FP32 square overflow. This is synthetic operator validation, not checkpoint +validation, a performance benchmark or evidence of execution on SM100. diff --git a/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/src/main.rs new file mode 100644 index 0000000000..220a45d7a2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/src/main.rs @@ -0,0 +1,455 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! RMSNorm validation with FP32 accumulation and explicit BF16 rounding boundaries. + +use cuda_core::LaunchConfig; +use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; +use cute_rs::BFloat16; + +/// Shared FP32/BF16 row kernels and bit-preserving BF16 copy kernel. +#[path = "../../common/rmsnorm.rs"] +mod rmsnorm; +use rmsnorm::{ROWS_PER_BLOCK, kernels}; + +/// Logical dimensions include head/hidden widths and awkward row boundaries. +const WIDTHS: [usize; 9] = [1, 7, 31, 32, 33, 128, 1024, 2048, 2051]; +/// Empty and partial final blocks must preserve full-warp participation. +const ROW_COUNTS: [usize; 5] = [0, 1, 3, 4, 5]; +/// Epsilon used by this synthetic validation, not a checkpoint configuration. +const EPSILON: f32 = 1e-6; + +/// Validate dense `[rows, cols]` input/output and a shared `[cols]` weight vector. +/// +/// Lengths are element counts. Epsilon must be finite and positive; zero rows +/// return `None`, while nonempty shapes return the number of four-warp blocks. +fn checked_work( + rows: usize, + cols: usize, + input: usize, + weight: usize, + output: usize, + epsilon: f32, +) -> Result, &'static str> { + if cols == 0 { + return Err("RMSNorm width must be positive"); + } + if !epsilon.is_finite() || epsilon <= 0.0 { + return Err("RMSNorm epsilon must be finite and positive"); + } + let elements = rows + .checked_mul(cols) + .ok_or("matrix extent overflows usize")?; + if input != elements || output != elements || weight != cols { + return Err("RMSNorm input, weight and output lengths must match the shape"); + } + if rows == 0 { + return Ok(None); + } + let blocks = rows.div_ceil(ROWS_PER_BLOCK); + if blocks > i32::MAX as usize { + return Err("row count exceeds the CUDA grid-x limit"); + } + Ok(Some(blocks as u32)) +} + +/// Independent host rounding oracle, separate from the device BF16 implementation. +fn reference_bf16(value: f32) -> u16 { + let bits = value.to_bits(); + let retained = (bits >> 16) as u16; + if value.is_nan() { + return retained | 0x0040; + } + let discarded = bits & 0xffff; + let increment = discarded > 0x8000 || (discarded == 0x8000 && retained & 1 != 0); + retained.wrapping_add(u16::from(increment)) +} + +/// Read a BF16 storage pattern without using the device conversion helper. +fn reference_widen(bits: u16) -> f32 { + f32::from_bits(u32::from(bits) << 16) +} + +/// FP64 square sum and normalization reference on the actual stored inputs. +fn reference_row(input: &[f32], weight: &[f32], epsilon: f32) -> Vec { + let squares: f64 = input.iter().map(|&value| f64::from(value).powi(2)).sum(); + let inverse_rms = (squares / input.len() as f64 + f64::from(epsilon)) + .sqrt() + .recip(); + input + .iter() + .zip(weight) + .map(|(&value, &scale)| f64::from(value) * inverse_rms * f64::from(scale)) + .collect() +} + +/// Require finite results within an absolute-plus-relative error budget. +fn close(actual: f64, expected: f64, atol: f64, rtol: f64) -> bool { + actual.is_finite() + && expected.is_finite() + && (actual - expected).abs() <= atol + rtol * expected.abs() +} + +/// Match the kernel's four complete warps per block. +fn launch_config(blocks: u32) -> LaunchConfig { + LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Check dynamic shape and epsilon before allocating and launching the FP32 kernel. +#[allow(clippy::too_many_arguments)] // Keep every checked launch-boundary argument explicit. +fn run_f32( + module: &kernels::LoadedModule, + stream: &CudaStream, + input: &[f32], + weight: &[f32], + output: &mut [f32], + rows: usize, + cols: usize, + epsilon: f32, +) -> Result<(), Box> { + let Some(blocks) = checked_work(rows, cols, input.len(), weight.len(), output.len(), epsilon)? + else { + return Ok(()); + }; + let input_device = DeviceBuffer::from_host(stream, input)?; + let weight_device = DeviceBuffer::from_host(stream, weight)?; + let mut output_device = DeviceBuffer::from_host(stream, output)?; + // SAFETY: exact allocation extents were checked; each complete warp owns + // one row and each lane writes distinct columns. Copies synchronize lifetime. + unsafe { + module.rmsnorm_f32( + stream, + launch_config(blocks), + &input_device, + &weight_device, + &mut output_device, + rows, + cols, + epsilon, + ) + }?; + output_device.copy_to_host(stream, output)?; + Ok(()) +} + +/// Check dynamic shape and epsilon before allocating and launching the BF16 kernel. +#[allow(clippy::too_many_arguments)] // Mirrors the FP32 boundary and kernel parameters. +fn run_bf16( + module: &kernels::LoadedModule, + stream: &CudaStream, + input: &[BFloat16], + weight: &[BFloat16], + output: &mut [BFloat16], + rows: usize, + cols: usize, + epsilon: f32, +) -> Result<(), Box> { + let Some(blocks) = checked_work(rows, cols, input.len(), weight.len(), output.len(), epsilon)? + else { + return Ok(()); + }; + let input_device = DeviceBuffer::from_host(stream, input)?; + let weight_device = DeviceBuffer::from_host(stream, weight)?; + let mut output_device = DeviceBuffer::from_host(stream, output)?; + // SAFETY: same full-warp ownership and checked exact extents as run_f32. + unsafe { + module.rmsnorm_bf16( + stream, + launch_config(blocks), + &input_device, + &weight_device, + &mut output_device, + rows, + cols, + epsilon, + ) + }?; + output_device.copy_to_host(stream, output)?; + Ok(()) +} + +/// Generate bounded rows so the FP64 oracle does not hide FP32 overflow behavior. +fn input_values(rows: usize, cols: usize, pattern: usize) -> (Vec, Vec) { + /// Bounded signed row samples; the tiny pattern scales these by `1e-7`. + const SIGNED: [f32; 9] = [-3.0, 0.25, -0.5, 1.0, -1.25, 2.0, 0.0, 4.0, -0.125]; + let input = (0..rows * cols) + .map(|index| match pattern { + 0 => 0.0, + 1 => 2.0, + 2 => SIGNED[(index * 7 + index / cols) % SIGNED.len()], + _ => SIGNED[index % SIGNED.len()] * 1e-7, + }) + .collect(); + let weight = (0..cols) + .map(|col| [0.5078125, 1.0, -0.75, 0.0, 1.5][col % 5]) + .collect(); + (input, weight) +} + +/// Verify pure BF16 storage copies, including signed zeros, infinities and NaN bits. +fn validate_copy( + module: &kernels::LoadedModule, + stream: &CudaStream, +) -> Result> { + let mut checked = 0; + for length in [0, 1, 7, 8, 9, 1031] { + let input: Vec = (0..length) + .map(|i| { + let bits = [ + 0x0000, 0x8000, 0x7f80, 0xff80, 0x7f81, 0xffc5, 0x0001, 0xffff, 0x3f80, + ][i % 9]; + BFloat16::from_bits(bits) + }) + .collect(); + let mut output = vec![BFloat16::from_bits(0x5555); length]; + if length != 0 { + let input_device = DeviceBuffer::from_host(stream, &input)?; + let mut output_device = DeviceBuffer::from_host(stream, &output)?; + // SAFETY: exact equal lengths and each thread owns one aligned + // eight-element vector or scalar tail; allocations remain live. + unsafe { + module.copy_bf16( + stream, + LaunchConfig::for_num_elems(length.div_ceil(8) as u32), + &input_device, + &mut output_device, + ) + }?; + output_device.copy_to_host(stream, &mut output)?; + } + for (index, (actual, expected)) in output.iter().zip(&input).enumerate() { + assert_eq!( + actual.to_bits(), + expected.to_bits(), + "BF16 copy length={length} index={index}" + ); + } + checked += length; + } + Ok(checked) +} + +/// Validate the standalone operators on device zero; this is not a model benchmark. +fn main() -> Result<(), Box> { + let context = CudaContext::new(0)?; + let stream = context.default_stream(); + let module = kernels::load(&context)?; + let mut checked = 0; + let mut max_error_f32 = 0.0_f64; + let mut max_error_bf16 = 0.0_f64; + assert!( + run_f32( + &module, + &stream, + &[1.0], + &[], + &mut [f32::NAN], + 1, + 1, + EPSILON + ) + .is_err() + ); + assert!( + run_bf16( + &module, + &stream, + &[BFloat16::from_bits(0)], + &[], + &mut [BFloat16::from_bits(0)], + 1, + 1, + EPSILON + ) + .is_err() + ); + for rows in ROW_COUNTS { + for cols in WIDTHS { + for pattern in 0..4 { + let (input, weight) = input_values(rows, cols, pattern); + let mut output = vec![f32::NAN; input.len()]; + run_f32( + &module, + &stream, + &input, + &weight, + &mut output, + rows, + cols, + EPSILON, + )?; + for row in 0..rows { + let expected = + reference_row(&input[row * cols..(row + 1) * cols], &weight, EPSILON); + for (col, &expected) in expected.iter().enumerate() { + let actual = f64::from(output[row * cols + col]); + assert!( + close(actual, expected, 1e-6, 2e-5), + "FP32 rows={rows} cols={cols} pattern={pattern} row={row} col={col}: expected={expected} actual={actual}" + ); + max_error_f32 = max_error_f32.max((actual - expected).abs()); + } + } + let input: Vec = input + .into_iter() + .map(|v| BFloat16::from_bits(reference_bf16(v))) + .collect(); + let weight: Vec = weight + .into_iter() + .map(|v| BFloat16::from_bits(reference_bf16(v))) + .collect(); + let mut output = vec![BFloat16::from_bits(0x7fc0); input.len()]; + run_bf16( + &module, + &stream, + &input, + &weight, + &mut output, + rows, + cols, + EPSILON, + )?; + for row in 0..rows { + let stored: Vec = input[row * cols..(row + 1) * cols] + .iter() + .map(|v| reference_widen(v.to_bits())) + .collect(); + let normalized = reference_row(&stored, &vec![1.0; cols], EPSILON); + for col in 0..cols { + let normalized = reference_widen(reference_bf16(normalized[col] as f32)); + let expected = f64::from(reference_widen(reference_bf16( + normalized * reference_widen(weight[col].to_bits()), + ))); + let actual = f64::from(reference_widen(output[row * cols + col].to_bits())); + assert!( + close(actual, expected, 1e-7, 0.016), + "BF16 rows={rows} cols={cols} pattern={pattern} row={row} col={col}: expected={expected} actual={actual}" + ); + max_error_bf16 = max_error_bf16.max((actual - expected).abs()); + } + } + checked += input.len(); + } + } + } + + // An exact separate fixture rejects a kernel that casts only once, even + // though a loose BF16 tolerance alone could accept that semantic error. + let mut fixture = [BFloat16::from_bits(0); 2]; + run_bf16( + &module, + &stream, + &[BFloat16::from_bits(0x3f80), BFloat16::from_bits(0x4000)], + &[BFloat16::from_bits(0x3f02); 2], + &mut fixture, + 1, + 2, + EPSILON, + )?; + assert_eq!( + fixture[0].to_bits(), + 0x3ea5, + "cast-before-weight fixture; cast-once would produce 0x3ea4" + ); + let copied = validate_copy(&module, &stream)?; + println!( + "PASS: FP32/BF16 RMSNorm, {} shapes, four patterns, {checked} outputs per dtype", + WIDTHS.len() * ROW_COUNTS.len() + ); + println!("FP32 max absolute error={max_error_f32:.8e}; atol=1e-6, rtol=2e-5"); + println!( + "BF16 max absolute error={max_error_bf16:.8e}; atol=1e-7, rtol=0.016; exact cast-before-weight fixture passed" + ); + println!( + "BF16 typed vector/scalar copies preserved {copied} storage patterns; empty rows skipped." + ); + Ok(()) +} + +/// Host checks for RMSNorm shapes, rounding boundaries, and reference arithmetic. +#[cfg(test)] +mod tests { + use super::*; + + /// Reject malformed row storage, nonpositive epsilon, and overflowing matrix extents. + #[test] + fn boundary_rejects_shape_epsilon_and_overflow_errors() { + for (rows, cols, input, weight, output, epsilon) in [ + (1, 0, 0, 0, 0, 1e-6), + (2, 3, 5, 3, 6, 1e-6), + (2, 3, 6, 2, 6, 1e-6), + (2, 3, 6, 3, 5, 1e-6), + (0, 3, 1, 3, 0, 1e-6), + (0, 3, 0, 2, 0, 1e-6), + (1, 1, 1, 1, 1, 0.0), + (1, 1, 1, 1, 1, -1.0), + (1, 1, 1, 1, 1, f32::NAN), + (1, 1, 1, 1, 1, f32::INFINITY), + (usize::MAX, 2, 0, 2, 0, 1e-6), + ] { + assert!(checked_work(rows, cols, input, weight, output, epsilon).is_err()); + } + } + + /// Check empty launches, four-row block tails, and the signed CUDA grid limit. + #[test] + fn boundary_skips_empty_rows_and_covers_partial_blocks() { + assert_eq!(checked_work(0, 128, 0, 128, 0, 1e-6), Ok(None)); + for (rows, blocks) in [(1, 1), (3, 1), (4, 1), (5, 2)] { + assert_eq!( + checked_work(rows, 33, rows * 33, 33, rows * 33, 1e-6), + Ok(Some(blocks)) + ); + } + if usize::BITS > 32 { + let rows = (i32::MAX as usize + 1) * 4; + assert!(checked_work(rows, 1, rows, 1, rows, 1e-6).is_err()); + } + } + + /// Verify independent BF16 tie handling, signed zero, and a retained NaN encoding. + #[test] + fn rounding_reference_uses_ties_to_even_and_preserves_nan() { + assert_eq!(reference_bf16(f32::from_bits(0x3f80_8000)), 0x3f80); + assert_eq!(reference_bf16(f32::from_bits(0x3f81_8000)), 0x3f82); + assert_eq!(reference_bf16(-0.0), 0x8000); + assert!(reference_widen(reference_bf16(f32::from_bits(0x7f80_0001))).is_nan()); + } + + /// Check analytic zero/constant-row RMSNorm values against the FP64 reference. + #[test] + fn reference_normalizes_zero_and_constant_rows() { + assert_eq!(reference_row(&[0.0, 0.0], &[1.0, -2.0], 1e-6), [0.0, -0.0]); + let out = reference_row(&[2.0, 2.0], &[1.0, -2.0], 1e-6); + assert!((out[0] - 1.0).abs() < 2e-7); + assert!((out[1] + 2.0).abs() < 3e-7); + } + + /// Require the intermediate BF16 cast that changes the weighted result by one bit. + #[test] + fn bf16_fixture_distinguishes_cast_before_weight_from_cast_once() { + let normalized = reference_row(&[1.0, 2.0], &[1.0, 1.0], 1e-6)[0] as f32; + let weight = reference_widen(0x3f02); + let correct = reference_bf16(reference_widen(reference_bf16(normalized)) * weight); + let cast_once = reference_bf16(normalized * weight); + assert_eq!(correct, 0x3ea5); + assert_eq!(cast_once, 0x3ea4); + assert_ne!(correct, cast_once); + } + + /// Ensure finite-value tolerance checks reject NaNs, infinity, and excessive error. + #[test] + fn acceptance_rejects_nonfinite_or_incorrect_outputs() { + assert!(close(1.000_01, 1.0, 1e-6, 2e-5)); + assert!(!close(1.001, 1.0, 1e-6, 2e-5)); + assert!(!close(f64::NAN, 1.0, 1e-6, 2e-5)); + assert!(!close(f64::INFINITY, 1.0, 1e-6, 2e-5)); + assert!(!close(1.0, f64::NAN, 1e-6, 2e-5)); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.lock new file mode 100644 index 0000000000..70e3b9f5c9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.lock @@ -0,0 +1,806 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rope_cache_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", +] + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.toml new file mode 100644 index 0000000000..c86f029632 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "rope_cache_cute" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" +description = "Checked BF16 rotary positions, owned KV cache and grouped-query attention on SM120" + +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core" } diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/README.md b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/README.md new file mode 100644 index 0000000000..5332ff690e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/README.md @@ -0,0 +1,49 @@ +# Rotary positions, KV cache, and BF16 attention + +This cuda-oxide example adds position information to queries and keys, saves +earlier keys and values in a cache, and reuses them for attention. It starts +with projected, normalized queries and keys, unrotated values, and explicit +token positions: + +```text +query → rotary pairs ───────────────────────────────┐ +key → rotary pairs → append to owned key cache ───┼→ attention → output +value ───────────────→ append to owned value cache ┘ +``` + +Rotary embeddings pair the first half of each head with the second half. +`StridedTensor` loads each pair into a register tile; both products round to +BF16 before their sum. Standard RoPE and three-axis interleaved MRoPE share +this kernel, using different host-generated coefficient tables. + +## The CuTe flow + +`AttentionShape` checks query/KV head grouping, and `CacheShape` checks append +capacity. The owned cache stores `[token, head, coordinate]` values and exposes +only its completed prefix to attention. Appends and attention calls wait for +GPU completion before returning. + +For attention, CuTe cooperative copies stage 64-key tiles in shared memory. +Typed layouts assign dot-product and output coordinates to threads. A second +kernel merges the split results. Query-key scores, softmax probabilities, and +weighted values are computed in FP32; only the output rounds to BF16. Attention uses 128-wide heads and supports +causal or noncausal windows, including grouped-query attention. + +## Run it + +From `cuda-oxide/`, with the pinned toolchain, CUDA toolkit, the CUTLASS compiler +(`cargo oxide toolchain install cutlass`) and an RTX 5090: + +```bash +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run rope_cache_cute --arch sm_120a +``` + +A successful run prints `PASS RoPE/cache` and `PASS BF16 attention` cases, +followed by reference comparison counts. It checks whole and chunked appends, +cache guards, exact rotary rounding, and prefill versus incremental attention. + +See [contracts and validation](reference/validation.md) for MRoPE axis mapping, +cache readiness, numerical tolerances, sanitizer commands, and recorded GPU +evidence. Host trigonometry has not been matched bit-for-bit with CUDA/PyTorch; +checkpoint inference and model performance are not established by this example. diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/reference/validation.md b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/reference/validation.md new file mode 100644 index 0000000000..87c3c3b048 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/reference/validation.md @@ -0,0 +1,112 @@ +# Contracts and validation + +[Back to the example](../README.md). + +## Geometry and rotary arithmetic + +This cuda-oxide example starts transformer integration with already projected +and normalized Q/K, unrotated V, and explicit positions. It uses `cute-rs` for +BF16 storage and checked attention geometry, and `ctm-rs` for kernel interfaces +and instruction primitives. It computes causal/noncausal attention over the +completed cache prefix; it does not yet run a model. + +`AttentionShape` keeps query heads, KV heads and head width independent +of hidden width. `CacheShape` checks capacity and append token ranges. +The example's `KvCache` privately owns K/V allocations; `RotatedQuery` +retains query dimensions. `HostTables` prevents mixing different rotary head +widths. Dynamic input lengths, grouped heads, byte extents, contexts, positions +and launch sizes are checked before execution. + +The runner rotates split-half pairs: coordinate `i` pairs with `i+D/2`, not +with its adjacent coordinate. Both BF16 products round before their BF16 sum. +The predictor uses ordinary RoPE; Talker interleaved MRoPE uses `[24,20,20]` +sections at D=128. Height frequencies occupy indices 1,4,…,58 and width +frequencies occupy 2,5,…,59; all remaining frequencies use temporal positions. +Three-axis positions remain separate from physical cache slots. + +## Cache lifecycle and attention schedule + +The cache stores rotated K and unrotated V in dense `[token,head,dimension]` +order. An append borrows inputs and the cache until GPU completion, then advances +the valid-token cursor. Input readiness is established with a context wait; all +launched work is synchronized even on enqueue-error paths. Validation failure +does not change the cursor or completed prefix. Reset discards the logical +prefix without promising to erase old allocation bytes. Empty appends launch +no kernels. This is a conservative synchronous API, without paging or eviction. + +`Runner::attend` borrows a `RotatedQuery` and `KvCache` and +returns owned `AttentionOutput` storage. The head width is fixed by the +Rust signature. `AttentionWindow` checks query counts, visible key counts and +causal offsets; a nonempty query requires a nonempty cache. The causal offset +counts physical cache slots, independently of rotary positions. Attention +preserves the cache and retains all input borrows through completion, including +enqueue-error paths. Empty queries launch nothing. + +The attention schedule adapts the SM120 decode example: one CTA per query/head +and 64-key split, followed by a stable merge. It supports Q16/KV8/D128 plus other +valid MHA/GQA/MQA head groupings, arbitrary sequence tails and cached query +offsets. Each partial CTA uses 128 threads and 35,104 shared bytes; each merge +uses 32 threads. Scratch has `queries * query_heads * ceil(keys/64) * 130` FP32 +elements, plus guards. This shared prefill/decode baseline is intended for +correct model composition; tensor-core prefill specialization remains future work. + +CuTe cooperative copies transport BF16 bits as `u16` inside a private BF16 tile +wrapper, restoring the BF16 type before arithmetic. The copy exporter supports +16-bit integer storage with the same byte-oriented copy and tail zero-fill as +FP16. QK, softmax, PV and split merging remain FP32; only the final output narrows +to BF16. This is a fused-attention numerical contract, not a claim of bitwise +agreement with an eager implementation that rounds scores or probabilities. + +## Host and GPU checks + +From `cuda-oxide/`: + +```bash +cargo test --offline -p cute-rs --features host --all-targets +cargo test --offline -p cute-rs --features host --doc +cargo test --offline --manifest-path cute-rs-preview/examples/rope_cache_cute/Cargo.toml +CARGO_NET_OFFLINE=true CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run rope_cache_cute --arch sm_120a +/usr/local/cuda/bin/compute-sanitizer --tool memcheck --error-exitcode 1 \ + cute-rs-preview/examples/rope_cache_cute/target/release/rope_cache_cute +/usr/local/cuda/bin/compute-sanitizer --tool synccheck --error-exitcode 1 \ + cute-rs-preview/examples/rope_cache_cute/target/release/rope_cache_cute +/usr/local/cuda/bin/compute-sanitizer --tool racecheck --error-exitcode 1 \ + cute-rs-preview/examples/rope_cache_cute/target/release/rope_cache_cute +``` + +## Acceptance criteria + +Validation covers Q16/KV8/D128 for both rotary modes, unequal position axes, +the last supported positions, Q6/KV2/D6 tails, empty input, cache exhaustion, +invalid lengths, reset and pending writes on another stream. One-shot and +chunked appends must match. An independent FP64 oracle with integer BF16 +rounding checks Q/K bits exactly; V bits must remain unchanged. Intermediate +snapshots check every completed prefix, unwritten slot and trailing guard. +An exact fixture requires `0x3eb4` where incorrectly rounding only once produces +`0x3eb5`. + +Attention checks every output against a separate FP64 dot/softmax/PV reference +with final integer BF16 rounding (`atol=2e-5`, `rtol=0.008`). Cases cover causal +and noncausal attention, MHA/GQA/MQA, 1/63/64/65/129 keys, unequal rotary axes, +nonzero query offsets, chunked and single-token decoding, amplified scores, +empty batches, reset rejection, unchanged cache storage and allocation guards. +An exact zero-query fixture crosses a split boundary. A poisoned future cache +tests that causal masking excludes both partial tiles and entire future splits. + +## Evidence and limitations + +The host computes FP32 angles and trigonometry, then narrows coefficients to +BF16. GPU validation checks rotation against those stored coefficients. Host +trigonometry has not been compared bit-for-bit with CUDA/PyTorch table generation; +checkpoint/reference execution is still required for model correctness. Local +execution targets RTX 5090 / SM120; no SM100 execution or performance claim is made. + +Validated on RTX 5090 with CUTLASS 4.7 on 2026-10-15: 21 example host tests, +582,652 rotary/cache reference values and 2,464,256 attention reference output +comparisons. The largest attention absolute difference was `0.0078125`, within +the declared tolerance; full prefill and incremental outputs also passed their +comparison. Memcheck and synccheck reported zero errors; racecheck reported +zero hazards, errors or warnings. +The pending-input case exercises ordering but is not an isolated proof of the +context wait: allocation may also serialize prior input writes. diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention.rs new file mode 100644 index 0000000000..2a7b7aa9e9 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention.rs @@ -0,0 +1,247 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Checked synchronous attention over borrowed, completed BF16 cache prefixes. +use crate::attention_kernel::{PARTIAL_SHARED_BYTES, STATE_SIZE, kernels}; +use crate::operator::{GUARD_LEN, Result}; +use cuda_core::{CudaStream, DeviceBuffer, LaunchConfig1D}; +use cute_rs::BFloat16; +use cute_rs::attention::{AttentionShape, AttentionWindow}; +use std::sync::Arc; + +/// Host-checked launch sizes for the fixed 128-wide split-KV schedule. +#[derive(Debug)] +struct Geometry { + /// One CTA for each `(query token, query head, 64-key split)`. + partial_blocks: u32, + /// One CTA for each `(query token, query head)` output row. + merge_blocks: u32, + /// Number of 64-key partitions covering the visible cache allocation prefix. + splits: u32, + /// FP32 scratch element count, excluding guards: `partial_blocks * STATE_SIZE`. + partial_elements: usize, + /// BF16 result element count, excluding guards: `query_tokens * Q * 128`. + output_elements: usize, +} + +/// Validate head grouping, ABI widths, grid limits, and guarded scratch bytes. +/// +/// Returns `None` for empty queries after checking shape validity. Counts describe +/// 64-key splits over the supplied window; all buffer extents are element counts. +fn geometry(window: AttentionWindow) -> Result> { + let shape = AttentionShape::::new()?; + let output_elements = shape.query_len(window.query_tokens())?; + shape.kv_len(window.key_tokens())?; + for value in [Q, K, window.query_tokens(), window.key_tokens()] { + u32::try_from(value).map_err(|_| "attention dimension exceeds kernel ABI")?; + } + if window.query_tokens() == 0 { + return Ok(None); + } + let splits = window.key_tokens().div_ceil(64); + let merge_blocks = window + .query_tokens() + .checked_mul(Q) + .ok_or("query grid overflow")?; + let partial_blocks = merge_blocks + .checked_mul(splits) + .ok_or("split grid overflow")?; + if partial_blocks > i32::MAX as usize || merge_blocks > i32::MAX as usize { + return Err("attention grid exceeds CUDA grid-x range".into()); + } + let partial_elements = partial_blocks + .checked_mul(STATE_SIZE) + .ok_or("scratch overflow")?; + let guarded = partial_elements + .checked_add(GUARD_LEN) + .ok_or("scratch guard overflow")?; + let bytes = guarded + .checked_mul(size_of::()) + .ok_or("scratch byte overflow")?; + if bytes > isize::MAX as usize { + return Err("attention scratch exceeds addressable byte extent".into()); + } + Ok(Some(Geometry { + partial_blocks: partial_blocks as u32, + merge_blocks: merge_blocks as u32, + splits: splits as u32, + partial_elements, + output_elements, + })) +} + +/// Owned BF16 output; query-head count is retained and head width is always 128. +pub struct AttentionOutput { + /// Dense `[tokens, Q, 128]` completed BF16 results, followed by validation guards. + buffer: DeviceBuffer, + /// Number of query token rows, excluding allocation guards. + tokens: usize, +} + +impl AttentionOutput { + /// Read-only allocation including trailing validation guards. + pub fn buffer(&self) -> &DeviceBuffer { + &self.buffer + } + + /// Return the number of completed query token rows. + pub fn tokens(&self) -> usize { + self.tokens + } +} + +/// Bind the split and merge kernels to one stream and context. +pub(crate) struct AttentionRunner { + /// Partial and merge kernels compiled with their fixed CTA launch contracts. + module: kernels::LoadedModule, + /// Stream ordering partial and merge work in the owning context. + stream: Arc, +} + +impl AttentionRunner { + /// Load the split-attention bundle in the supplied stream's CUDA context. + pub fn new(stream: Arc) -> Result { + // SAFETY: this binary's named embedded bundle is compiled from this + // non-generic cuda_module, including its declared launch contracts. + Ok(Self { + module: unsafe { kernels::load(stream.context())? }, + stream, + }) + } + + /// Run FP32 split attention and return dense `[query_tokens, Q, 128]` BF16 output. + /// + /// The cache owner supplies initialized K/V prefixes in `[key_tokens, K, 128]` + /// order. `query_start` is a physical causal token offset; `None` exposes all + /// keys. Checked buffers share this runner's context. A context wait orders + /// prior writes, and an unconditional stream wait retains all borrows through + /// completion even after enqueue errors. Only the final result rounds to BF16; + /// empty queries allocate guarded output without launching attention kernels. + #[allow(clippy::too_many_arguments)] + pub fn run( + &self, + query: &DeviceBuffer, + query_tokens: usize, + key: &DeviceBuffer, + value: &DeviceBuffer, + key_tokens: usize, + query_start: Option, + ) -> Result> { + let window = match query_start { + Some(start) => AttentionWindow::causal(query_tokens, key_tokens, start)?, + None => AttentionWindow::noncausal(query_tokens, key_tokens)?, + }; + let shape = AttentionShape::::new()?; + let query_len = shape.query_len(query_tokens)?; + let kv_len = shape.kv_len(key_tokens)?; + if query.len() < query_len || key.len() < kv_len || value.len() < kv_len { + return Err("attention allocations do not contain the requested prefix".into()); + } + for buffer in [query, key, value] { + if !Arc::ptr_eq(buffer.context(), self.stream.context()) { + return Err("attention inputs must share the runner context".into()); + } + } + let geometry = geometry::(window)?; + self.stream.context().synchronize()?; + let mut output = AttentionOutput { + buffer: crate::operator::allocate(&self.stream, query_len)?, + tokens: query_tokens, + }; + let Some(g) = geometry else { + return Ok(output); + }; + debug_assert_eq!(g.output_elements, query_len); + // NaN sentinels reveal incomplete scratch writes; keep a trailing guard. + let mut partials = DeviceBuffer::from_host( + &self.stream, + &vec![f32::from_bits(0x7fc1_0001); g.partial_elements + GUARD_LEN], + )?; + let partial_grid = self.module.prepare_partial(LaunchConfig1D::new( + g.partial_blocks, + 128, + PARTIAL_SHARED_BYTES, + ))?; + let merge_grid = self + .module + .prepare_merge(LaunchConfig1D::new(g.merge_blocks, 32, 0))?; + // Checked lengths, dimensions, contexts, grids and scratch + // extents. The cache owner guarantees the prefix is initialized. Each + // CTA owns one split or query/head output. Every allocation/borrow stays + // live until the unconditional stream wait, including enqueue errors. + let enqueued = (|| -> Result<()> { + self.module.partial( + &self.stream, + &partial_grid, + query, + key, + value, + &mut partials, + query_tokens as u32, + key_tokens as u32, + Q as u32, + K as u32, + query_start.unwrap_or(0) as u32, + query_start.is_some(), + )?; + self.module.merge( + &self.stream, + &merge_grid, + &partials, + &mut output.buffer, + query_tokens as u32, + Q as u32, + g.splits, + )?; + Ok(()) + })(); + let completed = self.stream.synchronize(); + enqueued?; + completed?; + let mut scratch = vec![0.0; partials.len()]; + partials.copy_to_host(&self.stream, &mut scratch)?; + if !scratch[g.partial_elements..] + .iter() + .all(|x| x.to_bits() == 0x7fc1_0001) + { + return Err("attention scratch guard overwritten".into()); + } + Ok(output) + } +} + +/// CPU checks of split scratch sizing, empty queries, and launch integer limits. +#[cfg(test)] +mod tests { + use super::*; + + /// Check each query/head/split gets one independent 130-element FP32 state. + #[test] + fn splits_and_query_heads_determine_distinct_scratch_rows() { + let window = AttentionWindow::causal(3, 129, 126).unwrap(); + let g = geometry::<16, 8>(window).unwrap().unwrap(); + assert_eq!((g.partial_blocks, g.merge_blocks, g.splits), (144, 48, 3)); + assert_eq!(g.partial_elements, 144 * 130); + assert_eq!(g.output_elements, 3 * 16 * 128); + } + + /// Require valid query/KV grouping even when no attention work is launched. + #[test] + fn empty_queries_launch_nothing_but_invalid_grouping_is_rejected() { + let window = AttentionWindow::causal(0, 0, 0).unwrap(); + assert!(geometry::<16, 8>(window).unwrap().is_none()); + assert!(geometry::<3, 2>(window).is_err()); + } + + /// Reject dimensions and split grids that cannot be represented by the kernel ABI. + #[test] + fn reject_abi_and_grid_overflow_before_allocation() { + let too_long = AttentionWindow::noncausal(1, u32::MAX as usize + 1).unwrap(); + assert!(geometry::<16, 8>(too_long).is_err()); + let too_many_queries = AttentionWindow::noncausal(i32::MAX as usize, 1).unwrap(); + assert!(geometry::<16, 8>(too_many_queries).is_err()); + let split_grid = AttentionWindow::noncausal(2, u32::MAX as usize).unwrap(); + assert!(geometry::<16, 8>(split_grid).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_kernel.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_kernel.rs new file mode 100644 index 0000000000..c9925a8734 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_kernel.rs @@ -0,0 +1,402 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! BF16 grouped-query attention using parallel, stable 64-key reductions. +use core::marker::PhantomData; +use ctm_rs::{ + cuda_module, kernel, launch_bounds, launch_contract, + memory::{DisjointSlice, DynamicSharedArray}, + thread, +}; +use cute_rs::arch::NamedBarrier; +use cute_rs::markers::{C, ComposedLayout, CopyG2S, Layout, LeadingDim, RowMajor, Swizzle}; +use cute_rs::static_layout::StaticLayout; +use cute_rs::tensor::partition::{StridedTensor, StridedTensorMut}; +use cute_rs::{BFloat16, GmemMatrix, RegisterTile, SmemTile, copy_g2s}; + +/// Shared bytes for two 64-by-128 BF16 tiles and FP32 reduction state. +pub const PARTIAL_SHARED_BYTES: u32 = 35_104; +/// FP32 values per split: maximum, denominator, and 128 numerator coordinates. +pub const STATE_SIZE: usize = 130; + +/// Shared-memory storage layout for one complete BF16 key or value tile. +type KvLayout = ComposedLayout, 0, RowMajor<64, 128>>; +/// Eight rows per 128-thread copy; each thread transports eight BF16 values. +type CopyLayout = ComposedLayout, 0, RowMajor<8, 128>>; +/// Map 128 copy threads to eight rows and sixteen contiguous 8-element vectors per row. +type CopyThreads = Layout<(C<8>, C<16>), (C<16>, C<1>)>; +/// Map each copy thread's eight contiguous BF16 carriers to a 16-byte transaction. +type CopyValues = Layout<(C<1>, C<8>), (C<8>, C<1>)>; +/// Four neighboring threads own consecutive 32-coordinate pieces of a key row. +type DotThreads = Layout<(C<4>, C<32>), (C<32>, C<128>)>; +/// Each lane owns four adjacent numerator coordinates; warps own separate rows. +type NumeratorThreads = Layout<(C<32>, C<4>), (C<4>, C<128>)>; + +/// BF16 storage carried unchanged by the existing typed integer copy primitive. +/// +/// The private wrapper keeps the physical `u16` copy carrier separate from +/// numeric BF16 values. No FP16 reinterpretation or numeric cast occurs. +struct Bf16KvTile( + #[doc = "Shared BF16 bit carriers in the typed 64×128 layout."] SmemTile, +); + +/// Attention scale `1 / sqrt(128)` in base two for the exponential instruction. +const SCALE_LOG2: f32 = 0.088_388_346 * core::f32::consts::LOG2_E; + +/// Expand fixed iterations so per-lane numerator coordinates remain registers. +macro_rules! static_for { + ($index:ident in [$($value:expr),*] $body:block) => { + $({ const $index: usize = $value; $body })* + }; +} + +/// Return the greater score for the finite-score and negative-infinity reduction domain. +#[inline(always)] +fn maximum(a: f32, b: f32) -> f32 { + if a > b { a } else { b } +} + +impl Bf16KvTile { + /// Wrap shared storage as a typed 64×128 BF16 tile without initializing it. + /// + /// # Safety + /// `base` supplies 8192 live, suitably aligned shared-memory `u16` elements. + #[inline(always)] + unsafe fn new(base: *mut u16) -> Self { + Self(SmemTile { + base, + capacity: 8192, + layout: PhantomData, + }) + } + + /// Load `N` neighboring BF16 coordinates and widen them to FP32 registers. + /// + /// # Safety + /// The copy must be complete and CTA-synchronized; row < 64 and the + /// complete `column..column + N` interval must lie within 128 columns. + #[inline(always)] + unsafe fn row(&self, row: usize, column: usize) -> RegisterTile { + unsafe { self.0.load::(row, column) }.map(|bits| BFloat16::from_bits(bits).to_f32()) + } +} + +/// Stage only initialized and causally visible cache rows, zero-filling tails. +/// +/// # Safety +/// All 128 threads participate with their unique thread IDs. `source` contains +/// at least `visible * kv_heads * 128` initialized BF16 values, `head < kv_heads`, +/// and destination storage remains live until the copy has been waited on. +#[inline(always)] +#[allow(clippy::too_many_arguments)] +unsafe fn stage( + source: &[BFloat16], + head: usize, + kv_heads: usize, + visible: usize, + split: usize, + destination: &Bf16KvTile, + tid: u32, +) { + // BFloat16 is repr(transparent) over u16. This cast changes only the copy + // carrier; loads later restore the BF16 type before numerical conversion. + let source = GmemMatrix { + base: unsafe { source.as_ptr().add(head * 128) }.cast::(), + rows: visible, + cols: 128, + leading_dim: LeadingDim::<8> { + elements: kv_heads * 128, + }, + }; + static_for!(GROUP in [0, 1, 2, 3, 4, 5, 6, 7] { + let mut part = SmemTile:: { + base: unsafe { destination.0.base.add(GROUP * 1024) }, + capacity: 1024, + layout: PhantomData, + }; + unsafe { + copy_g2s::, CopyThreads, CopyValues, RowMajor<8, 128>, CopyLayout, u16>( + &source, (split * 8 + GROUP, 0), &mut part, tid, + ); + } + }); +} + +/// Fixed-launch BF16 split-attention kernels; host code proves extents and readiness. +#[cuda_module] +pub mod kernels { + use super::*; + + /// Produce one stable FP32 attention state per query/head and 64-key split. + /// + /// Q is dense `[query_tokens, query_heads, 128]`; K/V use + /// `[key_tokens, kv_heads, 128]`. Adjacent query-head groups share a KV head. + /// Each state stores an unscaled score maximum, scaled-exponential denominator, + /// and 128 weighted-value numerator coordinates. Causal offsets count physical + /// cache tokens; fully masked splits publish a neutral state without reading K/V. + /// + /// The launch contract requires 128 threads and 35,104 shared bytes aligned to + /// 1024 bytes. The host must validate positive integral head grouping, a nonempty + /// key prefix for nonempty queries, every buffer/scratch extent, and readiness. + /// All threads of each active CTA reach copies, waits, and barriers convergently; + /// the type system does not establish that participation or output disjointness. + #[launch_bounds(128)] + #[launch_contract(domain=1, block=(128,1,1), dynamic_shared=35104, + dynamic_shared_alignment=1024, min_compute_capability=(12,0))] + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn partial( + q: &[BFloat16], + k: &[BFloat16], + v: &[BFloat16], + mut partials: DisjointSlice, + query_tokens: u32, + key_tokens: u32, + query_heads: u32, + kv_heads: u32, + query_start: u32, + causal: bool, + ) { + let tid = thread::threadIdx_x() as usize; + let lane = tid % 32; + let warp_id = tid / 32; + let numerator_column = NumeratorThreads::crd2idx(lane); + let warp_numerator_offset = NumeratorThreads::crd2idx(tid); + let splits = (key_tokens as usize).div_ceil(64); + let work = thread::blockIdx_x() as usize; + let split = work % splits; + let query_head = (work / splits) % query_heads as usize; + let query_row = (work / splits) / query_heads as usize; + // This guard is CTA-uniform. The checked host grid supplies exactly + // the query rows, heads and splits described by these dimensions. + if query_row >= query_tokens as usize { + return; + } + let visible = if causal { + (query_start as usize + query_row + 1).min(key_tokens as usize) + } else { + key_tokens as usize + }; + let kv_head = query_head / (query_heads as usize / kv_heads as usize); + + // SAFETY: the host checks head grouping, all tensor/scratch extents, + // initialized KV prefix, context/readiness and the exact launch shape. + // Each CTA owns one 130-value state; each lane owns four numerator + // coordinates. All collectives below execute convergently. + unsafe { + let mut destination = StridedTensorMut::from_disjoint_slice(&mut partials) + .subview_unchecked(work * STATE_SIZE, STATE_SIZE); + if split * 64 >= visible { + // A causal query can have whole future splits. Publish the + // neutral reduction state without touching any cache storage. + if warp_id == 0 { + destination + .reborrow() + .subview_unchecked(2 + numerator_column, 4) + .store(RegisterTile::from_array([0.0; 4])); + if lane == 0 { + destination.store_at(0, f32::NEG_INFINITY); + destination.store_at(1, 0.0); + } + } + return; + } + + let shared = DynamicSharedArray::::get_raw().cast::(); + let keys = Bf16KvTile::new(shared); + let values = Bf16KvTile::new(shared.add(8192)); + let score_base = shared.add(16384).cast::(); + let maximum_base = score_base.add(64); + let sum_base = maximum_base.add(4); + let warp_base = sum_base.add(4); + let mut score_output = StridedTensorMut::from_raw_parts(score_base, 64, 1); + let mut maximum_output = StridedTensorMut::from_raw_parts(maximum_base, 4, 1); + let mut sum_output = StridedTensorMut::from_raw_parts(sum_base, 4, 1); + let mut warp_output = StridedTensorMut::from_raw_parts(warp_base, 512, 1); + stage( + k, + kv_head, + kv_heads as usize, + visible, + split, + &keys, + tid as u32, + ); + stage( + v, + kv_head, + kv_heads as usize, + visible, + split, + &values, + tid as u32, + ); + cute_rs::copy::cooperative::finish_async_copies::<1, 128>(); + + let query_base = (query_row * query_heads as usize + query_head) * 128; + // Four neighboring lanes sum 32 dimensions apiece. Two passes + // cover all 64 keys, including zero-filled tail rows. + static_for!(PASS in [0, 1] { + let origin = DotThreads::crd2idx(tid); + let row = origin / 128 + PASS * 32; + let column = origin % 128; + let query = StridedTensor::from_slice(q) + .subview_unchecked(query_base + column, 32); + let dot = keys.0.row_view::<32>(row, column) + .zip_fold(query, 0.0_f32, |dot, key, query| { + dot + query.to_f32() * BFloat16::from_bits(key).to_f32() + }); + let dot = cute_rs::arch::quad_sum_f32(dot); + if tid.is_multiple_of(4) { + score_output.store_at(row, if split * 64 + row < visible { + dot + } else { + f32::NEG_INFINITY + }); + } + }); + NamedBarrier::<1, 128>::sync(); + + let scores = StridedTensor::from_raw_parts(score_base, 64, 1); + let score = if tid < 64 { + scores.get_unchecked(tid) + } else { + f32::NEG_INFINITY + }; + let local_maximum = cute_rs::arch::warp_reduction_max_f32(score); + if lane == 0 { + maximum_output.store_at(warp_id, local_maximum); + } + NamedBarrier::<1, 128>::sync(); + let maxima = StridedTensor::from_raw_parts(maximum_base, 4, 1); + let maximum = maximum( + maximum(maxima.get_unchecked(0), maxima.get_unchecked(1)), + maximum(maxima.get_unchecked(2), maxima.get_unchecked(3)), + ); + // Nonempty splits have at least one visible score. Explicitly + // zero masked rows before they contribute to normalization/PV. + let probability = if tid < 64 && split * 64 + tid < visible { + cute_rs::arch::exp2_f32((score - maximum) * SCALE_LOG2) + } else { + 0.0 + }; + if tid < 64 { + score_output.store_at(tid, probability); + } + let local_sum = cute_rs::arch::warp_reduction_sum_f32(probability); + if lane == 0 { + sum_output.store_at(warp_id, local_sum); + } + NamedBarrier::<1, 128>::sync(); + let sums = StridedTensor::from_raw_parts(sum_base, 4, 1); + + // A warp integrates 16 keys into four dimensions per lane. + let mut numerator = RegisterTile::from_array([0.0f32; 4]); + static_for!(ROW in [0, 1, 2, 3, 4, 5, 6, 7, + 8, 9, 10, 11, 12, 13, 14, 15] { + let row = warp_id * 16 + ROW; + let probability = scores.get_unchecked(row); + numerator = numerator.zip_map(values.row::<4>(row, numerator_column), |sum, value| sum + probability * value); + }); + warp_output + .reborrow() + .subview_unchecked(warp_numerator_offset, 4) + .store(numerator); + NamedBarrier::<1, 128>::sync(); + let warp_outputs = StridedTensor::from_raw_parts(warp_base, 512, 1); + if warp_id == 0 { + let mut total = RegisterTile::from_array([0.0_f32; 4]); + static_for!(WARP in [0, 1, 2, 3] { + total = total.zip_map(warp_outputs.subview_unchecked(WARP * 128 + numerator_column, 4) + .load_unchecked::<4>(), |a, b| a + b); + }); + destination + .reborrow() + .subview_unchecked(2 + numerator_column, 4) + .store(total); + if lane == 0 { + destination.store_at(0, maximum); + destination.store_at( + 1, + sums.get_unchecked(0) + + sums.get_unchecked(1) + + sums.get_unchecked(2) + + sums.get_unchecked(3), + ); + } + } + } + } + + /// Merge completed split states in FP32, then round the normalized output to BF16. + /// + /// Scratch is `[query_tokens, query_heads, splits, STATE_SIZE]`; output is dense + /// `[query_tokens, query_heads, 128]`. One full 32-thread CTA owns each output + /// row. The caller must provide all initialized states from a completed partial + /// launch, positive `splits`, sufficient distinct output storage, and buffer + /// lifetimes through completion. Empty states have zero denominator and are + /// skipped; each lane stores its own four coordinates. + #[launch_bounds(32)] + #[launch_contract(domain=1, block=(32,1,1), min_compute_capability=(12,0))] + #[kernel] + pub fn merge( + partials: &[f32], + mut out: DisjointSlice, + query_tokens: u32, + query_heads: u32, + splits: u32, + ) { + let work = thread::blockIdx_x() as usize; + let lane = thread::threadIdx_x() as usize; + let numerator_column = NumeratorThreads::crd2idx(lane); + if work >= query_tokens as usize * query_heads as usize { + return; + } + // SAFETY: one complete warp owns each query/head output. Host-checked + // scratch extents include every split; the preceding launch completed + // on this stream, and output coordinates are pairwise disjoint. + unsafe { + let states = StridedTensor::from_slice(partials).subview_unchecked( + work * splits as usize * STATE_SIZE, + splits as usize * STATE_SIZE, + ); + let mut max_score = f32::NEG_INFINITY; + let mut split = 0; + while split < splits as usize { + let state = states.subview_unchecked(split * STATE_SIZE, STATE_SIZE); + if state.get_unchecked(1) != 0.0 { + max_score = maximum(max_score, state.get_unchecked(0)); + } + split += 1; + } + let mut denominator = 0.0f32; + let mut numerator = RegisterTile::from_array([0.0f32; 4]); + split = 0; + while split < splits as usize { + let state = states.subview_unchecked(split * STATE_SIZE, STATE_SIZE); + let sum = state.get_unchecked(1); + // A fully masked split is neutral. Skipping it also prevents + // evaluating (-infinity)-(-infinity) for an empty state. + if sum != 0.0 { + let scale = + cute_rs::arch::exp2_f32((state.get_unchecked(0) - max_score) * SCALE_LOG2); + denominator += sum * scale; + let partial = state + .subview_unchecked(2 + numerator_column, 4) + .load_unchecked::<4>(); + numerator = numerator.zip_map(partial, |sum, value| sum + value * scale); + } + split += 1; + } + let reciprocal = if denominator == 0.0 { + 0.0 + } else { + cute_rs::arch::reciprocal_ftz_f32(denominator) + }; + StridedTensorMut::from_disjoint_slice(&mut out) + .subview_unchecked(work * 128 + numerator_column, 4) + .store(numerator.map(|value| BFloat16::from_f32(value * reciprocal))); + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_reference.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_reference.rs new file mode 100644 index 0000000000..d094952f17 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_reference.rs @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent FP64 grouped-query attention oracle over BF16 storage. +use crate::reference::{round, widen}; +use cute_rs::BFloat16; + +/// Computes scaled dot-product attention in token/head/coordinate order. +/// +/// `query_start` selects causal attention at that absolute cache position; +/// `None` exposes every key to every query. Probabilities remain FP64, and +/// only the final output is rounded through FP32 to BF16. Inputs and output are +/// dense `[tokens, heads, dim]` arrays, with separate query and KV head counts. +/// Panics on inconsistent lengths, invalid grouping, or a causal range beyond +/// the cache. Future keys and values are excluded before evaluating arithmetic. +#[allow(clippy::too_many_arguments)] +pub fn attention( + query: &[BFloat16], + key: &[BFloat16], + value: &[BFloat16], + query_heads: usize, + kv_heads: usize, + dim: usize, + query_tokens: usize, + key_tokens: usize, + query_start: Option, +) -> Vec { + assert!(query_heads > 0, "query heads must be positive"); + assert!(kv_heads > 0, "KV heads must be positive"); + assert!(dim > 0, "head dimension must be positive"); + assert!( + query_heads.is_multiple_of(kv_heads), + "query heads must be a multiple of KV heads" + ); + let query_len = query_tokens + .checked_mul(query_heads) + .and_then(|x| x.checked_mul(dim)) + .expect("query extent overflow"); + let kv_len = key_tokens + .checked_mul(kv_heads) + .and_then(|x| x.checked_mul(dim)) + .expect("KV extent overflow"); + assert_eq!(query.len(), query_len, "query storage length mismatch"); + assert_eq!(key.len(), kv_len, "key storage length mismatch"); + assert_eq!(value.len(), kv_len, "value storage length mismatch"); + if let Some(start) = query_start { + let end = start + .checked_add(query_tokens) + .expect("causal query range overflow"); + assert!(end <= key_tokens, "causal query range exceeds key tokens"); + } + assert!( + query_tokens == 0 || key_tokens > 0, + "nonempty queries require at least one key" + ); + + let mut output = Vec::with_capacity(query_len); + let group_size = query_heads / kv_heads; + let scale = 1.0 / (dim as f64).sqrt(); + for query_token in 0..query_tokens { + // Restrict the iteration domain before touching K or V: masking a + // computed NaN score or multiplying a masked NaN value by zero is unsafe. + let visible_keys = query_start.map_or(key_tokens, |start| start + query_token + 1); + for query_head in 0..query_heads { + let kv_head = query_head / group_size; + let query_base = (query_token * query_heads + query_head) * dim; + let mut weights = Vec::with_capacity(visible_keys); + for key_token in 0..visible_keys { + let key_base = (key_token * kv_heads + kv_head) * dim; + let dot: f64 = (0..dim) + .map(|d| widen(query[query_base + d]) * widen(key[key_base + d])) + .sum(); + weights.push(dot * scale); + } + let maximum = weights.iter().copied().fold(f64::NEG_INFINITY, f64::max); + for weight in &mut weights { + *weight = (*weight - maximum).exp(); + } + let denominator: f64 = weights.iter().sum(); + for coordinate in 0..dim { + let numerator: f64 = weights + .iter() + .enumerate() + .map(|(key_token, weight)| { + let index = (key_token * kv_heads + kv_head) * dim + coordinate; + weight * widen(value[index]) + }) + .sum(); + output.push(round(numerator / denominator)); + } + } + } + output +} + +/// Analytic CPU checks of causal windows, head grouping, scale, and stable softmax. +#[cfg(test)] +mod tests { + use super::*; + + /// Widen oracle output for direct comparison with analytic FP64 fixture values. + fn values(output: Vec) -> Vec { + output.into_iter().map(widen).collect() + } + + /// Check uniform softmax over the causal prefix selected by a physical query offset. + #[test] + fn zero_queries_average_each_visible_prefix_at_a_nonzero_offset() { + let query = [round(0.0); 4]; + let key = [1.0, -2.0, 3.0, -4.0, 5.0, -6.0].map(round); + let value = [2.0, 4.0, 6.0, 8.0, 10.0, 12.0].map(round); + assert_eq!( + values(attention(&query, &key, &value, 1, 1, 2, 2, 3, Some(1))), + [4.0, 6.0, 6.0, 8.0] + ); + } + + /// Expose incorrect grouped-query head mapping with distinct KV values. + #[test] + fn adjacent_query_heads_share_the_correct_kv_head() { + let query = [round(0.0); 8]; + let key = [round(0.0); 8]; + let value = [2.0, 4.0, 10.0, 20.0, 6.0, 8.0, 30.0, 40.0].map(round); + assert_eq!( + values(attention(&query, &key, &value, 4, 2, 2, 1, 2, None)), + [4.0, 6.0, 4.0, 6.0, 20.0, 30.0, 20.0, 30.0] + ); + } + + /// Check that the dot-product scale uses head width rather than head count. + #[test] + fn scores_use_inverse_square_root_head_dimension() { + let query = [1.0, 0.0, 0.0, 0.0].map(round); + let key = [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0].map(round); + let value = [0.0, 0.0, 0.0, 0.0, 1.0, 2.0, -1.0, -2.0].map(round); + let actual = attention(&query, &key, &value, 1, 1, 4, 1, 2, None); + // The second probability is sigmoid(1 / sqrt(4)), about 0.622459. + assert_eq!( + actual.iter().map(|x| x.to_bits()).collect::>(), + [0x3f1f, 0x3f9f, 0xbf1f, 0xbf9f] + ); + } + + /// Check max-subtracted FP64 softmax on large but finite input scores. + #[test] + fn extreme_finite_scores_do_not_overflow_the_softmax() { + let query = [round(1024.0)]; + let key = [1024.0, 0.0, -1024.0].map(round); + let value = [7.0, -2.0, 3.0].map(round); + assert_eq!( + values(attention(&query, &key, &value, 1, 1, 1, 1, 3, None)), + [7.0] + ); + } + + /// Ensure poisoned future K/V elements cannot contaminate a causal reference output. + #[test] + fn masked_future_keys_and_values_are_never_evaluated() { + let query = [round(1.0); 2]; + let key = [0.0, 0.0, f64::NAN, f64::NAN].map(round); + let value = [3.0, -4.0, f64::NAN, f64::NAN].map(round); + assert_eq!( + values(attention(&query, &key, &value, 1, 1, 2, 1, 2, Some(0))), + [3.0, -4.0] + ); + } + + /// Check the empty-query result when no key/value tokens exist. + #[test] + fn empty_queries_allow_an_empty_cache() { + assert!(attention(&[], &[], &[], 4, 2, 128, 0, 0, Some(0)).is_empty()); + } + + /// Require the causal query range to lie within the available cache token prefix. + #[test] + #[should_panic(expected = "causal query range exceeds key tokens")] + fn causal_range_must_fit_the_cache() { + attention( + &[round(0.0)], + &[round(0.0); 2], + &[round(0.0); 2], + 1, + 1, + 1, + 1, + 2, + Some(2), + ); + } + + /// Reject a query-head count that is not an integer multiple of KV heads. + #[test] + #[should_panic(expected = "query heads must be a multiple of KV heads")] + fn fractional_head_grouping_is_rejected() { + attention(&[], &[], &[], 3, 2, 128, 0, 0, None); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_validation.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_validation.rs new file mode 100644 index 0000000000..0dffcabf7b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/attention_validation.rs @@ -0,0 +1,238 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! GPU acceptance for causal/noncausal GQA over an owned, incrementally filled cache. +use crate::operator::{AttentionInputs, KvCache, Result, Runner}; +use crate::rotary::RotarySpec; +use crate::{attention_reference, reference}; +use cuda_core::{CudaStream, DeviceBuffer}; +use cute_rs::BFloat16; + +/// Check exact constant-valued attention with zero queries across the 64-key split boundary. +fn uniform(runner: &Runner, stream: &CudaStream) -> Result { + let tokens = 65; + let q = DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0); tokens * 2 * 128])?; + let k = DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0); tokens * 128])?; + let v = DeviceBuffer::from_host(stream, &vec![BFloat16::from_f32(3.0); tokens * 128])?; + let mut cache = KvCache::<1, 128>::new(stream, tokens + 3)?; + let tables = RotarySpec::<128>::standard(1e6, 128)?.tables(&vec![[0; 3]; tokens])?; + let query = runner.append( + &mut cache, + AttentionInputs::<2, 1, 128>::new(&q, &k, &v, tokens)?, + &tables, + )?; + let output = runner.attend(&query, &cache, Some(0))?; + assert_eq!(output.tokens(), tokens); + let actual = crate::query_output(stream, output.buffer(), tokens * 2 * 128)?; + crate::assert_bits( + "uniform GQA", + &actual, + &vec![BFloat16::from_f32(3.0); actual.len()], + ); + println!("PASS uniform BF16 GQA across a 64-key split boundary"); + Ok(actual.len()) +} + +/// BF16 output comparison with one-ULP-scale relative tolerance near normal values. +fn compare(label: &str, actual: &[BFloat16], expected: &[BFloat16]) -> f64 { + assert_eq!(actual.len(), expected.len(), "{label} length"); + let mut maximum = 0.0_f64; + for (i, (&a, &b)) in actual.iter().zip(expected).enumerate() { + let a = reference::widen(a); + let b = reference::widen(b); + let error = (a - b).abs(); + assert!( + a.is_finite() && b.is_finite() && error <= 2e-5 + 0.008 * b.abs(), + "{label} [{i}]: actual={a}, expected={b}, error={error}" + ); + maximum = maximum.max(error); + } + maximum +} + +/// Compare full attention and incremental append/attention over identical positions. +fn sequence( + runner: &Runner, + stream: &CudaStream, + spec: &RotarySpec<128>, + positions: &[[u32; 3]], + chunks: &[usize], + scale: f64, +) -> Result { + let tokens = positions.len(); + assert_eq!(chunks.iter().sum::(), tokens); + let scaled = |seed, length| { + crate::fixture(length, seed) + .into_iter() + .map(|x| reference::round(reference::widen(x) * scale)) + .collect::>() + }; + let q = scaled(11, tokens * Q * 128); + let k = scaled(23, tokens * K * 128); + let v = crate::fixture(tokens * K * 128, 37); + let tables = spec.tables(positions)?; + let expected_q = reference::rotate(&q, tables.cos(), tables.sin(), Q, 128); + let expected_k = reference::rotate(&k, tables.cos(), tables.sin(), K, 128); + let q_dev = DeviceBuffer::from_host(stream, &q)?; + let k_dev = DeviceBuffer::from_host(stream, &k)?; + let v_dev = DeviceBuffer::from_host(stream, &v)?; + let mut cache = KvCache::::new(stream, tokens + 17)?; + let query = runner.append( + &mut cache, + AttentionInputs::::new(&q_dev, &k_dev, &v_dev, tokens)?, + &tables, + )?; + let before = cache.snapshot(stream)?; + let output = runner.attend(&query, &cache, Some(0))?; + assert_eq!(output.tokens(), tokens); + let actual = crate::query_output(stream, output.buffer(), q.len())?; + let expected = attention_reference::attention( + &expected_q, + &expected_k, + &v, + Q, + K, + 128, + tokens, + tokens, + Some(0), + ); + let mut maximum = compare("causal prefill", &actual, &expected); + let output = runner.attend(&query, &cache, None)?; + let noncausal = crate::query_output(stream, output.buffer(), q.len())?; + let expected_noncausal = attention_reference::attention( + &expected_q, + &expected_k, + &v, + Q, + K, + 128, + tokens, + tokens, + None, + ); + maximum = maximum.max(compare("noncausal", &noncausal, &expected_noncausal)); + + let mut incremental = KvCache::::new(stream, tokens + 17)?; + let mut start = 0; + for &count in chunks { + let query_range = start * Q * 128..(start + count) * Q * 128; + let key_range = start * K * 128..(start + count) * K * 128; + let q_chunk = DeviceBuffer::from_host(stream, &q[query_range.clone()])?; + let k_chunk = DeviceBuffer::from_host(stream, &k[key_range.clone()])?; + let v_chunk = DeviceBuffer::from_host(stream, &v[key_range])?; + let tables = spec.tables(&positions[start..start + count])?; + let query = runner.append( + &mut incremental, + AttentionInputs::::new(&q_chunk, &k_chunk, &v_chunk, count)?, + &tables, + )?; + let output = runner.attend(&query, &incremental, Some(start))?; + let result = crate::query_output(stream, output.buffer(), count * Q * 128)?; + maximum = maximum.max(compare( + "incremental oracle", + &result, + &expected[query_range.clone()], + )); + compare( + "prefill versus incremental", + &result, + &actual[query_range.clone()], + ); + // The same short query against the complete cache tests future masking + // and query/key length differences independently of the incremental cache. + let output = runner.attend(&query, &cache, Some(start))?; + let result = crate::query_output(stream, output.buffer(), count * Q * 128)?; + maximum = maximum.max(compare( + "cached offset", + &result, + &expected[query_range.clone()], + )); + let output = runner.attend(&query, &cache, None)?; + let result = crate::query_output(stream, output.buffer(), count * Q * 128)?; + maximum = maximum.max(compare( + "short noncausal query", + &result, + &expected_noncausal[query_range], + )); + start += count; + } + let after = cache.snapshot(stream)?; + crate::assert_bits("attention preserves keys", &after.0, &before.0); + crate::assert_bits("attention preserves values", &after.1, &before.1); + assert_eq!(cache.len(), tokens); + assert!(runner.attend(&query, &cache, Some(tokens + 1)).is_err()); + if tokens > 0 { + cache.reset(); + assert!(runner.attend(&query, &cache, None).is_err()); + assert_eq!(cache.len(), 0); + } + println!( + "PASS BF16 attention Q={Q} KV={K} D=128 tokens={tokens} scale={scale}: max_abs_error={maximum:.8e}" + ); + Ok(q.len() * 5) +} + +/// Entire future splits and a partial future tile contain poisoned K/V values. +fn poisoned_future(runner: &Runner, stream: &CudaStream) -> Result { + let zero = vec![BFloat16::from_bits(0); 128]; + let key = DeviceBuffer::from_host(stream, &zero)?; + let value = DeviceBuffer::from_host(stream, &vec![BFloat16::from_f32(2.0); 128])?; + let query_buffer = DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0); 256])?; + let spec = RotarySpec::<128>::standard(1e6, 256)?; + let mut cache = KvCache::<1, 128>::new(stream, 145)?; + let query = runner.append( + &mut cache, + AttentionInputs::<2, 1, 128>::new(&query_buffer, &key, &value, 1)?, + &spec.tables(&[[0; 3]])?, + )?; + let future = 128; + let poison = DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0x7fc1); future * 128])?; + let queries = DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0); future * 256])?; + runner.append( + &mut cache, + AttentionInputs::<2, 1, 128>::new(&queries, &poison, &poison, future)?, + &spec.tables(&vec![[0; 3]; future])?, + )?; + let output = runner.attend(&query, &cache, Some(0))?; + let actual = crate::query_output(stream, output.buffer(), 256)?; + crate::assert_bits( + "masked future poison", + &actual, + &vec![BFloat16::from_f32(2.0); 256], + ); + println!("PASS masked future K/V poison and neutral split states"); + Ok(actual.len()) +} + +/// Run GPU attention fixtures for MHA/GQA/MQA, causal masking, tails, and incremental caches. +/// +/// Prints comparison counts after the independent FP64 oracle, guard checks, and +/// full-prefill/incremental comparisons succeed; this is synthetic model-component evidence. +pub fn run(runner: &Runner, stream: &CudaStream) -> Result<()> { + let mut checked = uniform(runner, stream)?; + checked += poisoned_future(runner, stream)?; + let standard = RotarySpec::<128>::standard(1e6, 65536)?; + let multimodal = RotarySpec::<128>::interleaved(1e6, 32768, [24, 20, 20])?; + let positions = |count| (0..count).map(|i| [i * 7; 3]).collect::>(); + checked += sequence::<16, 8>(runner, stream, &standard, &positions(1), &[1], 1.0)?; + checked += sequence::<4, 4>(runner, stream, &standard, &positions(63), &[1, 61, 1], 1.0)?; + checked += sequence::<6, 2>(runner, stream, &standard, &positions(64), &[0, 63, 1], 1.0)?; + let axes: Vec<_> = (0..65).map(|i| [i * 3, i * 5 + 1, i * 11 + 7]).collect(); + checked += sequence::<16, 8>(runner, stream, &multimodal, &axes, &[0, 63, 1, 1], 1.0)?; + checked += sequence::<16, 8>( + runner, + stream, + &standard, + &positions(129), + &[1, 62, 1, 1, 64], + 4.0, + )?; + checked += sequence::<4, 1>(runner, stream, &standard, &positions(17), &[1; 17], 1.0)?; + checked += sequence::<16, 8>(runner, stream, &standard, &[], &[0], 1.0)?; + println!( + "PASS: {checked} BF16 attention output comparisons; prefill, chunks and single-token decode" + ); + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/kernel.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/kernel.rs new file mode 100644 index 0000000000..46e0437f6d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/kernel.rs @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Table-driven BF16 split-half rotation and contiguous cache append. +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; +use cute_rs::tensor::partition::{StridedTensor, StridedTensorMut}; +use cute_rs::{BFloat16, RegisterTile}; + +/// Independent-thread rotary and append operations over checked dense BF16 storage. +#[cuda_module] +pub mod kernels { + use super::*; + + /// Rotate dense `[tokens, heads, dim]` BF16 tensors using split-half pairs. + /// + /// Coefficients have shape `[tokens, dim / 2]` and are shared across heads. + /// Each product rounds to BF16 before the sum rounds again. One linear thread + /// owns each pair; `dim` must be positive and even, and `heads` positive. The + /// host must validate nonoverflowing extents, ready inputs, distinct output, + /// and allocation lifetimes. No warp or block collective is required. + #[kernel] + #[allow(clippy::too_many_arguments)] + pub fn rotate_bf16( + input: &[BFloat16], + cos: &[BFloat16], + sin: &[BFloat16], + mut output: DisjointSlice, + tokens: usize, + heads: usize, + dim: usize, + ) { + let pair = thread::index_1d().get(); + let half = dim / 2; + if pair < tokens * heads * half { + let row = pair / half; + let frequency = pair % half; + let table = (row / heads) * half + frequency; + let left = row * dim + frequency; + // SAFETY: the host checks even head width and dense extents; + // each pair contains coordinates separated by one half-head. + let values = unsafe { + StridedTensor::from_raw_parts(input.as_ptr().add(left), 2, half) + .load::<2>() + .map(BFloat16::to_f32) + }; + let c = StridedTensor::from_slice(cos).get(table).to_f32(); + let s = StridedTensor::from_slice(sin).get(table).to_f32(); + // Eager BF16 RoPE rounds each multiplication before addition. + // Keeping these explicit conversions also prevents FMA contraction + // across the observable intermediate rounding boundaries. + let cosine = values.map(|value| BFloat16::from_f32(value * c)); + let sine = values + .gather([1, 0]) + .zip_map(RegisterTile::from_array([true, false]), |value, negate| { + BFloat16::from_f32((if negate { -value } else { value }) * s) + }); + let rotated = cosine.zip_map(sine, |a, b| BFloat16::from_f32(a.to_f32() + b.to_f32())); + // SAFETY: checked dense extents, even head width and one thread + // per pair give two disjoint, in-bounds coordinates per thread. + unsafe { + StridedTensorMut::from_raw_parts(output.as_mut_ptr().add(left), 2, half) + .store(rotated); + } + } + } + + /// Copy rotated keys and unrotated values unchanged into matching cache intervals. + /// + /// `start` and `count` are BF16 element offsets/counts, not token counts. The + /// caller supplies `count` initialized source elements and writable cache + /// intervals `start..start + count` beyond the completed prefix. Each linear + /// thread owns one element; inputs and destinations must not alias and must + /// remain live through completion. No arithmetic conversion or collective occurs. + #[kernel] + pub fn append_bf16( + key: &[BFloat16], + value: &[BFloat16], + mut cache_key: DisjointSlice, + mut cache_value: DisjointSlice, + start: usize, + count: usize, + ) { + let index = thread::index_1d().get(); + if index < count { + // SAFETY: the host checked start+count against the cache capacity. + // Each thread owns one new slot beyond the completed prefix. + unsafe { + let key = StridedTensor::from_slice(key).subview(index, 1); + let value = StridedTensor::from_slice(value).subview(index, 1); + let mut key_destination = StridedTensorMut::from_disjoint_slice(&mut cache_key) + .subview_unchecked(start + index, 1); + let mut value_destination = StridedTensorMut::from_disjoint_slice(&mut cache_value) + .subview_unchecked(start + index, 1); + key.map_to(&mut key_destination, |value| value); + value.map_to(&mut value_destination, |value| value); + } + } + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/main.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/main.rs new file mode 100644 index 0000000000..680cd1e1cb --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/main.rs @@ -0,0 +1,295 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! BF16 rotary positions, owned KV cache and grouped-query attention validation. +/// Checked split-attention launches over a completed cache prefix. +mod attention; +/// CuTe shared-memory staging, 64-key partial attention, and split merging. +mod attention_kernel; +/// Independent FP64 grouped-query attention oracle. +mod attention_reference; +/// GPU fixtures for causal windows, split tails, and incremental attention. +mod attention_validation; +/// Register-tile rotary pairs and bit-preserving cache append kernels. +mod kernel; +/// Owned cache, reusable rotary tables, and synchronous execution boundary. +mod operator; +/// Independent integer BF16 rounding and split-half rotation oracle. +mod reference; +/// Checked host generation and storage of standard or interleaved rotary tables. +mod rotary; + +use cuda_core::{CudaContext, CudaStream, DeviceBuffer}; +use cute_rs::BFloat16; +use operator::{AttentionInputs, GUARD, KvCache, Result, Runner}; +use rotary::{HostTables, RotarySpec}; + +/// Deliberately asymmetric finite inputs distinguish heads, coordinates and roles. +fn fixture(len: usize, seed: usize) -> Vec { + (0..len) + .map(|i| reference::round(((i * 17 + (i / 7) * 3 + seed * 11) % 67) as f64 / 16.0 - 2.0)) + .collect() +} + +/// Compare complete BF16 encodings, reporting the first unequal element and label. +fn assert_bits(label: &str, actual: &[BFloat16], expected: &[BFloat16]) { + assert_eq!(actual.len(), expected.len(), "{label} length"); + for (i, (a, b)) in actual.iter().zip(expected).enumerate() { + assert_eq!(a.to_bits(), b.to_bits(), "{label} at {i}"); + } +} + +/// Require every unwritten element to retain the BF16 sentinel bit pattern. +fn assert_guard(label: &str, values: &[BFloat16]) { + assert!( + values.iter().all(|v| v.to_bits() == GUARD), + "{label} overwritten" + ); +} + +/// Copy a result to the host, validate trailing guards, and return logical elements. +fn query_output( + stream: &CudaStream, + output: &DeviceBuffer, + logical: usize, +) -> Result> { + let mut values = vec![BFloat16::from_bits(0); output.len()]; + output.copy_to_host(stream, &mut values)?; + assert_guard("query guard", &values[logical..]); + values.truncate(logical); + Ok(values) +} + +/// Compare whole and chunked rotary/cache appends against exact BF16 reference bits. +/// +/// `positions` supplies one axis triple per token and `chunks` must sum to that +/// count. Checks completed prefixes, unwritten slots, capacity rejection, and reset; +/// returns the number of logical Q/K/V reference values in the full sequence. +fn validate( + runner: &Runner, + stream: &CudaStream, + spec: &RotarySpec, + positions: &[[u32; 3]], + chunks: &[usize], +) -> Result { + let tokens = positions.len(); + assert_eq!(chunks.iter().sum::(), tokens); + let q = fixture(tokens * Q * D, 1); + let k = fixture(tokens * K * D, 2); + let v = fixture(tokens * K * D, 3); + let tables = spec.tables(positions)?; + let expected_q = reference::rotate(&q, tables.cos(), tables.sin(), Q, D); + let expected_k = reference::rotate(&k, tables.cos(), tables.sin(), K, D); + let q_dev = DeviceBuffer::from_host(stream, &q)?; + let k_dev = DeviceBuffer::from_host(stream, &k)?; + let v_dev = DeviceBuffer::from_host(stream, &v)?; + let mut whole = KvCache::::new(stream, tokens)?; + assert_eq!(whole.capacity(), tokens); + let output = runner.append( + &mut whole, + AttentionInputs::::new(&q_dev, &k_dev, &v_dev, tokens)?, + &tables, + )?; + assert_eq!(output.tokens(), tokens); + assert_bits( + "one-shot query", + &query_output(stream, output.buffer(), q.len())?, + &expected_q, + ); + let whole_snapshot = whole.snapshot(stream)?; + assert_bits("one-shot key", &whole_snapshot.0[..k.len()], &expected_k); + assert_bits("unrotated value", &whole_snapshot.1[..v.len()], &v); + assert_guard("key allocation guard", &whole_snapshot.0[k.len()..]); + assert_guard("value allocation guard", &whole_snapshot.1[v.len()..]); + + let mut cache = KvCache::::new(stream, tokens)?; + let mut start = 0; + for &chunk in chunks { + let end = start + chunk; + let q_dev = DeviceBuffer::from_host(stream, &q[start * Q * D..end * Q * D])?; + let k_dev = DeviceBuffer::from_host(stream, &k[start * K * D..end * K * D])?; + let v_dev = DeviceBuffer::from_host(stream, &v[start * K * D..end * K * D])?; + let tables = spec.tables(&positions[start..end])?; + let output = runner.append( + &mut cache, + AttentionInputs::::new(&q_dev, &k_dev, &v_dev, chunk)?, + &tables, + )?; + assert_eq!(cache.len(), end); + assert_bits( + "chunk query", + &query_output(stream, output.buffer(), chunk * Q * D)?, + &expected_q[start * Q * D..end * Q * D], + ); + let (key, value) = cache.snapshot(stream)?; + assert_bits( + "visible keys", + &key[..end * K * D], + &expected_k[..end * K * D], + ); + assert_bits("visible values", &value[..end * K * D], &v[..end * K * D]); + assert_guard("unwritten key slots", &key[end * K * D..]); + assert_guard("unwritten value slots", &value[end * K * D..]); + start = end; + } + let snapshot = cache.snapshot(stream)?; + assert_bits("chunked/whole keys", &snapshot.0, &whole_snapshot.0); + assert_bits("chunked/whole values", &snapshot.1, &whole_snapshot.1); + + // A full cache must reject another token without changing its cursor/data. + let q1 = DeviceBuffer::from_host(stream, &fixture(Q * D, 4))?; + let k1 = DeviceBuffer::from_host(stream, &fixture(K * D, 5))?; + let v1 = DeviceBuffer::from_host(stream, &fixture(K * D, 6))?; + let table1 = spec.tables(&[[0; 3]])?; + assert!( + runner + .append( + &mut cache, + AttentionInputs::::new(&q1, &k1, &v1, 1)?, + &table1 + ) + .is_err() + ); + assert_eq!(cache.len(), tokens); + let after_error = cache.snapshot(stream)?; + assert_bits("rejected append key", &after_error.0, &snapshot.0); + assert_bits("rejected append value", &after_error.1, &snapshot.1); + let empty = DeviceBuffer::from_host(stream, &[])?; + assert!(AttentionInputs::::new(&q1, &k1, &empty, 1).is_err()); + assert!( + runner + .append( + &mut cache, + AttentionInputs::::new(&empty, &empty, &empty, 0)?, + &table1 + ) + .is_err() + ); + assert_eq!(cache.len(), tokens); + + cache.reset(); + assert_eq!(cache.len(), 0); + if tokens > 0 { + let output = runner.append( + &mut cache, + AttentionInputs::::new(&q1, &k1, &v1, 1)?, + &table1, + )?; + assert_bits( + "position-zero identity", + &query_output(stream, output.buffer(), Q * D)?, + &fixture(Q * D, 4), + ); + assert_eq!(cache.len(), 1); + let reset = cache.snapshot(stream)?; + assert_bits("reset key", &reset.0[..K * D], &fixture(K * D, 5)); + assert_bits("reset value", &reset.1[..K * D], &fixture(K * D, 6)); + } + println!( + "PASS RoPE/cache Q={Q} KV={K} D={D} tokens={tokens}: one-shot, chunks, guards and rejection" + ); + Ok(tokens * (Q + 2 * K) * D) +} + +/// Require the exact rotary result that distinguishes product rounding from one final cast. +fn check_product_rounding(runner: &Runner, stream: &CudaStream) -> Result<()> { + let q = DeviceBuffer::from_host(stream, &[0x3f81, 0x3f02].map(BFloat16::from_bits))?; + let table = HostTables::<2>::from_coefficients( + 1, + vec![BFloat16::from_bits(0x3f35)], + vec![BFloat16::from_bits(0x3f35)], + )?; + let mut cache = KvCache::<1, 2>::new(stream, 1)?; + let result = runner.append( + &mut cache, + AttentionInputs::<1, 1, 2>::new(&q, &q, &q, 1)?, + &table, + )?; + assert_bits( + "separate product rounding", + &query_output(stream, result.buffer(), 2)?, + &[0x3eb4, 0x3f89].map(BFloat16::from_bits), + ); + println!("PASS exact product-rounding discriminator (cast-once would produce 0x3eb5)"); + Ok(()) +} + +/// Prior safe asynchronous writes must finish before the safe runner reads input. +fn check_pending_input(runner: &Runner, stream: &CudaStream) -> Result<()> { + use std::sync::mpsc; + use std::time::Duration; + let producer = stream.context().new_stream()?; + let mut input = DeviceBuffer::from_host(stream, &[BFloat16::from_bits(0x3f80); 2])?; + let mut cache = KvCache::<1, 2>::new(stream, 1)?; + let table = RotarySpec::<2>::standard(1e6, 8)?.tables(&[[0; 3]])?; + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + producer.launch_host_function(move || { + let _ = entered_tx.send(()); + // Driver callbacks must not call CUDA APIs. Bound failure-path waits. + let _ = release_rx.recv_timeout(Duration::from_secs(10)); + })?; + entered_rx.recv_timeout(Duration::from_secs(10))?; + input.zero_async(&producer)?; + assert!(!producer.query()?, "producer must still be pending"); + let release = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(250)); + let _ = release_tx.send(()); + }); + let result = runner.append( + &mut cache, + AttentionInputs::<1, 1, 2>::new(&input, &input, &input, 1)?, + &table, + ); + let completion = producer.synchronize(); + release.join().expect("producer release thread panicked"); + let output = result?; + completion?; + let zero = [BFloat16::from_bits(0); 2]; + assert_bits( + "pending query", + &query_output(stream, output.buffer(), 2)?, + &zero, + ); + let snapshot = cache.snapshot(stream)?; + assert_bits("pending key", &snapshot.0[..2], &zero); + assert_bits("pending value", &snapshot.1[..2], &zero); + assert_eq!(cache.len(), 1); + println!("PASS pending input writes on a separate stream"); + Ok(()) +} + +/// Validate rotary modes, prepared tables, cache lifecycle, and attention on CUDA device 0. +fn main() -> Result<()> { + let context = CudaContext::new(0)?; + let stream = context.default_stream(); + let runner = Runner::new(stream.clone())?; + check_product_rounding(&runner, &stream)?; + check_pending_input(&runner, &stream)?; + operator::validate_prepared_tables()?; + let standard = RotarySpec::<128>::standard(1e6, 65536)?; + let mut positions: Vec<_> = (0..71).map(|i| [i * 13; 3]).collect(); + positions[70] = [65535; 3]; + let mut checked = validate::<16, 8, 128>( + &runner, + &stream, + &standard, + &positions, + &[0, 1, 3, 0, 61, 6], + )?; + let multimodal = RotarySpec::<128>::interleaved(1e6, 32768, [24, 20, 20])?; + let mut positions: Vec<_> = (0..71).map(|i| [i * 11, i * 17 + 3, i * 19 + 7]).collect(); + positions[70] = [32767, 32766, 32765]; + checked += + validate::<16, 8, 128>(&runner, &stream, &multimodal, &positions, &[3, 0, 1, 61, 6])?; + let small = RotarySpec::<6>::standard(1e6, 32)?; + let positions: Vec<_> = (0..17).map(|i| [i; 3]).collect(); + checked += validate::<6, 2, 6>(&runner, &stream, &small, &positions, &[0, 1, 3, 13])?; + checked += validate::<6, 2, 6>(&runner, &stream, &small, &[], &[0])?; + println!( + "PASS: {checked} BF16 query/key/value reference values; stage 4a synthetic validation" + ); + attention_validation::run(&runner, &stream)?; + Ok(()) +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/operator.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/operator.rs new file mode 100644 index 0000000000..7a1bed95ef --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/operator.rs @@ -0,0 +1,621 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Synchronous BF16 rotary composition with privately owned key/value caches. +use crate::{kernel::kernels, rotary::HostTables}; +use cuda_core::LaunchConfig; +use cuda_core::{CudaStream, DeviceBuffer}; +use cute_rs::{ + BFloat16, + attention::{AttentionShape, CacheShape}, +}; +use std::sync::Arc; + +/// Host result carrying CUDA errors and rejected shape, position, or context contracts. +pub type Result = std::result::Result>; +/// BF16 quiet-NaN bits used as unwritten-storage and trailing-guard sentinels. +pub const GUARD: u16 = 0x7fc1; +/// Number of BF16 sentinel elements after each logical allocation. +pub const GUARD_LEN: usize = 8; + +/// Immutable, completed rotary coefficients that may be reused across layers. +/// +/// Construction preserves the checked host table's head width and token count. +/// The allocations stay private: callers cannot mutate them or enqueue writes +/// between appends, and each synchronous append retains this borrow to completion. +pub struct DeviceTables { + /// Number of coefficient rows; each row has `D / 2` elements. + tokens: usize, + /// Completed, private `[tokens, D / 2]` BF16 cosine allocation. + cos: DeviceBuffer, + /// Completed, private `[tokens, D / 2]` BF16 sine allocation. + sin: DeviceBuffer, +} + +impl DeviceTables { + /// Upload both checked `[tokens, D / 2]` tables and complete the host copies. + /// + /// The resulting immutable allocations can be reused on another stream in the + /// same context; their head width and token count remain tied to the host tables. + pub fn new(stream: &CudaStream, tables: &HostTables) -> Result { + stream.context().bind_to_thread()?; + Ok(Self { + tokens: tables.tokens(), + cos: DeviceBuffer::from_host(stream, tables.cos())?, + sin: DeviceBuffer::from_host(stream, tables.sin())?, + }) + } + + /// Return the token-row count carried by these prepared coefficient tables. + pub fn tokens(&self) -> usize { + self.tokens + } +} + +/// Check the BF16 allocation extent including its trailing sentinel elements. +fn guarded_len(logical_len: usize) -> Result { + let elements = logical_len + .checked_add(GUARD_LEN) + .ok_or("guard extent overflow")?; + let bytes = elements + .checked_mul(size_of::()) + .ok_or("guard byte extent overflow")?; + if bytes > isize::MAX as usize { + return Err("guarded BF16 allocation exceeds the slice address range".into()); + } + Ok(elements) +} + +/// One thread handles one rotation pair or one appended BF16 element. +fn launch_grid(work_items: usize) -> Result> { + if work_items == 0 { + return Ok(None); + } + let blocks = work_items.div_ceil(128); + if blocks > i32::MAX as usize { + return Err("work exceeds the CUDA grid-x range".into()); + } + Ok(Some(LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + })) +} + +/// Allocate sentinel-filled storage only after checking the full guarded extent. +pub(crate) fn allocate(stream: &CudaStream, logical_len: usize) -> Result> { + let elements = guarded_len(logical_len)?; + stream.context().bind_to_thread()?; + Ok(DeviceBuffer::from_host( + stream, + &vec![BFloat16::from_bits(GUARD); elements], + )?) +} + +/// Borrow dense BF16 inputs with distinct query/KV head dimensions. +/// +/// Exact lengths reject accidental reinterpretation of trailing storage as +/// tokens. Context and pending producer work are checked by the runner. +pub struct AttentionInputs<'a, const Q: usize, const K: usize, const D: usize> { + /// Unrotated `[tokens, Q, D]` queries; pending writes are ordered by the runner. + query: &'a DeviceBuffer, + /// Unrotated `[tokens, K, D]` keys to rotate before appending. + key: &'a DeviceBuffer, + /// `[tokens, K, D]` values copied unchanged into the cache. + value: &'a DeviceBuffer, + /// Checked grouping and element-count arithmetic for the const head dimensions. + shape: AttentionShape, + /// Number of new token rows across all three inputs. + tokens: usize, +} + +impl<'a, const Q: usize, const K: usize, const D: usize> AttentionInputs<'a, Q, K, D> { + /// Borrow dense `[tokens, Q, D]` queries and `[tokens, K, D]` keys/values. + /// + /// Reject invalid head grouping, extent overflow, and any nonexact buffer length. + /// Context identity and input readiness are established later by the runner. + pub fn new( + query: &'a DeviceBuffer, + key: &'a DeviceBuffer, + value: &'a DeviceBuffer, + tokens: usize, + ) -> Result { + let shape = AttentionShape::new()?; + if query.len() != shape.query_len(tokens)? + || key.len() != shape.kv_len(tokens)? + || value.len() != shape.kv_len(tokens)? + { + return Err( + "attention input lengths must exactly match the token/head dimensions".into(), + ); + } + Ok(Self { + query, + key, + value, + shape, + tokens, + }) + } +} + +/// Own cache storage and publish only the prefix completed by synchronous appends. +/// +/// Buffer ownership and head dimensions cannot be replaced by callers. A reset +/// discards the logical prefix; it does not expose or clear old allocation data. +pub struct KvCache { + /// Rotated keys in `[capacity, K, D]` order, followed by guard elements. + key: DeviceBuffer, + /// Unrotated values with the same dense layout and trailing guards. + value: DeviceBuffer, + /// Checked fixed token capacity and per-token width in BF16 elements. + shape: CacheShape, + /// Completed token prefix; only successful synchronized appends advance it. + valid_tokens: usize, +} + +/// Guarded K/V storage exclusively owned by a private asynchronous execution region. +/// +/// Typed `CacheShape` preserves head count/width and fixed token capacity. The +/// region retains this owner through completion; no conversion back to the public +/// synchronous cache is provided while pending device work may exist. +#[allow( + dead_code, + reason = "Used by model consumers; this standalone example stays synchronous" +)] +pub(crate) struct OwnedKvCache { + /// Initialized rotated keys in `[capacity, K, D]` order plus guard suffix. + pub(crate) key: DeviceBuffer, + /// Initialized unrotated values with the same shape and guard suffix. + pub(crate) value: DeviceBuffer, + /// Checked dimensions retained independently of each buffer's element count. + shape: CacheShape, + /// Last completed prefix, never advanced merely by enqueueing an append. + valid_tokens: usize, +} + +#[allow( + dead_code, + reason = "Used by model consumers; this standalone example stays synchronous" +)] +impl OwnedKvCache { + /// Return the last published completed prefix in token rows. + pub(crate) fn len(&self) -> usize { + self.valid_tokens + } + + /// Return the fixed token capacity, excluding trailing guards. + pub(crate) fn capacity(&self) -> usize { + self.shape.capacity() + } + + /// Advance metadata only after the enclosing region validates every layer. + /// + /// # Safety + /// The caller must prove `self.len() < end <= self.capacity()`, initialized + /// K/V storage through `end`, and successful completion of all device users. + /// No device operation or allocation occurs here, so a stack can validate all + /// layers first and then publish their cursors without a fallible midpoint. + pub(crate) unsafe fn publish_completed(&mut self, end: usize) { + self.valid_tokens = end; + } +} + +impl KvCache { + /// Allocate guarded `[capacity, K, D]` key/value buffers with an empty valid prefix. + pub fn new(stream: &CudaStream, capacity: usize) -> Result { + let shape = CacheShape::new(capacity)?; + Ok(Self { + key: allocate(stream, shape.allocation_len())?, + value: allocate(stream, shape.allocation_len())?, + shape, + valid_tokens: 0, + }) + } + + /// Return the number of tokens published by completed appends. + pub fn len(&self) -> usize { + self.valid_tokens + } + + /// Report whether no completed append has published a token yet. + #[allow( + dead_code, + reason = "Pairs with `len` for model consumers; this standalone example never asks" + )] + pub fn is_empty(&self) -> bool { + self.valid_tokens == 0 + } + + /// Return the fixed token capacity, excluding allocation guards. + pub fn capacity(&self) -> usize { + self.shape.capacity() + } + + /// Validate completed private storage before transfer into an execution region. + /// + /// This performs no synchronization or mutation; public cache operations have + /// already completed every append. The enclosing stack fences context readiness + /// once, after every layer has passed these shape/context checks. + #[allow( + dead_code, + reason = "Used by model consumers; this standalone example stays synchronous" + )] + pub(crate) fn validate_execution(&self, stream: &CudaStream) -> Result<()> { + let allocation_len = guarded_len(self.shape.allocation_len())?; + if self.valid_tokens > self.shape.capacity() + || self.key.len() != allocation_len + || self.value.len() != allocation_len + { + return Err("execution cache requires complete guarded K/V storage".into()); + } + if !Arc::ptr_eq(self.key.context(), stream.context()) + || !Arc::ptr_eq(self.value.context(), stream.context()) + { + return Err("execution cache must share the supplied stream context".into()); + } + Ok(()) + } + + /// Consume completed storage without copying, clearing, or changing its prefix. + /// + /// Public cache operations finish before returning and expose no mutable + /// backing storage, so consumption cannot abandon a pending append. The caller + /// establishes region readiness once for the complete stack before submission. + #[allow( + dead_code, + reason = "Used by model consumers; this standalone example stays synchronous" + )] + pub(crate) fn into_execution(self, stream: &CudaStream) -> Result> { + self.validate_execution(stream)?; + Ok(OwnedKvCache { + key: self.key, + value: self.value, + shape: self.shape, + valid_tokens: self.valid_tokens, + }) + } + + /// Return all capacity and sentinel elements for this example's validation. + /// Only the first `len()` tokens represent the current cache prefix. + pub fn snapshot(&self, stream: &CudaStream) -> Result<(Vec, Vec)> { + if !Arc::ptr_eq(self.key.context(), stream.context()) + || !Arc::ptr_eq(self.value.context(), stream.context()) + { + return Err("snapshot stream must share the cache context".into()); + } + stream.context().synchronize()?; + let mut key = vec![BFloat16::from_bits(0); self.key.len()]; + let mut value = vec![BFloat16::from_bits(0); self.value.len()]; + self.key.copy_to_host(stream, &mut key)?; + self.value.copy_to_host(stream, &mut value)?; + Ok((key, value)) + } + + /// Discard the completed prefix; the next append starts at token zero. + pub fn reset(&mut self) { + self.valid_tokens = 0; + } +} + +/// Owned rotated query storage, retaining its query-head and head-width types. +pub struct RotatedQuery { + /// Completed dense `[tokens, Q, D]` rotated queries, plus trailing guards. + buffer: DeviceBuffer, + /// Number of valid query token rows, excluding guards. + tokens: usize, +} + +impl RotatedQuery { + /// Read-only buffer including the trailing validation sentinel elements. + pub fn buffer(&self) -> &DeviceBuffer { + &self.buffer + } + + /// Return the number of rotated query token rows. + pub fn tokens(&self) -> usize { + self.tokens + } +} + +/// Bind a loaded kernel module to one stream and context. +pub struct Runner { + /// Rotation and append kernels loaded in the execution context. + module: kernels::LoadedModule, + /// Split-attention executor sharing this runner's stream and context. + attention: crate::attention::AttentionRunner, + /// Stream whose completion makes returned queries and cache prefixes readable. + stream: Arc, +} + +impl Runner { + /// Load rotary/cache and attention kernels in the supplied stream's context. + pub fn new(stream: Arc) -> Result { + Ok(Self { + module: kernels::load(stream.context())?, + attention: crate::attention::AttentionRunner::new(stream.clone())?, + stream, + }) + } + + /// Attend over the completed cache prefix with fixed 128-wide BF16 heads. + /// `query_start` is a physical causal cache offset, independent of RoPE + /// positions. `None` selects all available keys for every query. + /// Queries and cache must share this runner's context. Pending work is + /// ordered before reading, and input borrows remain live through completion, + /// including enqueue failures. Empty queries launch no attention kernels. + pub fn attend( + &self, + query: &RotatedQuery, + cache: &KvCache, + query_start: Option, + ) -> Result> { + self.attention.run::( + &query.buffer, + query.tokens, + &cache.key, + &cache.value, + cache.valid_tokens, + query_start, + ) + } + + /// Rotate Q/K and append completed K/V, retaining every borrow through completion. + /// + /// Coefficients are explicit inputs. Their token count and head dimension + /// are checked, while their semantic positions remain the caller's choice. + /// The current synchronous runner orders prior work in its whole context. + pub fn append( + &self, + cache: &mut KvCache, + inputs: AttentionInputs<'_, Q, K, D>, + tables: &HostTables, + ) -> Result> { + let tables = DeviceTables::new(&self.stream, tables)?; + self.append_prepared(cache, inputs, &tables) + } + + /// Reuse completed coefficients without changing the synchronous append + /// contract. Input readiness, cache publication and error completion remain + /// the same as [`Self::append`]; the tables may originate on another stream + /// in this context because their constructor completed both uploads. + pub fn append_prepared( + &self, + cache: &mut KvCache, + inputs: AttentionInputs<'_, Q, K, D>, + tables: &DeviceTables, + ) -> Result> { + if !D.is_multiple_of(2) { + return Err("split-half rotary head width must be even".into()); + } + if tables.tokens() != inputs.tokens { + return Err("rotary table token count does not match attention inputs".into()); + } + for buffer in [ + inputs.query, + inputs.key, + inputs.value, + &cache.key, + &cache.value, + &tables.cos, + &tables.sin, + ] { + if !Arc::ptr_eq(buffer.context(), self.stream.context()) { + return Err( + "attention inputs, cache and tables must share the runner context".into(), + ); + } + } + let range = cache + .shape + .append_range(cache.valid_tokens, inputs.tokens)?; + let query_len = inputs.shape.query_len(inputs.tokens)?; + let kv_len = inputs.shape.kv_len(inputs.tokens)?; + let query_grid = launch_grid(query_len / 2)?; + let key_grid = launch_grid(kv_len / 2)?; + let append_grid = launch_grid(kv_len)?; + let start_element = range + .start + .checked_mul(cache.shape.width()) + .ok_or("cache offset overflow")?; + // Safe asynchronous writes may be pending on another stream. Establish + // readiness before any kernel reads the borrowed input allocations. + self.stream.context().synchronize()?; + let mut query = RotatedQuery { + buffer: allocate(&self.stream, query_len)?, + tokens: inputs.tokens, + }; + if inputs.tokens == 0 { + return Ok(query); + } + let query_grid = query_grid.ok_or("missing query rotation grid")?; + let key_grid = key_grid.ok_or("missing key rotation grid")?; + let append_grid = append_grid.ok_or("missing cache append grid")?; + let mut rotated_key = allocate(&self.stream, kv_len)?; + // SAFETY: lengths, contexts, grids and capacity were checked above. + // Inputs remain borrowed; every output is a distinct owned allocation. + // Rotation threads own disjoint pairs, and append threads own disjoint + // elements beyond the completed cache prefix. One stream orders all + // launches; unconditional synchronization retains their buffers on error. + let enqueued = (|| -> Result<()> { + unsafe { + self.module.rotate_bf16( + &self.stream, + query_grid, + inputs.query, + &tables.cos, + &tables.sin, + &mut query.buffer, + inputs.tokens, + Q, + D, + )?; + self.module.rotate_bf16( + &self.stream, + key_grid, + inputs.key, + &tables.cos, + &tables.sin, + &mut rotated_key, + inputs.tokens, + K, + D, + )?; + self.module.append_bf16( + &self.stream, + append_grid, + &rotated_key, + inputs.value, + &mut cache.key, + &mut cache.value, + start_element, + kv_len, + )?; + } + Ok(()) + })(); + let completed = self.stream.synchronize(); + enqueued?; + completed?; + cache.valid_tokens = range.end; + Ok(query) + } +} + +/// GPU validation lives in the executable rather than the CPU test harness, +/// which does not link the compiled kernel artifact. +pub(crate) fn validate_prepared_tables() -> Result<()> { + prepared_tables_reuse_completed_coefficients_across_streams_exactly()?; + prepared_table_rejection_preserves_the_completed_cache_prefix()?; + println!("PASS prepared rotary table reuse, exact rounding, guards and rejection"); + Ok(()) +} + +/// Check cross-stream table reuse, exact rotary rounding, and cache/query guards. +fn prepared_tables_reuse_completed_coefficients_across_streams_exactly() -> Result<()> { + let context = cuda_core::CudaContext::new(0)?; + let upload = context.new_stream()?; + let stream = context.new_stream()?; + let runner = Runner::new(stream.clone())?; + let host = HostTables::<2>::from_coefficients( + 1, + vec![BFloat16::from_bits(0x3f35)], + vec![BFloat16::from_bits(0x3f35)], + )?; + let tables = DeviceTables::new(&upload, &host)?; + drop(host); + let input = DeviceBuffer::from_host(&stream, &[0x3f81, 0x3f02].map(BFloat16::from_bits))?; + let mut cache = KvCache::<1, 2>::new(&stream, 2)?; + for prefix in 1..=2 { + let query = runner.append_prepared( + &mut cache, + AttentionInputs::<1, 1, 2>::new(&input, &input, &input, 1)?, + &tables, + )?; + let actual = query.buffer().to_host_vec(&stream)?; + // Discriminates separate BF16 product rounding from cast-once rotation. + assert_eq!([actual[0].to_bits(), actual[1].to_bits()], [0x3eb4, 0x3f89]); + assert!(actual[2..].iter().all(|value| value.to_bits() == GUARD)); + assert_eq!(cache.len(), prefix); + } + let (key, value) = cache.snapshot(&stream)?; + let bits = |values: &[BFloat16]| { + values + .iter() + .map(|value| value.to_bits()) + .collect::>() + }; + assert_eq!(bits(&key[..4]), [0x3eb4, 0x3f89, 0x3eb4, 0x3f89]); + assert_eq!(bits(&value[..4]), [0x3f81, 0x3f02, 0x3f81, 0x3f02]); + assert!( + key[4..] + .iter() + .chain(&value[4..]) + .all(|value| value.to_bits() == GUARD) + ); + Ok(()) +} + +/// Check that wrong table token counts or contexts leave the published cache unchanged. +fn prepared_table_rejection_preserves_the_completed_cache_prefix() -> Result<()> { + let context = cuda_core::CudaContext::new(0)?; + let stream = context.default_stream(); + let runner = Runner::new(stream.clone())?; + let input = DeviceBuffer::from_host(&stream, &[BFloat16::from_f32(1.0); 2])?; + let mut cache = KvCache::<1, 2>::new(&stream, 2)?; + let rotary = crate::rotary::RotarySpec::<2>::standard(1e6, 8)?; + let host = rotary.tables(&[[0; 3]])?; + runner.append( + &mut cache, + AttentionInputs::<1, 1, 2>::new(&input, &input, &input, 1)?, + &host, + )?; + let before = cache.snapshot(&stream)?; + let wrong_tokens = DeviceTables::new(&stream, &rotary.tables(&[[0; 3]; 2])?)?; + let other = cuda_core::CudaContext::new(0)?; + let wrong_context = DeviceTables::new(&other.default_stream(), &host)?; + for tables in [&wrong_tokens, &wrong_context] { + assert!( + runner + .append_prepared( + &mut cache, + AttentionInputs::<1, 1, 2>::new(&input, &input, &input, 1)?, + tables, + ) + .is_err() + ); + assert_eq!(cache.len(), 1); + let after = cache.snapshot(&stream)?; + assert!( + before + .0 + .iter() + .chain(&before.1) + .zip(after.0.iter().chain(&after.1)) + .all(|(left, right)| left.to_bits() == right.to_bits()) + ); + } + Ok(()) +} + +/// CPU checks for guarded allocation extents and scalar launch geometry. +#[cfg(test)] +mod tests { + use super::*; + + /// Check eight-element guards against both integer and slice-byte overflow. + #[test] + fn guards_cover_empty_and_nonempty_buffers_without_exceeding_slice_extent() { + assert_eq!(guarded_len(0).unwrap(), 8); + assert_eq!(guarded_len(17).unwrap(), 25); + let largest = isize::MAX as usize / size_of::(); + assert_eq!(guarded_len(largest - GUARD_LEN).unwrap(), largest); + assert!(guarded_len(largest - GUARD_LEN + 1).is_err()); + assert!(guarded_len(usize::MAX).is_err()); + } + + /// Check 128-thread block rounding and the empty-work no-launch policy. + #[test] + fn launch_grid_covers_partial_blocks_and_skips_empty_work() { + assert!(launch_grid(0).unwrap().is_none()); + for (elements, blocks) in [(1, 1), (127, 1), (128, 1), (129, 2), (257, 3)] { + let grid = launch_grid(elements).unwrap().unwrap(); + assert_eq!(grid.grid_dim, (blocks, 1, 1)); + assert_eq!(grid.block_dim, (128, 1, 1)); + assert_eq!(grid.shared_mem_bytes, 0); + } + } + + /// Require a representable CUDA grid before narrowing the block count to `u32`. + #[test] + fn launch_grid_rejects_cuda_grid_overflow_before_casting() { + let largest = i32::MAX as usize * 128; + assert_eq!( + launch_grid(largest).unwrap().unwrap().grid_dim.0, + i32::MAX as u32 + ); + assert!(launch_grid(largest + 1).is_err()); + assert!(launch_grid(usize::MAX).is_err()); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/reference.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/reference.rs new file mode 100644 index 0000000000..915d4fbb29 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/reference.rs @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Independent BF16 storage arithmetic and coordinate-wise rotation oracle. +use cute_rs::BFloat16; + +/// Narrow FP64 through FP32, then round to BF16 with an independent integer even-tie oracle. +/// +/// Signed zero is retained and NaNs are quieted while preserving retained payload bits. +pub fn round(value: f64) -> BFloat16 { + let value = value as f32; + let bits = value.to_bits(); + let retained = (bits >> 16) as u16; + if value.is_nan() { + return BFloat16::from_bits(retained | 0x0040); + } + let discarded = bits & 0xffff; + let increment = discarded > 0x8000 || (discarded == 0x8000 && retained & 1 != 0); + BFloat16::from_bits(retained.wrapping_add(u16::from(increment))) +} + +/// Expand BF16 storage bits exactly to FP64 through an FP32 bit reconstruction. +pub fn widen(value: BFloat16) -> f64 { + f64::from(f32::from_bits(u32::from(value.to_bits()) << 16)) +} + +/// Rotate dense `[tokens, heads, dim]` values using `[tokens, dim / 2]` tables. +/// +/// The fixture must provide positive even `dim`, positive `heads`, and complete +/// matching buffers. Coordinate `i` pairs with `i + dim / 2`; each BF16 product +/// rounds before the final BF16 sum. Table values are shared by heads of a token. +pub fn rotate( + input: &[BFloat16], + cos: &[BFloat16], + sin: &[BFloat16], + heads: usize, + dim: usize, +) -> Vec { + let half = dim / 2; + input + .iter() + .enumerate() + .map(|(index, &x)| { + let coordinate = index % dim; + let token = index / (heads * dim); + let frequency = coordinate % half; + let table = token * half + frequency; + let rotated = if coordinate < half { + -widen(input[index + half]) + } else { + widen(input[index - half]) + }; + let product = round(widen(x) * widen(cos[table])); + let other = round(rotated * widen(sin[table])); + round(widen(product) + widen(other)) + }) + .collect() +} + +/// CPU fixtures checking split-half coordinates and independent BF16 rounding. +#[cfg(test)] +mod tests { + use super::*; + + /// Use a quarter-turn to expose pairing the wrong rotary coordinates. + #[test] + fn quarter_turn_rotates_split_halves_not_adjacent_pairs() { + let x = [1.0, 2.0, 3.0, 4.0].map(round); + let y = rotate(&x, &[round(0.0); 2], &[round(1.0); 2], 1, 4); + assert_eq!( + y.iter().copied().map(widen).collect::>(), + [-3.0, -4.0, 1.0, 2.0] + ); + } + + /// Check both tie directions, signed zero, infinity, and retained NaN status. + #[test] + fn rounding_retains_even_ties_and_special_bits() { + assert_eq!( + round(f64::from(f32::from_bits(0x3f80_8000))).to_bits(), + 0x3f80 + ); + assert_eq!( + round(f64::from(f32::from_bits(0x3f81_8000))).to_bits(), + 0x3f82 + ); + assert_eq!(round(-0.0).to_bits(), 0x8000); + assert_eq!(round(f64::INFINITY).to_bits(), 0x7f80); + assert_ne!(round(f64::NAN).to_bits() & 0x7f, 0); + } + + /// Require the exact two-product rounding result that a single final cast misses. + #[test] + fn each_product_rounds_before_the_sum() { + let x = [0x3f81, 0x3f02].map(BFloat16::from_bits); + let c = [BFloat16::from_bits(0x3f35)]; + let y = rotate(&x, &c, &c, 1, 2); + // Cast-once would produce 0x3eb5 for the first coordinate. + assert_eq!(y[0].to_bits(), 0x3eb4); + assert_eq!(y[1].to_bits(), 0x3f89); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/rotary.rs b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/rotary.rs new file mode 100644 index 0000000000..fbbdbdf446 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/src/rotary.rs @@ -0,0 +1,398 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Checked host coefficients for split-half BF16 rotary embeddings. +use cute_rs::BFloat16; + +/// Rejected rotary geometry, positions, or coefficient storage. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RopeError { + /// Head width is zero or odd, so split-half pairs cannot be formed. + InvalidDimension, + /// The frequency base is nonfinite or nonpositive. + InvalidTheta, + /// The exclusive position limit is zero. + InvalidMaxPosition, + /// Axis counts or interleaved frequency indices do not fit the half-head. + InvalidSections, + /// Element counts overflow or BF16 bytes exceed the addressable slice extent. + ExtentOverflow, + /// Coefficient storage differs from `tokens * (D / 2)` elements. + LengthMismatch, + /// At least one axis position reaches or exceeds the exclusive limit. + PositionOutOfBounds, + /// Standard RoPE received different positions for its three axes. + UnequalAxes, + /// A supplied sine or cosine coefficient contains infinity or NaN. + NonFiniteCoefficient, +} + +impl core::fmt::Display for RopeError { + /// Describe the rejected rotary contract without exposing implementation details. + fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + formatter.write_str(match self { + Self::InvalidDimension => "rotary head width must be positive and even", + Self::InvalidTheta => "rotary theta must be finite and positive", + Self::InvalidMaxPosition => "rotary position limit must be positive", + Self::InvalidSections => "interleaved rotary sections do not fit the half-head", + Self::ExtentOverflow => "rotary table extent is not addressable", + Self::LengthMismatch => { + "rotary coefficient lengths must equal tokens times half-head width" + } + Self::PositionOutOfBounds => "rotary position must be below its configured limit", + Self::UnequalAxes => "standard rotary positions must agree across all three axes", + Self::NonFiniteCoefficient => "rotary coefficients must be finite", + }) + } +} + +impl core::error::Error for RopeError {} + +/// Frequency-to-position mapping for standard RoPE or interleaved three-axis MRoPE. +#[derive(Clone, Copy, Debug)] +enum Axes { + /// Every frequency uses the same position, requiring equal input axes. + Standard, + /// Counts for temporal, height, and width frequencies, summing to `D / 2`. + Interleaved(#[doc = "Half-head frequency counts in temporal, height, width order."] [usize; 3]), +} + +/// Positional geometry, independent of hidden width and the number of heads. +#[derive(Debug)] +pub struct RotarySpec { + /// Positive finite base used by `theta^(-2 * frequency / D)`. + theta: f32, + /// Exclusive upper bound on each integer temporal/height/width position. + max_position: u32, + /// Rule selecting which position axis supplies each half-head frequency. + axes: Axes, +} + +/// One contiguous half-head per token; both rotated halves share these values. +/// +/// The const head width prevents passing a table to a differently sized rotary +/// operation. Construction checks dynamic lengths and finite coefficient values. +#[derive(Debug)] +pub struct HostTables { + /// Number of token rows, each containing `D / 2` coefficient pairs. + tokens: usize, + /// Dense `[tokens, D / 2]` cosine coefficients rounded to BF16. + cos: Vec, + /// Dense `[tokens, D / 2]` sine coefficients rounded to BF16. + sin: Vec, +} + +/// Validate a positive even head width and return its number of rotation pairs. +fn half_width() -> Result { + if D == 0 || !D.is_multiple_of(2) { + return Err(RopeError::InvalidDimension); + } + Ok(D / 2) +} + +/// Return `tokens * (D / 2)` elements after checking BF16 byte addressability. +fn table_length(tokens: usize) -> Result { + let length = tokens + .checked_mul(half_width::()?) + .ok_or(RopeError::ExtentOverflow)?; + let bytes = length + .checked_mul(core::mem::size_of::()) + .ok_or(RopeError::ExtentOverflow)?; + if bytes > isize::MAX as usize { + return Err(RopeError::ExtentOverflow); + } + Ok(length) +} + +impl RotarySpec { + /// Use one position for every frequency; unequal axis inputs are rejected. + pub fn standard(theta: f32, max_position: u32) -> Result { + half_width::()?; + if !theta.is_finite() || theta <= 0.0 { + return Err(RopeError::InvalidTheta); + } + if max_position == 0 { + return Err(RopeError::InvalidMaxPosition); + } + Ok(Self { + theta, + max_position, + axes: Axes::Standard, + }) + } + + /// Select height/width frequencies at indices `1+3j`/`2+3j`, leaving + /// remaining frequencies temporal. Counts are `[temporal, height, width]` + /// and must sum to `D / 2`; every selected index must fit that half-head. + pub fn interleaved( + theta: f32, + max_position: u32, + sections: [usize; 3], + ) -> Result { + let mut spec = Self::standard(theta, max_position)?; + let count = sections.into_iter().try_fold(0usize, |sum, size| { + sum.checked_add(size).ok_or(RopeError::ExtentOverflow) + })?; + if count != D / 2 { + return Err(RopeError::InvalidSections); + } + for (axis, count) in sections.into_iter().enumerate().skip(1) { + if count > 0 { + let last = (count - 1) + .checked_mul(3) + .and_then(|index| index.checked_add(axis)) + .ok_or(RopeError::ExtentOverflow)?; + if last >= D / 2 { + return Err(RopeError::InvalidSections); + } + } + } + spec.axes = Axes::Interleaved(sections); + Ok(spec) + } + + /// Compute FP32 frequencies, angles and trigonometry, then round tables to + /// BF16. Positions are `[temporal, height, width]` per token, independent of + /// physical cache offsets. Both output arrays have shape `[tokens, D / 2]`. + /// Host trigonometry is not claimed bit-identical to CUDA or PyTorch. + pub fn tables(&self, positions: &[[u32; 3]]) -> Result, RopeError> { + let length = table_length::(positions.len())?; + for position in positions { + if position.iter().any(|&value| value >= self.max_position) { + return Err(RopeError::PositionOutOfBounds); + } + if matches!(self.axes, Axes::Standard) + && (position[0] != position[1] || position[0] != position[2]) + { + return Err(RopeError::UnequalAxes); + } + } + let mut cos = Vec::with_capacity(length); + let mut sin = Vec::with_capacity(length); + for position in positions { + for index in 0..D / 2 { + let axis = match self.axes { + Axes::Standard => 0, + Axes::Interleaved(sections) => { + let candidate = index % 3; + if candidate != 0 && index / 3 < sections[candidate] { + candidate + } else { + 0 + } + } + }; + let inverse_frequency = 1.0 / self.theta.powf(2.0 * index as f32 / D as f32); + let angle = position[axis] as f32 * inverse_frequency; + cos.push(BFloat16::from_f32(angle.cos())); + sin.push(BFloat16::from_f32(angle.sin())); + } + } + HostTables::from_coefficients(positions.len(), cos, sin) + } +} + +impl HostTables { + /// Accept independently supplied coefficients without weakening shape or + /// dtype checks. Coefficients need not be generated by [`RotarySpec`]. + pub fn from_coefficients( + tokens: usize, + cos: Vec, + sin: Vec, + ) -> Result { + let length = table_length::(tokens)?; + if cos.len() != length || sin.len() != length { + return Err(RopeError::LengthMismatch); + } + if cos + .iter() + .chain(&sin) + .any(|value| !value.to_f32().is_finite()) + { + return Err(RopeError::NonFiniteCoefficient); + } + Ok(Self { tokens, cos, sin }) + } + + /// Return the number of token rows represented by the coefficient tables. + pub fn tokens(&self) -> usize { + self.tokens + } + + /// Borrow dense `[tokens, D / 2]` BF16 cosine coefficients in frequency order. + pub fn cos(&self) -> &[BFloat16] { + &self.cos + } + + /// Borrow dense `[tokens, D / 2]` BF16 sine coefficients in frequency order. + pub fn sin(&self) -> &[BFloat16] { + &self.sin + } +} + +/// CPU checks of axis selection, geometry, table bounds, and coefficient bits. +#[cfg(test)] +mod tests { + use super::*; + + /// Extract coefficient encodings for exact tests independent of numeric equality. + fn bits(values: &[BFloat16]) -> Vec { + values.iter().map(|value| value.to_bits()).collect() + } + + /// Check literal BF16 tables, exclusive position bounds, equal axes, and empty input. + #[test] + fn standard_positions_keep_tokens_and_split_half_frequency_order() { + let spec = RotarySpec::<4>::standard(16.0, 3).unwrap(); + let tables = spec.tables(&[[0; 3], [1; 3], [2; 3]]).unwrap(); + assert_eq!(tables.tokens(), 3); + // Frequencies [1, 1/4]; literal BF16 encodings of angles [0,0,1,.25,2,.5]. + assert_eq!( + bits(tables.cos()), + [0x3f80, 0x3f80, 0x3f0a, 0x3f78, 0xbed5, 0x3f61] + ); + assert_eq!(bits(tables.sin()), [0, 0, 0x3f57, 0x3e7d, 0x3f69, 0x3ef5]); + assert_eq!( + spec.tables(&[[3; 3]]).unwrap_err(), + RopeError::PositionOutOfBounds + ); + assert_eq!( + spec.tables(&[[1, 2, 1]]).unwrap_err(), + RopeError::UnequalAxes + ); + let empty = spec.tables(&[]).unwrap(); + assert_eq!(empty.tokens(), 0); + assert!(empty.cos().is_empty() && empty.sin().is_empty()); + } + + /// Check the `[24, 20, 20]` MRoPE mask, including its temporal-only final frequencies. + #[test] + fn interleaved_selection_uses_the_literal_qwen_axis_mask() { + let spec = RotarySpec::<128>::interleaved(1.0, 4, [24, 20, 20]).unwrap(); + let tables = spec.tables(&[[0, 1, 2], [2, 0, 1]]).unwrap(); + // Half-head indices 60..64 stay temporal; the mask is not a uniform + // three-way cycle and does not describe adjacent rotation pairs. + /// Expected temporal/height/width axis for each of the 64 half-head frequencies. + const AXIS: [usize; 64] = [ + 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, + 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, + 1, 2, 0, 0, 0, 0, + ]; + for (token, (cos, sin)) in [ + ([0x3f80, 0x3f0a, 0xbed5], [0, 0x3f57, 0x3f69]), + ([0xbed5, 0x3f80, 0x3f0a], [0x3f69, 0, 0x3f57]), + ] + .into_iter() + .enumerate() + { + for (i, axis) in AXIS.into_iter().enumerate() { + assert_eq!(tables.cos()[token * 64 + i].to_bits(), cos[axis]); + assert_eq!(tables.sin()[token * 64 + i].to_bits(), sin[axis]); + } + } + for position in [[4, 0, 0], [0, 4, 0], [0, 0, 4]] { + assert_eq!( + spec.tables(&[position]).unwrap_err(), + RopeError::PositionOutOfBounds + ); + } + } + + /// Reject odd/zero head widths, bad frequency bases, and invalid section placement. + #[test] + fn invalid_geometry_or_interleaved_counts_are_rejected() { + assert_eq!( + RotarySpec::<0>::standard(1.0, 1).unwrap_err(), + RopeError::InvalidDimension + ); + assert_eq!( + RotarySpec::<3>::standard(1.0, 1).unwrap_err(), + RopeError::InvalidDimension + ); + for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert_eq!( + RotarySpec::<4>::standard(theta, 1).unwrap_err(), + RopeError::InvalidTheta + ); + } + assert_eq!( + RotarySpec::<4>::standard(1.0, 0).unwrap_err(), + RopeError::InvalidMaxPosition + ); + assert_eq!( + RotarySpec::<128>::interleaved(1.0, 4, [23, 20, 20]).unwrap_err(), + RopeError::InvalidSections + ); + assert_eq!( + RotarySpec::<8>::interleaved(1.0, 4, [1, 2, 1]).unwrap_err(), + RopeError::InvalidSections + ); + assert_eq!( + RotarySpec::<8>::interleaved(1.0, 4, [1, 1, 2]).unwrap_err(), + RopeError::InvalidSections + ); + assert_eq!( + RotarySpec::<128>::interleaved(1.0, 4, [usize::MAX, 1, 0]).unwrap_err(), + RopeError::ExtentOverflow + ); + assert!(RotarySpec::<8>::interleaved(1.0, 4, [4, 0, 0]).is_ok()); + assert!(RotarySpec::<8>::interleaved(1.0, 4, [2, 1, 1]).is_ok()); + } + + /// Reject mismatched or unaddressable tables and nonfinite supplied coefficients. + #[test] + fn coefficient_boundary_checks_length_overflow_and_finiteness() { + let one = BFloat16::from_bits(0x3f80); + let zero = BFloat16::from_bits(0); + assert_eq!( + HostTables::<4>::from_coefficients(1, vec![one], vec![zero; 2]).unwrap_err(), + RopeError::LengthMismatch + ); + assert_eq!( + HostTables::<4>::from_coefficients(1, vec![one; 2], vec![zero; 3]).unwrap_err(), + RopeError::LengthMismatch + ); + assert_eq!( + HostTables::<4>::from_coefficients(usize::MAX, vec![], vec![]).unwrap_err(), + RopeError::ExtentOverflow + ); + assert_eq!( + HostTables::<2>::from_coefficients(isize::MAX as usize / 2 + 1, vec![], vec![]) + .unwrap_err(), + RopeError::ExtentOverflow + ); + assert_eq!( + HostTables::<0>::from_coefficients(0, vec![], vec![]).unwrap_err(), + RopeError::InvalidDimension + ); + assert_eq!( + HostTables::<3>::from_coefficients(0, vec![], vec![]).unwrap_err(), + RopeError::InvalidDimension + ); + for invalid in [0x7f80, 0xff80, 0x7fc1] { + assert_eq!( + HostTables::<2>::from_coefficients( + 1, + vec![BFloat16::from_bits(invalid)], + vec![zero] + ) + .unwrap_err(), + RopeError::NonFiniteCoefficient + ); + assert_eq!( + HostTables::<2>::from_coefficients( + 1, + vec![one], + vec![BFloat16::from_bits(invalid)] + ) + .unwrap_err(), + RopeError::NonFiniteCoefficient + ); + } + let valid = HostTables::<4>::from_coefficients(1, vec![one; 2], vec![zero; 2]).unwrap(); + assert_eq!(valid.tokens(), 1); + assert_eq!(bits(valid.cos()), [0x3f80; 2]); + assert_eq!(bits(valid.sin()), [0; 2]); + } +} diff --git a/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.lock new file mode 100644 index 0000000000..c3723ec22a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.lock @@ -0,0 +1,807 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[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", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[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 = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[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 = "ctm-rs" +version = "0.1.0" +dependencies = [ + "cuda-device", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.5.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.5.0" +dependencies = [ + "anyhow", + "bytemuck", + "cuda-bindings", + "cuda-core-derive", + "cuda-repr", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.5.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "bytemuck", + "cuda-macros", + "cuda-repr", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-repr" +version = "0.5.0" + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "cute-layout" +version = "0.1.0" + +[[package]] +name = "cute-rs" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cute-layout", +] + +[[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.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[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 = "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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "zerocopy", +] + +[[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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +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.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[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.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[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" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[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" +dependencies = [ + "libloading", + "thiserror", +] + +[[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 = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +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" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[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 = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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 = "sm100_fmha_prefill_cute" +version = "0.1.0" +dependencies = [ + "ctm-rs", + "cuda-core", + "cuda-device", + "cuda-host", + "cute-rs", + "serde_json", +] + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[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.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.toml b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.toml new file mode 100644 index 0000000000..e7da546e0b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "sm100_fmha_prefill_cute" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +description = "Source-matched Blackwell SM100 FP16 FMHA prefill in CuTe Rust" + +[workspace] + +[dependencies] +ctm-rs = { path = "../../ctm-rs" } +cute-rs = { path = "../../cute-rs", features = ["host"] } +cuda-device = { path = "../../../crates/cuda-device" } +cuda-host = { path = "../../../crates/cuda-host" } +cuda-core = { path = "../../../../cuda-core" } +serde_json = "1" diff --git a/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/README.md b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/README.md new file mode 100644 index 0000000000..f1f63a233f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/README.md @@ -0,0 +1,88 @@ +# FP16 FMHA prefill with CuTe-RS — SM100a (B200) + +This cuda-oxide example computes `softmax(Q @ K.T / sqrt(128)) @ V` on +data center Blackwell GPUs, using the public CUTLASS 4.7 FMHA schedule. + +The kernel uses `cute-rs` for typed layouts, structured MMA and TMA operations, +and `ctm-rs` for instruction primitives, kernel macros and storage types. +Barrier rings distinguish TMA producers from ordinary arrivals; packed FP32 +arithmetic and TMEM allocation retain their CTM types. +The manifest retains `cuda-device` because macro expansions still refer to +that crate internally. CuTe operations retain their CuTe dialect lowering +through `cutlass-mlir`; the primitives-only compiler profile does not cover +this example. + +Direct instruction calls and workload-local fragment handling remain legacy +migration work under the [CuTe kernel architecture](../../AGENTS.md). Tensor +computation belongs in reusable `cute-rs` operations. + +The measurements below validate both implementations on the current branch. +This kernel requires +SM100 hardware; the RTX 5090 / SM120 cannot execute it. + +- FP16 inputs/output, FP32 accumulation, and FP16 probabilities before PV. +- Contiguous BSHD tensors, head dimension **128**, runtime batch and lengths. +- MHA or GQA (`Hq % Hkv == 0`), dense attention (default) or top-left causal + masking. Partial query and KV tiles are supported. + +From `cuda-oxide/`, once per checkout: + +```bash +cargo oxide toolchain install cutlass +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Then: + +```bash +# Compile without launching a GPU kernel +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide build sm100_fmha_prefill_cute --arch sm_100a + +# Check both implementations against the independent oracle on B200 +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide run sm100_fmha_prefill_cute --arch sm_100a -- --shape 1,257,385,4,2 --causal --check-only + +# Dense +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide run sm100_fmha_prefill_cute --arch sm_100a -- --shape 1,1024,1024,8,8 --warmup 100 --iters 100 --windows 5 + +# Dense persistent +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide run sm100_fmha_prefill_cute --arch sm_100a -- --shape 4,4096,4096,8,8 --warmup 100 --iters 100 --windows 5 + +# Causal grouped-query attention (GQA) +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir CUDA_OXIDE_CUTLASS_OPT_LEVEL=2 cargo oxide run sm100_fmha_prefill_cute --arch sm_100a -- --shape 4,4096,4096,8,2 --causal --warmup 100 --iters 100 --windows 5 +``` + +The default runner compares Rust and pinned CuTeDSL on identical inputs and +checks both against an independent FP32 reference before timing. It finishes with: + +```text +CuTeDSL:

+ Qwen/Qwen3-8B dual-precision architecture map +

+ +👉 **[Read the Visual Architecture Guide (`visual.md`)](visual.md)** for a diagram-by-diagram walkthrough of the 36-layer GQA decoder and NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`) block scaling. + +--- + +## RTX 5090 (sm_120a) Live Benchmark: Full 36 Layers (cute-rs vs. CuTeDSL) + +

+ RTX 5090 BF16 vs NVFP4 benchmark comparison +

+ +| Implementation & Precision Mode | 36-Layer Weight Footprint | Compression | Prefill (`32` tok) | Decode (`B = 1`) | +| --- | ---: | ---: | ---: | ---: | +| **cute-rs NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`)** | **`3,726.59 MiB` (`3.64 GiB`)** | **`3.56×`** | **`50.57 ms` (`632.8 tok/s`)** | **`8.77 ms/tok` (`114.0 tok/s`)** | +| **CuTeDSL (`cutlass.cute`) NVFP4 Reference** | `3,726.59 MiB` (`3.64 GiB`) | `3.56×` | `53.08 ms` (`602.9 tok/s`) | `13.20 ms/tok` (`75.8 tok/s`) | +| **cute-rs BF16 (`m16n8k16` Tensor Core)** | `13,248.59 MiB` (`12.94 GiB`) | `1.00×` | `95.89 ms` (`333.7 tok/s`) | `14.32 ms/tok` (`69.8 tok/s`) | + +### Why cute-rs NVFP4 Wins (At a Glance) + +```text +1. 3.56× Less Memory Traffic per Token (NVFP4 vs. BF16) + BF16 weights : [████████████████████████████] 12.94 GiB read/token ──► 69.8 tok/s + NVFP4 weights : [████████] 3.64 GiB read/token ──► 114.0 tok/s (1.63× faster) + +2. Coalesced 32-Lane Warp Reduction + Async Stream Execution (cute-rs vs. CuTeDSL) + CuTeDSL (1 thread/row) : 32 threads read 32 different rows (strided) ──► 75.8 tok/s decode + cute-rs (1 warp/row) : 32 lanes read 1 row along K (1 KB coalesced) ──► 114.0 tok/s decode (1.50× faster) +``` + +- **NVFP4 vs. BF16 (`1.89×` prefill, `1.63×` decode):** Reading `3.64 GiB` instead of `12.94 GiB` per token cuts memory traffic by `3.56×`, and `Nvfp4Tile64::load()` lowers directly to Blackwell hardware `nvgpu.cvt_fpext`. +- **cute-rs (1 warp/output) vs. CuTeDSL (1 thread/output):** Assigning a 32-lane warp to each output row coalesces global `K`-tile loads and saturates all 170 SMs on the RTX 5090, while stream-ordered async execution over pre-allocated `Nvfp4Scratch` buffers eliminates per-layer host synchronization (**`632.8` vs. `602.9 tok/s` prefill**, **`114.0` vs. `75.8 tok/s` decode**). + +--- + +## Quick Start + +From `cuda-oxide/`, once per checkout: + +```bash +cargo oxide toolchain install cutlass +python3 -m venv cute-rs-preview/.venv +cute-rs-preview/.venv/bin/pip install -r cute-rs-preview/requirements.txt +``` + +Then: + +```bash +# Quick 4-layer validation + benchmark on RTX 5090 (sm_120a) +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run qwen3_8b_cute --arch sm_120a -- --layers 4 --prompt-len 16 --decode-steps 8 + +# Full 36-layer Qwen/Qwen3-8B benchmark (synthesizes weights on GPU in <50 ms) +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run qwen3_8b_cute --arch sm_120a -- --full --prompt-len 32 --decode-steps 16 + +# Load a local HuggingFace Qwen/Qwen3-8B safetensors directory +CUDA_OXIDE_DEVICE_BACKEND=cutlass-mlir \ + cargo oxide run qwen3_8b_cute --arch sm_120a -- --full --checkpoint /path/to/Qwen3-8B + +# Run CPU unit tests and rustdoc doctests +cargo test --manifest-path cute-rs-preview/models/qwen3_8b_cute/Cargo.toml + +# Run the CuTeDSL (cutlass.cute) Python reference benchmark on RTX 5090 +cute-rs-preview/.venv/bin/python cute-rs-preview/models/qwen3_8b_cute/reference/benchmark_cutedsl.py --layers 36 --prompt-len 32 --decode-steps 16 +``` + +--- + +## Module Map + +| File | Responsibility | +| --- | --- | +| [`src/config.rs`](src/config.rs) | `Qwen/Qwen3-8B` dimensions (`HIDDEN = 4096`, `INTERMEDIATE = 12288`, `LAYERS = 36`, `Q_HEADS = 32`, `KV_HEADS = 8`, `HEAD_DIM = 128`) and memory footprint helpers. | +| [`src/nvfp4.rs`](src/nvfp4.rs) | True NVFP4 (`Float4E2M1FN` + `Float8E4M3FN` scale, `SF_VEC = 16`) quantizer, canonical `(32, 4)` Blackwell scale-atom layout, and hardware-lowered `cute.scaled_view_load` (`nvgpu.cvt_fpext`) GPU projection kernels. | +| [`src/checkpoint.rs`](src/checkpoint.rs) | Multi-shard `model.safetensors.index.json` loader for `Qwen/Qwen3-8B` plus fast GPU synthetic weight generator. | +| [`src/benchmark.rs`](src/benchmark.rs) | `Qwen3Engine` executing both BF16 and NVFP4 stacks with KV caching, prefill/decode timing, and cosine-similarity verification. | +| [`reference/benchmark_cutedsl.py`](reference/benchmark_cutedsl.py) | CuTeDSL (`cutlass.cute`) Python reference implementation and RTX 5090 benchmark harness. | +| [`visual.md`](visual.md) | Visual walkthrough with adaptive light/dark SVG diagrams. | diff --git a/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/reference/benchmark_cutedsl.py b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/reference/benchmark_cutedsl.py new file mode 100644 index 0000000000..4d3987100f --- /dev/null +++ b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/reference/benchmark_cutedsl.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CuTeDSL (cutlass.cute) reference benchmark for Qwen/Qwen3-8B (BF16 & NVFP4 E2M1+E4M3).""" + +import argparse +import importlib.metadata +import json +from pathlib import Path +import time + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils.blockscaled_layout as blockscaled_utils +import torch +import torch.nn.functional as F + +HIDDEN = 4096 +INTERMEDIATE = 12288 +Q_HEADS = 32 +KV_HEADS = 8 +HEAD_DIM = 128 +Q_WIDTH = Q_HEADS * HEAD_DIM +KV_WIDTH = KV_HEADS * HEAD_DIM +LAYERS = 36 +EPSILON = 1e-6 +ROPE_THETA = 1_000_000.0 + +MMA_TILER_MNK = (128, 1, 64) +SF_VEC_SIZE = 16 + + +class CuTeDslNvfp4GemvKernel: + """CuTeDSL NVFP4 (Float4E2M1FN + Float8E4M3FN, SF_VEC = 16) block-scaled projection kernel.""" + + def __init__(self): + self.threads_per_cta = 128 + + @cute.jit + def __call__( + self, + a_tensor: cute.Tensor, + b_tensor: cute.Tensor, + sfa_tensor: cute.Tensor, + sfb_tensor: cute.Tensor, + c_tensor: cute.Tensor, + stream: cuda.CUstream, + ): + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, SF_VEC_SIZE + ) + sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, SF_VEC_SIZE + ) + sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + grid = ( + cute.ceil_div(c_tensor.shape[0], 128), + 1, + c_tensor.shape[2], + ) + self.kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(1, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mA_mkl: cute.Tensor, + mB_nkl: cute.Tensor, + mSFA_mkl: cute.Tensor, + mSFB_nkl: cute.Tensor, + mC_mnl: cute.Tensor, + ): + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(MMA_TILER_MNK, (None, 0, None)), (None, None, None) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(MMA_TILER_MNK, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(MMA_TILER_MNK, (0, None, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_nkl, cute.slice_(MMA_TILER_MNK, (0, None, None)), (None, None, None) + ) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(MMA_TILER_MNK, (None, None, 0)), (None, None, None) + ) + + tCgC = gC_mnl[tidx, None, bidx, bidy, bidz] + tCgC = cute.make_tensor(tCgC.iterator, 1) + res = cute.zeros_like(tCgC, cutlass.Float32) + + k_tile_cnt = gA_mkl.layout[3].shape + for k_tile in range(k_tile_cnt): + tAgA = gA_mkl[tidx, None, bidx, k_tile, 0] + tBgB = gB_nkl[None, None, bidy, k_tile, bidz] + tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, 0] + tBgSFB = gSFB_nkl[None, None, bidy, k_tile, bidz] + + a_val_nvfp4 = tAgA.load() + b_val_nvfp4 = tBgB.load() + sfa_val_fp8 = tAgSFA.load() + sfb_val_fp8 = tBgSFB.load() + + vec_num = MMA_TILER_MNK[2] // SF_VEC_SIZE + a_val = a_val_nvfp4.to(cutlass.Float32) + b_val = b_val_nvfp4.to(cutlass.Float32) + sfa_shp = (cute.size(sfa_val_fp8.shape) // vec_num, vec_num) + sfa_val = sfa_val_fp8.to(cutlass.Float32).reshape(sfa_shp) + sfb_shp = (cute.size(sfb_val_fp8.shape) // vec_num, vec_num) + sfb_val = sfb_val_fp8.to(cutlass.Float32).reshape(sfb_shp) + + for i in cutlass.range_constexpr(vec_num): + for j in cutlass.range_constexpr(SF_VEC_SIZE): + res += ( + a_val[i * SF_VEC_SIZE + j] + * sfa_val[0, i] + * b_val[i * SF_VEC_SIZE + j] + * sfb_val[0, i] + ) + tCgC.store(res.to(cutlass.BFloat16)) + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def make_cutedsl_args(out_dim: int, in_dim: int, rows: int): + """Allocate fresh NVFP4 (E2M1 + E4M3, SF_VEC = 16) weight/scale/activation views.""" + a_ref = cutlass_torch.matrix(1, out_dim, in_dim, False, cutlass.Float32) + b_ref = cutlass_torch.matrix(rows, 1, in_dim, False, cutlass.Float32) + c_ref = cutlass_torch.matrix(rows, out_dim, 1, True, cutlass.Float32) + + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, cutlass.Float4E2M1FN, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, cutlass.Float4E2M1FN, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, cutlass.BFloat16, is_dynamic_layout=True, assumed_align=16 + ) + + a_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=32) + b_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=32) + c_tensor.mark_compact_shape_dynamic(mode=0, stride_order=(2, 1, 0), divisibility=16) + + sf_k = ceil_div(in_dim, SF_VEC_SIZE) + sfa_mma_shape = (1, ceil_div(out_dim, 128), ceil_div(sf_k, 4), 32, 4, 4) + sfb_mma_shape = (rows, 1, ceil_div(sf_k, 4), 32, 4, 4) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + sfa_f32_cpu = cutlass_torch.create_and_permute_torch_tensor( + sfa_mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=0, max_val=1), + ) + sfb_f32_cpu = cutlass_torch.create_and_permute_torch_tensor( + sfb_mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=0, max_val=1), + ) + sfa_tensor, _ = cutlass_torch.cute_tensor_like( + sfa_f32_cpu, cutlass.Float8E4M3FN, is_dynamic_layout=True, assumed_align=16 + ) + sfb_tensor, _ = cutlass_torch.cute_tensor_like( + sfb_f32_cpu, cutlass.Float8E4M3FN, is_dynamic_layout=True, assumed_align=16 + ) + return (a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor), c_torch + + +def build_cutedsl_projection(out_dim: int, in_dim: int, rows: int, stream: cuda.CUstream): + """Compile a CuTeDSL NVFP4 (E2M1 + E4M3, SF_VEC = 16) projection for [rows, out_dim, in_dim].""" + args, c_torch = make_cutedsl_args(out_dim, in_dim, rows) + kernel = CuTeDslNvfp4GemvKernel() + compiled = cute.compile(kernel, *args, stream) + return compiled, args, c_torch + + +@torch.compile +def rms_norm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + variance = x.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + return (x.to(torch.float32) * torch.rsqrt(variance + EPSILON) * weight.to(torch.float32)).to( + torch.bfloat16 + ) + + +@torch.compile +def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + half = HEAD_DIM // 2 + x1 = x[..., :half].to(torch.float32) + x2 = x[..., half:].to(torch.float32) + c = cos.unsqueeze(1).to(torch.float32) + s = sin.unsqueeze(1).to(torch.float32) + out = torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1) + return out.to(torch.bfloat16) + + +@torch.compile +def quantize_activation_step(x: torch.Tensor): + x_blocks = x.reshape(-1, 16).to(torch.float32) + amax = x_blocks.abs().amax(dim=-1, keepdim=True) + scale = (amax * (1.0 / 6.0)).clamp(min=1.0 / 512.0, max=448.0).to(torch.float8_e4m3fn) + scaled = (x_blocks / scale.to(torch.float32)).clamp(-6.0, 6.0) + return scaled, scale + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layers", type=int, default=36) + parser.add_argument("--prompt-len", type=int, default=32) + parser.add_argument("--decode-steps", type=int, default=16) + parser.add_argument( + "--output", type=Path, default=Path(__file__).parent / "cutedsl_comparison.json" + ) + args = parser.parse_args() + + device = torch.device("cuda", 0) + torch.cuda.set_device(device) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + print( + f"CuTeDSL Reference | Qwen/Qwen3-8B | H={HIDDEN}, I={INTERMEDIATE}, " + f"Q={Q_HEADS}, KV={KV_HEADS}, layers={args.layers}/{LAYERS}" + ) + + shapes = { + "q_o_4096x4096": (Q_WIDTH, HIDDEN), + "kv_1024x4096": (KV_WIDTH, HIDDEN), + "gate_up_12288x4096": (INTERMEDIATE, HIDDEN), + "down_4096x12288": (HIDDEN, INTERMEDIATE), + } + + decode_kernels = {} + prefill_kernels = {} + projection_bench = {} + + for name, (out_dim, in_dim) in shapes.items(): + compiled_dec, dec_args, dec_out = build_cutedsl_projection(out_dim, in_dim, 1, stream) + compiled_pre, pre_args, pre_out = build_cutedsl_projection( + out_dim, in_dim, args.prompt_len, stream + ) + decode_kernels[name] = (compiled_dec, dec_args, dec_out) + prefill_kernels[name] = (compiled_pre, pre_args, pre_out) + + for _ in range(20): + compiled_dec(*dec_args, stream) + compiled_pre(*pre_args, stream) + torch.cuda.synchronize() + + t0 = time.perf_counter() + for _ in range(100): + compiled_dec(*dec_args, stream) + torch.cuda.synchronize() + dec_us = (time.perf_counter() - t0) * 1e6 / 100.0 + + t1 = time.perf_counter() + for _ in range(50): + compiled_pre(*pre_args, stream) + torch.cuda.synchronize() + pre_us = (time.perf_counter() - t1) * 1e6 / 50.0 + + projection_bench[name] = { + "shape": [out_dim, in_dim], + "decode_1tok_us": round(dec_us, 2), + "prefill_32tok_us": round(pre_us, 2), + } + print( + f" [CuTeDSL NVFP4 Kernel] {name:<20} ({out_dim:>5}x{in_dim:<5}): " + f"decode (1 tok) = {dec_us:>6.2f} us | prefill ({args.prompt_len} tok) = {pre_us:>7.2f} us" + ) + + norm_w = torch.ones(HIDDEN, dtype=torch.bfloat16, device=device) + head_norm_w = torch.ones(HEAD_DIM, dtype=torch.bfloat16, device=device) + inv_freq = 1.0 / ( + ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32, device=device) / HEAD_DIM) + ) + + max_seq = args.prompt_len + args.decode_steps + 16 + k_caches = [ + torch.zeros((max_seq, KV_HEADS, HEAD_DIM), dtype=torch.bfloat16, device=device) + for _ in range(args.layers) + ] + v_caches = [ + torch.zeros((max_seq, KV_HEADS, HEAD_DIM), dtype=torch.bfloat16, device=device) + for _ in range(args.layers) + ] + + # Allocate 4 distinct layer weight sets (~414 MiB > 96 MB RTX 5090 L2 cache) so + # 36-layer execution streams weights from GDDR7 DRAM instead of staying L2-resident. + num_weight_sets = 4 + decode_layer_sets = [] + prefill_layer_sets = [] + for _ in range(num_weight_sets): + dec_set = {} + pre_set = {} + for name, (out_dim, in_dim) in shapes.items(): + dec_set[name] = (decode_kernels[name][0], *make_cutedsl_args(out_dim, in_dim, 1)) + pre_set[name] = ( + prefill_kernels[name][0], + *make_cutedsl_args(out_dim, in_dim, args.prompt_len), + ) + decode_layer_sets.append(dec_set) + prefill_layer_sets.append(pre_set) + + def run_cutedsl_nvfp4_forward(x: torch.Tensor, start_pos: int, layer_sets: list): + rows = x.shape[0] + pos = torch.arange(start_pos, start_pos + rows, dtype=torch.float32, device=device) + freqs = torch.outer(pos, inv_freq) + cos = torch.cos(freqs).to(torch.bfloat16) + sin = torch.sin(freqs).to(torch.bfloat16) + + hidden = x + for layer_idx in range(args.layers): + kernel_set = layer_sets[layer_idx % len(layer_sets)] + q_fn, q_args, q_out = kernel_set["q_o_4096x4096"] + kv_fn, kv_args, kv_out = kernel_set["kv_1024x4096"] + gu_fn, gu_args, gu_out = kernel_set["gate_up_12288x4096"] + dn_fn, dn_args, dn_out = kernel_set["down_4096x12288"] + + normed = rms_norm(hidden, norm_w) + quantize_activation_step(normed) + q_fn(*q_args, stream) + kv_fn(*kv_args, stream) + kv_fn(*kv_args, stream) + + q = rms_norm(q_out.reshape(rows, Q_HEADS, HEAD_DIM), head_norm_w) + k = rms_norm(kv_out.reshape(rows, KV_HEADS, HEAD_DIM), head_norm_w) + v = kv_out.reshape(rows, KV_HEADS, HEAD_DIM) + + q = apply_rope(q, cos, sin) + k = apply_rope(k, cos, sin) + + k_caches[layer_idx][start_pos : start_pos + rows] = k + v_caches[layer_idx][start_pos : start_pos + rows] = v + + k_live = k_caches[layer_idx][: start_pos + rows].permute(1, 0, 2).unsqueeze(0) + v_live = v_caches[layer_idx][: start_pos + rows].permute(1, 0, 2).unsqueeze(0) + q_live = q.permute(1, 0, 2).unsqueeze(0) + attn_out = F.scaled_dot_product_attention( + q_live, k_live, v_live, is_causal=(rows > 1), enable_gqa=True + ) + attn_flat = attn_out.squeeze(0).permute(1, 0, 2).reshape(rows, HIDDEN) + + quantize_activation_step(attn_flat) + q_fn(*q_args, stream) + residual = hidden + q_out.reshape(rows, HIDDEN) + + post = rms_norm(residual, norm_w) + quantize_activation_step(post) + gu_fn(*gu_args, stream) + gate = gu_out.reshape(rows, INTERMEDIATE) + gu_fn(*gu_args, stream) + up = gu_out.reshape(rows, INTERMEDIATE) + swiglu = F.silu(gate) * up + + quantize_activation_step(swiglu) + dn_fn(*dn_args, stream) + hidden = residual + dn_out.reshape(rows, HIDDEN) + + return rms_norm(hidden, norm_w) + + x_prefill = torch.randn((args.prompt_len, HIDDEN), dtype=torch.bfloat16, device=device) + x_decode = torch.randn((1, HIDDEN), dtype=torch.bfloat16, device=device) + + _ = run_cutedsl_nvfp4_forward(x_prefill, 0, prefill_layer_sets) + _ = run_cutedsl_nvfp4_forward(x_decode, args.prompt_len, decode_layer_sets) + torch.cuda.synchronize() + + t_pre = time.perf_counter() + _ = run_cutedsl_nvfp4_forward(x_prefill, 0, prefill_layer_sets) + torch.cuda.synchronize() + cutedsl_prefill_ms = (time.perf_counter() - t_pre) * 1000.0 + cutedsl_prefill_tps = args.prompt_len / (cutedsl_prefill_ms / 1000.0) + + t_dec = time.perf_counter() + for step in range(args.decode_steps): + x_decode = run_cutedsl_nvfp4_forward(x_decode, args.prompt_len + step, decode_layer_sets) + torch.cuda.synchronize() + cutedsl_decode_total_ms = (time.perf_counter() - t_dec) * 1000.0 + cutedsl_decode_ms = cutedsl_decode_total_ms / args.decode_steps + cutedsl_decode_tps = args.decode_steps / (cutedsl_decode_total_ms / 1000.0) + + print( + f" [CuTeDSL NVFP4 (E2M1+E4M3, SF=16)] prefill ({args.prompt_len} tok): " + f"{cutedsl_prefill_ms:7.2f} ms ({cutedsl_prefill_tps:7.1f} tok/s) | " + f"decode ({args.decode_steps} steps): {cutedsl_decode_ms:6.2f} ms/tok ({cutedsl_decode_tps:6.1f} tok/s)" + ) + + report = { + "model": "Qwen/Qwen3-8B", + "gpu": torch.cuda.get_device_name(0), + "cutlass_dsl_version": importlib.metadata.version("nvidia-cutlass-dsl"), + "torch_version": torch.__version__, + "layers": args.layers, + "prompt_len": args.prompt_len, + "decode_steps": args.decode_steps, + "cutedsl_nvfp4_projections": projection_bench, + "cutedsl_nvfp4_e2e": { + "prefill_ms": round(cutedsl_prefill_ms, 2), + "prefill_tok_per_sec": round(cutedsl_prefill_tps, 1), + "decode_ms_per_tok": round(cutedsl_decode_ms, 2), + "decode_tok_per_sec": round(cutedsl_decode_tps, 1), + }, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() + diff --git a/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/benchmark.rs b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/benchmark.rs new file mode 100644 index 0000000000..c562187423 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/benchmark.rs @@ -0,0 +1,984 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! End-to-end `Qwen3-8B` BF16 and NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) engine and benchmark suite. +//! +//! Measures prefill (`tokens/sec`), autoregressive decode (`tokens/sec`), weight memory +//! savings, and numerical cosine similarity between BF16 and NVFP4 on RTX 5090 (`sm_120a`). + +use crate::checkpoint::{ + OwnedNvfp4Layer, Qwen3Bf16Stack, Qwen3Nvfp4Stack, load_qwen3_bf16_checkpoint, + synthesize_qwen3_bf16, +}; +use crate::config::{ + EPSILON, HEAD_DIM, HIDDEN, INTERMEDIATE, KV_HEADS, KV_WIDTH, MAX_POSITION, Q_HEADS, Q_WIDTH, + ROPE_THETA, +}; +use crate::model::Batch; +use crate::norm_kernel; +use crate::nvfp4::Nvfp4Runner; +use crate::operator::Result; +use crate::rotary::RotarySpec; +use cuda_core::{CudaContext, CudaStream, DeviceBuffer, LaunchConfig1D}; +use cute_rs::BFloat16; +use std::path::Path; +use std::sync::Arc; +use std::time::Instant; + +/// Summary metrics for one precision mode (`BF16` or `NVFP4`) on `Qwen3-8B`. +#[derive(Clone, Debug)] +pub struct ModeMetrics { + /// Precision label (`"BF16"` or `"NVFP4 (E2M1 + E4M3, SF=16)"`). + pub label: &'static str, + /// Number of executed decoder layers. + pub layers: usize, + /// Total device weight footprint in bytes. + pub weight_bytes: usize, + /// Number of prompt tokens processed during prefill. + pub prompt_tokens: usize, + /// Prefill latency in milliseconds. + pub prefill_ms: f64, + /// Prefill throughput in tokens per second. + pub prefill_tok_per_sec: f64, + /// Number of autoregressive decode steps measured. + pub decode_steps: usize, + /// Mean per-token decode latency in milliseconds. + pub decode_ms_per_tok: f64, + /// Autoregressive decode throughput in tokens per second. + pub decode_tok_per_sec: f64, +} + +/// Complete comparison report between `Qwen3-8B` BF16 and NVFP4 execution. +#[derive(Clone, Debug)] +pub struct BenchmarkReport { + /// Metrics for the BF16 Tensor-Core baseline. + pub bf16: ModeMetrics, + /// Metrics for the NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) block-scaled path. + pub nvfp4: ModeMetrics, + /// Cosine similarity between the final normalized hidden states of BF16 and NVFP4. + pub cosine_similarity: f64, + /// Compression ratio of total layer weights (`bf16.weight_bytes / nvfp4.weight_bytes`). + pub compression_ratio: f64, +} + +/// Pre-allocated device KV-cache for one NVFP4 `Qwen3-8B` layer (`[capacity, KV_HEADS, HEAD_DIM]`). +pub struct Nvfp4KvCache { + /// Rotated key cache `[capacity, KV_HEADS, HEAD_DIM]`. + pub key: DeviceBuffer, + /// Unrotated value cache `[capacity, KV_HEADS, HEAD_DIM]`. + pub value: DeviceBuffer, + /// Current valid token prefix length. + pub len: usize, + /// Maximum token capacity. + pub capacity: usize, +} + +impl Nvfp4KvCache { + /// Allocate a zero-initialized KV cache for `capacity` tokens on `stream`. + pub fn new(stream: &CudaStream, capacity: usize) -> Result { + let elems = capacity * KV_WIDTH + 8; + let zero = vec![BFloat16::from_bits(0); elems]; + Ok(Self { + key: DeviceBuffer::from_host(stream, &zero)?, + value: DeviceBuffer::from_host(stream, &zero)?, + len: 0, + capacity, + }) + } +} + +/// Pre-allocated device scratch buffers for zero-allocation, stream-ordered 36-layer NVFP4 execution. +struct Nvfp4Scratch { + /// Hidden states entering the first layer `[rows, HIDDEN]`; a run takes both hidden + /// buffers out of their options and alternates them as layer input and output. + hidden_a: Option>, + /// The other hidden-state buffer `[rows, HIDDEN]` of that alternation. + hidden_b: Option>, + /// Input RMSNorm output `[rows, HIDDEN]`, the operand of the Q/K/V projections. + normalized: DeviceBuffer, + /// Query projection `[rows, Q_WIDTH]`. + q: DeviceBuffer, + /// Key projection `[rows, KV_WIDTH]`. + k: DeviceBuffer, + /// Value projection `[rows, KV_WIDTH]`. + v: DeviceBuffer, + /// Per-head QK-normed queries `[rows, Q_WIDTH]`. + qn: DeviceBuffer, + /// Per-head QK-normed keys `[rows, KV_WIDTH]`. + kn: DeviceBuffer, + /// Rotated queries `[rows, Q_WIDTH]`. + rot_q: DeviceBuffer, + /// Rotated keys `[rows, KV_WIDTH]`, appended to the layer's KV cache. + rot_k: DeviceBuffer, + /// Split attention partial states `[rows, Q_HEADS, splits, STATE_SIZE]` in `f32`. + partials: DeviceBuffer, + /// Merged attention output `[rows, Q_WIDTH]`. + attended: DeviceBuffer, + /// Output projection `[rows, HIDDEN]`. + projected: DeviceBuffer, + /// Layer input plus the attention output `[rows, HIDDEN]`. + residual: DeviceBuffer, + /// Post-attention RMSNorm output `[rows, HIDDEN]`, the operand of the gate and up projections. + post: DeviceBuffer, + /// Gate projection `[rows, INTERMEDIATE]`. + gate: DeviceBuffer, + /// Up projection `[rows, INTERMEDIATE]`. + up: DeviceBuffer, + /// `SiLU(gate)` `[rows, INTERMEDIATE]`. + silu: DeviceBuffer, + /// `SiLU(gate) * up` `[rows, INTERMEDIATE]`, the operand of the down projection. + product: DeviceBuffer, + /// Down projection `[rows, HIDDEN]`, added to `residual` to form the layer output. + down: DeviceBuffer, +} + +impl Nvfp4Scratch { + /// Allocate zeroed scratch for up to `max_rows` token rows and the attention splits of + /// `max_seq` cached tokens. Every buffer carries eight guard elements past its logical length. + fn new(stream: &CudaStream, max_rows: usize, max_seq: usize) -> Result { + let splits = max_seq.div_ceil(64); + let partial_elems = max_rows * Q_HEADS * splits * crate::attention_kernel::STATE_SIZE + 8; + let z_bf16 = + |n: usize| DeviceBuffer::from_host(stream, &vec![BFloat16::from_bits(0); n + 8]); + Ok(Self { + hidden_a: Some(z_bf16(max_rows * HIDDEN)?), + hidden_b: Some(z_bf16(max_rows * HIDDEN)?), + normalized: z_bf16(max_rows * HIDDEN)?, + q: z_bf16(max_rows * Q_WIDTH)?, + k: z_bf16(max_rows * KV_WIDTH)?, + v: z_bf16(max_rows * KV_WIDTH)?, + qn: z_bf16(max_rows * Q_WIDTH)?, + kn: z_bf16(max_rows * KV_WIDTH)?, + rot_q: z_bf16(max_rows * Q_WIDTH)?, + rot_k: z_bf16(max_rows * KV_WIDTH)?, + partials: DeviceBuffer::from_host(stream, &vec![0.0f32; partial_elems])?, + attended: z_bf16(max_rows * Q_WIDTH)?, + projected: z_bf16(max_rows * HIDDEN)?, + residual: z_bf16(max_rows * HIDDEN)?, + post: z_bf16(max_rows * HIDDEN)?, + gate: z_bf16(max_rows * INTERMEDIATE)?, + up: z_bf16(max_rows * INTERMEDIATE)?, + silu: z_bf16(max_rows * INTERMEDIATE)?, + product: z_bf16(max_rows * INTERMEDIATE)?, + down: z_bf16(max_rows * HIDDEN)?, + }) + } +} + +/// Combined `Qwen3-8B` execution engine supporting both BF16 and NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`) stacks. +pub struct Qwen3Engine { + /// Shared execution stream. + stream: Arc, + /// NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) linear projection runner. + nvfp4_runner: Nvfp4Runner, + /// Loaded RMSNorm module for stream-ordered execution. + norms: norm_kernel::kernels::LoadedModule, + /// Loaded rotary/append module for stream-ordered execution. + rotary_mod: crate::kernel::kernels::LoadedModule, + /// Loaded split-KV attention module for stream-ordered execution. + attn_mod: crate::attention_kernel::kernels::LoadedModule, + /// Loaded elementwise add module for stream-ordered execution. + elem_mod: crate::element_kernel::kernels::LoadedModule, + /// Loaded BF16 MMA & SwiGLU activation module for stream-ordered execution. + ffn_mod: crate::ffn_kernel::kernels::LoadedModule, + /// Pre-allocated scratch buffers for zero-sync 36-layer execution. + scratch: std::sync::Mutex, + /// Standard 1D RoPE specification (`theta = 1_000_000.0`, `D = 128`). + rotary: RotarySpec, +} + +/// Make a 128-thread 1D launch configuration for `work` items. +fn grid1d(work: usize) -> cuda_core::LaunchConfig { + cuda_core::LaunchConfig { + grid_dim: (work.div_ceil(128) as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + } +} + +impl Qwen3Engine { + /// Load all BF16, NVFP4, RMSNorm, RoPE, and GQA modules onto `stream`. + pub fn new(stream: Arc) -> Result { + let ctx = stream.context(); + Ok(Self { + nvfp4_runner: Nvfp4Runner::new(stream.clone())?, + norms: norm_kernel::kernels::load(ctx)?, + rotary_mod: crate::kernel::kernels::load(ctx)?, + attn_mod: unsafe { crate::attention_kernel::kernels::load(ctx)? }, + elem_mod: crate::element_kernel::kernels::load(ctx)?, + ffn_mod: unsafe { crate::ffn_kernel::kernels::load(ctx)? }, + scratch: std::sync::Mutex::new(Nvfp4Scratch::new(&stream, 256, 512)?), + rotary: RotarySpec::standard(ROPE_THETA, MAX_POSITION)?, + stream, + }) + } + + /// Borrow the execution stream. + #[must_use] + pub fn stream(&self) -> &Arc { + &self.stream + } + + /// Enqueue one BF16 projection asynchronously on `self.stream`. + fn enqueue_project_bf16( + &self, + input: &DeviceBuffer, + weight: &DeviceBuffer, + output: &mut DeviceBuffer, + rows: usize, + width: usize, + outputs: usize, + ) -> Result<()> { + unsafe { + if rows <= 2 { + let launch = self + .ffn_mod + .prepare_decode_linear_bf16(LaunchConfig1D::new( + outputs.div_ceil(8) as u32, + 32, + 0, + ))?; + self.ffn_mod.decode_linear_bf16( + &self.stream, + &launch, + input, + weight, + output, + rows, + width, + outputs, + width, + width, + )?; + } else { + let tiles = rows.div_ceil(16) * outputs.div_ceil(8); + let cfg = cuda_core::LaunchConfig { + grid_dim: (tiles.div_ceil(4) as u32, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + }; + self.ffn_mod.linear_bf16( + &self.stream, + cfg, + input, + weight, + output, + rows, + width, + outputs, + width, + width, + )?; + } + } + Ok(()) + } + + /// Enqueue one `Qwen3-8B` NVFP4 decoder layer asynchronously on `self.stream`. + #[allow(clippy::too_many_arguments)] + fn enqueue_layer_nvfp4( + &self, + hidden_in: &DeviceBuffer, + hidden_out: &mut DeviceBuffer, + rows: usize, + layer: &OwnedNvfp4Layer, + cache: &mut Nvfp4KvCache, + cos: &DeviceBuffer, + sin: &DeviceBuffer, + s: &mut Nvfp4Scratch, + ) -> Result<()> { + let start = cache.len; + let key_tokens = start + rows; + let splits = key_tokens.div_ceil(64); + let heads = rows * Q_HEADS; + let stream = &self.stream; + + unsafe { + crate::norm_specialized::enqueue( + &self.norms, + stream, + hidden_in, + &layer.input_norm, + &mut s.normalized, + rows, + HIDDEN, + EPSILON, + )?; + self.nvfp4_runner.enqueue_qkv( + &s.normalized, + rows, + &layer.q_proj, + &layer.k_proj, + &layer.v_proj, + &mut s.q, + &mut s.k, + &mut s.v, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.q, + &layer.q_norm, + &mut s.qn, + rows * Q_HEADS, + HEAD_DIM, + EPSILON, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.k, + &layer.k_norm, + &mut s.kn, + rows * KV_HEADS, + HEAD_DIM, + EPSILON, + )?; + self.rotary_mod.rotate_bf16( + stream, + grid1d(rows * (Q_WIDTH / 2)), + &s.qn, + cos, + sin, + &mut s.rot_q, + rows, + Q_HEADS, + HEAD_DIM, + )?; + self.rotary_mod.rotate_bf16( + stream, + grid1d(rows * (KV_WIDTH / 2)), + &s.kn, + cos, + sin, + &mut s.rot_k, + rows, + KV_HEADS, + HEAD_DIM, + )?; + self.rotary_mod.append_bf16( + stream, + grid1d(rows * KV_WIDTH), + &s.rot_k, + &s.v, + &mut cache.key, + &mut cache.value, + start * KV_WIDTH, + rows * KV_WIDTH, + )?; + let partial_launch = self.attn_mod.prepare_partial(LaunchConfig1D::new( + (heads * splits) as u32, + 128, + crate::attention_kernel::PARTIAL_SHARED_BYTES, + ))?; + self.attn_mod.partial( + stream, + &partial_launch, + &s.rot_q, + &cache.key, + &cache.value, + &mut s.partials, + rows as u32, + key_tokens as u32, + Q_HEADS as u32, + KV_HEADS as u32, + start as u32, + true, + )?; + let merge_launch = + self.attn_mod + .prepare_merge(LaunchConfig1D::new(heads as u32, 32, 0))?; + self.attn_mod.merge( + stream, + &merge_launch, + &s.partials, + &mut s.attended, + rows as u32, + Q_HEADS as u32, + splits as u32, + )?; + self.nvfp4_runner.enqueue_project( + &s.attended, + rows, + &layer.o_proj, + &mut s.projected, + )?; + self.elem_mod.add_bf16( + stream, + grid1d(rows * HIDDEN), + hidden_in, + &s.projected, + &mut s.residual, + rows * HIDDEN, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.residual, + &layer.post_attention_norm, + &mut s.post, + rows, + HIDDEN, + EPSILON, + )?; + self.nvfp4_runner.enqueue_gate_up( + &s.post, + rows, + &layer.gate_proj, + &layer.up_proj, + &mut s.gate, + &mut s.up, + )?; + self.ffn_mod.activate_bf16( + stream, + grid1d((rows * INTERMEDIATE).div_ceil(8)), + &s.gate, + &s.up, + &mut s.silu, + &mut s.product, + rows * INTERMEDIATE, + )?; + self.nvfp4_runner + .enqueue_project(&s.product, rows, &layer.down_proj, &mut s.down)?; + self.elem_mod.add_bf16( + stream, + grid1d(rows * HIDDEN), + &s.residual, + &s.down, + hidden_out, + rows * HIDDEN, + )?; + } + cache.len = key_tokens; + Ok(()) + } + + /// Execute the `Qwen3-8B` BF16 stack over `input` at 1D token `positions`. + /// Enqueue one `Qwen3-8B` BF16 decoder layer asynchronously on `self.stream`. + #[allow(clippy::too_many_arguments)] + fn enqueue_layer_bf16( + &self, + hidden_in: &DeviceBuffer, + hidden_out: &mut DeviceBuffer, + rows: usize, + layer: &crate::checkpoint::OwnedBf16Layer, + cache: &mut Nvfp4KvCache, + cos: &DeviceBuffer, + sin: &DeviceBuffer, + s: &mut Nvfp4Scratch, + ) -> Result<()> { + let start = cache.len; + let key_tokens = start + rows; + let splits = key_tokens.div_ceil(64); + let heads = rows * Q_HEADS; + let stream = &self.stream; + + unsafe { + crate::norm_specialized::enqueue( + &self.norms, + stream, + hidden_in, + &layer.input_norm, + &mut s.normalized, + rows, + HIDDEN, + EPSILON, + )?; + self.enqueue_project_bf16( + &s.normalized, + &layer.q_proj, + &mut s.q, + rows, + HIDDEN, + Q_WIDTH, + )?; + self.enqueue_project_bf16( + &s.normalized, + &layer.k_proj, + &mut s.k, + rows, + HIDDEN, + KV_WIDTH, + )?; + self.enqueue_project_bf16( + &s.normalized, + &layer.v_proj, + &mut s.v, + rows, + HIDDEN, + KV_WIDTH, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.q, + &layer.q_norm, + &mut s.qn, + rows * Q_HEADS, + HEAD_DIM, + EPSILON, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.k, + &layer.k_norm, + &mut s.kn, + rows * KV_HEADS, + HEAD_DIM, + EPSILON, + )?; + self.rotary_mod.rotate_bf16( + stream, + grid1d(rows * (Q_WIDTH / 2)), + &s.qn, + cos, + sin, + &mut s.rot_q, + rows, + Q_HEADS, + HEAD_DIM, + )?; + self.rotary_mod.rotate_bf16( + stream, + grid1d(rows * (KV_WIDTH / 2)), + &s.kn, + cos, + sin, + &mut s.rot_k, + rows, + KV_HEADS, + HEAD_DIM, + )?; + self.rotary_mod.append_bf16( + stream, + grid1d(rows * KV_WIDTH), + &s.rot_k, + &s.v, + &mut cache.key, + &mut cache.value, + start * KV_WIDTH, + rows * KV_WIDTH, + )?; + let partial_launch = self.attn_mod.prepare_partial(LaunchConfig1D::new( + (heads * splits) as u32, + 128, + crate::attention_kernel::PARTIAL_SHARED_BYTES, + ))?; + self.attn_mod.partial( + stream, + &partial_launch, + &s.rot_q, + &cache.key, + &cache.value, + &mut s.partials, + rows as u32, + key_tokens as u32, + Q_HEADS as u32, + KV_HEADS as u32, + start as u32, + true, + )?; + let merge_launch = + self.attn_mod + .prepare_merge(LaunchConfig1D::new(heads as u32, 32, 0))?; + self.attn_mod.merge( + stream, + &merge_launch, + &s.partials, + &mut s.attended, + rows as u32, + Q_HEADS as u32, + splits as u32, + )?; + self.enqueue_project_bf16( + &s.attended, + &layer.o_proj, + &mut s.projected, + rows, + Q_WIDTH, + HIDDEN, + )?; + self.elem_mod.add_bf16( + stream, + grid1d(rows * HIDDEN), + hidden_in, + &s.projected, + &mut s.residual, + rows * HIDDEN, + )?; + crate::norm_specialized::enqueue( + &self.norms, + stream, + &s.residual, + &layer.post_attention_norm, + &mut s.post, + rows, + HIDDEN, + EPSILON, + )?; + self.enqueue_project_bf16( + &s.post, + &layer.gate_proj, + &mut s.gate, + rows, + HIDDEN, + INTERMEDIATE, + )?; + self.enqueue_project_bf16( + &s.post, + &layer.up_proj, + &mut s.up, + rows, + HIDDEN, + INTERMEDIATE, + )?; + self.ffn_mod.activate_bf16( + stream, + grid1d((rows * INTERMEDIATE).div_ceil(8)), + &s.gate, + &s.up, + &mut s.silu, + &mut s.product, + rows * INTERMEDIATE, + )?; + self.enqueue_project_bf16( + &s.product, + &layer.down_proj, + &mut s.down, + rows, + INTERMEDIATE, + HIDDEN, + )?; + self.elem_mod.add_bf16( + stream, + grid1d(rows * HIDDEN), + &s.residual, + &s.down, + hidden_out, + rows * HIDDEN, + )?; + } + cache.len = key_tokens; + Ok(()) + } + + /// Execute the `Qwen3-8B` BF16 stack over `input` asynchronously on `self.stream`. + pub fn forward_bf16( + &self, + input: &Batch, + stack: &Qwen3Bf16Stack, + caches: &mut [Nvfp4KvCache], + cos: &DeviceBuffer, + sin: &DeviceBuffer, + ) -> Result> { + if caches.len() != stack.layers.len() { + return Err("BF16 layer count does not match KV cache count".into()); + } + let rows = input.rows(); + let mut s = self.scratch.lock().map_err(|_| "poisoned scratch")?; + let mut hidden_a = s.hidden_a.take().ok_or("missing hidden_a scratch")?; + let mut hidden_b = s.hidden_b.take().ok_or("missing hidden_b scratch")?; + + self.enqueue_layer_bf16( + input.buffer(), + &mut hidden_a, + rows, + &stack.layers[0], + &mut caches[0], + cos, + sin, + &mut s, + )?; + for (idx, (layer, cache)) in stack.layers[1..].iter().zip(&mut caches[1..]).enumerate() { + if idx % 2 == 0 { + self.enqueue_layer_bf16( + &hidden_a, + &mut hidden_b, + rows, + layer, + cache, + cos, + sin, + &mut s, + )?; + } else { + self.enqueue_layer_bf16( + &hidden_b, + &mut hidden_a, + rows, + layer, + cache, + cos, + sin, + &mut s, + )?; + } + } + let final_in = if stack.layers.len() % 2 == 1 { + &hidden_a + } else { + &hidden_b + }; + let mut out = self.nvfp4_runner.allocate_batch::(rows)?; + unsafe { + crate::norm_specialized::enqueue( + &self.norms, + &self.stream, + final_in, + &stack.final_norm, + out.buffer_mut(), + rows, + HIDDEN, + EPSILON, + )?; + } + s.hidden_a = Some(hidden_a); + s.hidden_b = Some(hidden_b); + Ok(out) + } + + /// Execute the `Qwen3-8B` NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) stack over `input`. + /// + /// Enqueues all 36 layers asynchronously onto `self.stream` over pre-allocated + /// device scratch buffers without per-layer allocation or host synchronization. + pub fn forward_nvfp4( + &self, + input: &Batch, + stack: &Qwen3Nvfp4Stack, + caches: &mut [Nvfp4KvCache], + cos: &DeviceBuffer, + sin: &DeviceBuffer, + ) -> Result> { + if caches.len() != stack.layers.len() { + return Err("NVFP4 layer count does not match KV cache count".into()); + } + let rows = input.rows(); + let mut s = self.scratch.lock().map_err(|_| "poisoned scratch")?; + let mut hidden_a = s.hidden_a.take().ok_or("missing hidden_a scratch")?; + let mut hidden_b = s.hidden_b.take().ok_or("missing hidden_b scratch")?; + + self.enqueue_layer_nvfp4( + input.buffer(), + &mut hidden_a, + rows, + &stack.layers[0], + &mut caches[0], + cos, + sin, + &mut s, + )?; + for (idx, (layer, cache)) in stack.layers[1..].iter().zip(&mut caches[1..]).enumerate() { + if idx % 2 == 0 { + self.enqueue_layer_nvfp4( + &hidden_a, + &mut hidden_b, + rows, + layer, + cache, + cos, + sin, + &mut s, + )?; + } else { + self.enqueue_layer_nvfp4( + &hidden_b, + &mut hidden_a, + rows, + layer, + cache, + cos, + sin, + &mut s, + )?; + } + } + let final_in = if stack.layers.len() % 2 == 1 { + &hidden_a + } else { + &hidden_b + }; + let mut out = self.nvfp4_runner.allocate_batch::(rows)?; + unsafe { + crate::norm_specialized::enqueue( + &self.norms, + &self.stream, + final_in, + &stack.final_norm, + out.buffer_mut(), + rows, + HIDDEN, + EPSILON, + )?; + } + s.hidden_a = Some(hidden_a); + s.hidden_b = Some(hidden_b); + Ok(out) + } +} + +/// Compute cosine similarity between two equal-length BF16 slices in `f64`. +#[must_use] +pub fn cosine_similarity_bf16(a: &[BFloat16], b: &[BFloat16]) -> f64 { + let mut dot = 0.0f64; + let mut norm_a = 0.0f64; + let mut norm_b = 0.0f64; + for (&x, &y) in a.iter().zip(b.iter()) { + let xf = x.to_f32() as f64; + let yf = y.to_f32() as f64; + dot += xf * yf; + norm_a += xf * xf; + norm_b += yf * yf; + } + let denom = (norm_a * norm_b).sqrt(); + if denom == 0.0 { 0.0 } else { dot / denom } +} + +/// Run the end-to-end `Qwen3-8B` BF16 vs. NVFP4 (`Float4E2M1FN` + `Float8E4M3FN`, `SF_VEC = 16`) benchmark on GPU 0. +pub fn run_benchmark( + num_layers: usize, + prompt_len: usize, + decode_steps: usize, + checkpoint_path: Option<&Path>, +) -> Result { + let context = CudaContext::new(0)?; + let stream = context.default_stream(); + let engine = Qwen3Engine::new(stream.clone())?; + + let bf16_stack = if let Some(path) = checkpoint_path { + load_qwen3_bf16_checkpoint(&stream, path, num_layers)? + } else { + synthesize_qwen3_bf16(&stream, num_layers)? + }; + let nvfp4_stack = Qwen3Nvfp4Stack::from_bf16_gpu(&stream, &engine.nvfp4_runner, &bf16_stack)?; + + let capacity = prompt_len + decode_steps + 16; + let mut bf16_caches: Vec = (0..num_layers) + .map(|_| Nvfp4KvCache::new(&stream, capacity)) + .collect::>()?; + let mut nvfp4_caches: Vec = (0..num_layers) + .map(|_| Nvfp4KvCache::new(&stream, capacity)) + .collect::>()?; + + let prompt_host: Vec = (0..prompt_len * HIDDEN) + .map(|i| BFloat16::from_f32(((i % 97) as f32 - 48.0) * 0.03125)) + .collect(); + let prompt_batch = Batch::::from_host(&stream, prompt_len, &prompt_host)?; + let prompt_positions: Vec = (0..prompt_len as u32).collect(); + + let prompt_pos3: Vec<[u32; 3]> = prompt_positions.iter().map(|&p| [p, p, p]).collect(); + let prompt_tables = engine.rotary.tables(&prompt_pos3)?; + let prompt_cos = DeviceBuffer::from_host(&stream, prompt_tables.cos())?; + let prompt_sin = DeviceBuffer::from_host(&stream, prompt_tables.sin())?; + + let mut decode_cos_sin = Vec::with_capacity(decode_steps); + for step in 0..decode_steps { + let p = (prompt_len + step) as u32; + let t = engine.rotary.tables(&[[p, p, p]])?; + decode_cos_sin.push(( + DeviceBuffer::from_host(&stream, t.cos())?, + DeviceBuffer::from_host(&stream, t.sin())?, + )); + } + + // Warm up both paths and reset caches before timed runs. + let _ = engine.forward_bf16( + &prompt_batch, + &bf16_stack, + &mut bf16_caches, + &prompt_cos, + &prompt_sin, + )?; + for c in &mut bf16_caches { + c.len = 0; + } + let _ = engine.forward_nvfp4( + &prompt_batch, + &nvfp4_stack, + &mut nvfp4_caches, + &prompt_cos, + &prompt_sin, + )?; + for c in &mut nvfp4_caches { + c.len = 0; + } + stream.synchronize()?; + + // Timed BF16 prefill + decode. + let t0 = Instant::now(); + let bf16_prefill_out = engine.forward_bf16( + &prompt_batch, + &bf16_stack, + &mut bf16_caches, + &prompt_cos, + &prompt_sin, + )?; + stream.synchronize()?; + let bf16_prefill_ms = t0.elapsed().as_secs_f64() * 1000.0; + + let step_host: Vec = prompt_host[..HIDDEN].to_vec(); + let mut step_in = Batch::::from_host(&stream, 1, &step_host)?; + let t_dec0 = Instant::now(); + for (cos, sin) in &decode_cos_sin { + step_in = engine.forward_bf16(&step_in, &bf16_stack, &mut bf16_caches, cos, sin)?; + } + stream.synchronize()?; + let bf16_decode_total_ms = t_dec0.elapsed().as_secs_f64() * 1000.0; + + // Timed NVFP4 (E2M1 + E4M3, SF=16) prefill + decode. + let t1 = Instant::now(); + let nvfp4_prefill_out = engine.forward_nvfp4( + &prompt_batch, + &nvfp4_stack, + &mut nvfp4_caches, + &prompt_cos, + &prompt_sin, + )?; + stream.synchronize()?; + let nvfp4_prefill_ms = t1.elapsed().as_secs_f64() * 1000.0; + + let mut nv_step_in = Batch::::from_host(&stream, 1, &step_host)?; + let t_dec1 = Instant::now(); + for (cos, sin) in &decode_cos_sin { + nv_step_in = + engine.forward_nvfp4(&nv_step_in, &nvfp4_stack, &mut nvfp4_caches, cos, sin)?; + } + stream.synchronize()?; + let nvfp4_decode_total_ms = t_dec1.elapsed().as_secs_f64() * 1000.0; + + let bf16_vec = bf16_prefill_out.read(&stream)?; + let nvfp4_vec = nvfp4_prefill_out.read(&stream)?; + let cosine_similarity = cosine_similarity_bf16(&bf16_vec, &nvfp4_vec); + + let bf16_bytes = bf16_stack.byte_len(); + let nvfp4_bytes = nvfp4_stack.byte_len(); + + Ok(BenchmarkReport { + bf16: ModeMetrics { + label: "BF16 (Tensor Core m16n8k16)", + layers: num_layers, + weight_bytes: bf16_bytes, + prompt_tokens: prompt_len, + prefill_ms: bf16_prefill_ms, + prefill_tok_per_sec: (prompt_len as f64) / (bf16_prefill_ms / 1000.0), + decode_steps, + decode_ms_per_tok: bf16_decode_total_ms / (decode_steps as f64), + decode_tok_per_sec: (decode_steps as f64) / (bf16_decode_total_ms / 1000.0), + }, + nvfp4: ModeMetrics { + label: "NVFP4 (E2M1 + E4M3, SF_VEC = 16)", + layers: num_layers, + weight_bytes: nvfp4_bytes, + prompt_tokens: prompt_len, + prefill_ms: nvfp4_prefill_ms, + prefill_tok_per_sec: (prompt_len as f64) / (nvfp4_prefill_ms / 1000.0), + decode_steps, + decode_ms_per_tok: nvfp4_decode_total_ms / (decode_steps as f64), + decode_tok_per_sec: (decode_steps as f64) / (nvfp4_decode_total_ms / 1000.0), + }, + cosine_similarity, + compression_ratio: (bf16_bytes as f64) / (nvfp4_bytes as f64), + }) +} diff --git a/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/checkpoint.rs b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/checkpoint.rs new file mode 100644 index 0000000000..890aa1953b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/models/qwen3_8b_cute/src/checkpoint.rs @@ -0,0 +1,477 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Multi-shard safetensors loader and fast GPU weight initializer for `Qwen/Qwen3-8B`. +//! +//! Supports both: +//! - Loading a HuggingFace `Qwen/Qwen3-8B` directory (`model.safetensors.index.json` or +//! single `model.safetensors` file), and +//! - Fast deterministic GPU initialization (`init_synthetic_bf16_layer` / `init_synthetic_nvfp4_layer`) +//! so all 36 layers (`H = 4096, I = 12288, Q = 32, KV = 8, D = 128`) can be benchmarked on RTX 5090 +//! in milliseconds without requiring an external 16 GB download. + +use crate::config::{HEAD_DIM, HIDDEN, INTERMEDIATE, KV_WIDTH, Q_WIDTH}; +use crate::nvfp4::Nvfp4Projection; +use crate::operator::Result; +use ctm_rs::memory::DisjointSlice; +use ctm_rs::{cuda_module, kernel, thread}; +use cuda_core::LaunchConfig; +use cuda_core::{CudaStream, DeviceBuffer}; +use cute_rs::BFloat16; +use cute_rs::decoder::{ + AttentionWeights, DecoderWeights, FinalNorm, InputNorm, KeyNorm, NormWeights, + PostAttentionNorm, QueryNorm, +}; +use cute_rs::ffn::{ + AttentionOutput, Down, FfnWeights, Gate, Key, ProjectionWeights, Query, Up, Value, +}; +use serde::Deserialize; +use std::collections::BTreeMap; +use std::path::Path; + +/// Device kernels for fast deterministic initialization of `Qwen3-8B` BF16 and NVFP4 weights. +#[allow(missing_docs, clippy::missing_docs_in_private_items)] +#[cuda_module] +pub mod init_kernels { + use super::*; + + /// Populate a BF16 weight buffer with bounded, deterministic pseudo-orthogonal values. + #[kernel] + pub fn fill_bf16(mut out: DisjointSlice, length: usize, seed: u32, scale: f32) { + let tid = thread::index_1d().get(); + if tid >= length { + return; + } + let mut x = (tid as u32).wrapping_mul(0x9e37_79b9) ^ seed.wrapping_mul(0x85eb_ca6b); + x ^= x >> 16; + x = x.wrapping_mul(0x7feb_352d); + x ^= x >> 15; + let signed = ((x & 0xffff) as i32 - 32768) as f32 * (1.0 / 32768.0); + let mut tensor = + cute_rs::tensor::partition::StridedTensorMut::from_disjoint_slice(&mut out); + unsafe { + tensor.store_at(tid, BFloat16::from_f32(signed * scale)); + } + } + + /// Populate an RMSNorm weight buffer with values near `1.0`. + #[kernel] + pub fn fill_norm_bf16(mut out: DisjointSlice, length: usize) { + let tid = thread::index_1d().get(); + if tid >= length { + return; + } + let mut tensor = + cute_rs::tensor::partition::StridedTensorMut::from_disjoint_slice(&mut out); + unsafe { + tensor.store_at(tid, BFloat16::from_f32(1.0)); + } + } +} + +/// Owned device buffers for one `Qwen3-8B` decoder layer in BF16. +pub struct OwnedBf16Layer { + /// Pre-attention RMSNorm scale (`[4096]`). + pub input_norm: DeviceBuffer, + /// Post-attention RMSNorm scale (`[4096]`). + pub post_attention_norm: DeviceBuffer, + /// Query projection (`[4096, 4096]`). + pub q_proj: DeviceBuffer, + /// Key projection (`[1024, 4096]`). + pub k_proj: DeviceBuffer, + /// Value projection (`[1024, 4096]`). + pub v_proj: DeviceBuffer, + /// Output projection (`[4096, 4096]`). + pub o_proj: DeviceBuffer, + /// Per-head query RMSNorm scale (`[128]`). + pub q_norm: DeviceBuffer, + /// Per-head key RMSNorm scale (`[128]`). + pub k_norm: DeviceBuffer, + /// SwiGLU gate projection (`[12288, 4096]`). + pub gate_proj: DeviceBuffer, + /// SwiGLU up projection (`[12288, 4096]`). + pub up_proj: DeviceBuffer, + /// SwiGLU down projection (`[4096, 12288]`). + pub down_proj: DeviceBuffer, +} + +impl OwnedBf16Layer { + /// Borrow this layer's buffers as a role-checked `cute_rs::decoder::DecoderWeights` view. + pub fn view( + &self, + ) -> Result> { + let attention = AttentionWeights::new( + ProjectionWeights::::new(&self.q_proj, HIDDEN)?, + ProjectionWeights::::new(&self.k_proj, HIDDEN)?, + ProjectionWeights::::new(&self.v_proj, HIDDEN)?, + ProjectionWeights::